Efficient graph APIs #2869
Replies: 18 comments 27 replies
-
|
Here's an ARG-type algorithm that fits well: def descendant_span(ts, arg, u):
"""
Return an array giving the total sequence lenght over which
each node in the tree sequence descends from the specified node.
"""
total_descending = np.zeros(ts.num_nodes)
stack = [(u, 0, ts.sequence_length)]
total_descending[u] = ts.sequence_length # NOTE questionable quick hack
while len(stack) > 0:
u, left, right = stack.pop()
# NOTE: if we had an index here sorted by left coord
# we could binary search to first match, and could
# break once e_right > left (I think?)
for e in range(*arg.child_range[u]):
e_left = ts.edges_left[e]
e_right = ts.edges_right[e]
if e_right > left and right > e_left:
inter_left = max(e_left, left)
inter_right = min(e_right, right)
e_child = ts.edges_child[e]
total_descending[e_child] += inter_right - inter_left
stack.append((e_child, inter_left, inter_right))
return total_descending
def descendant_span_tree(ts, u):
total_descending = np.zeros(ts.num_nodes)
for tree in ts.trees():
descendants = tree.preorder(u)
total_descending[descendants] += tree.span
return total_descending
for u in range(ts.num_nodes):
d1 = descendant_span(ts, arg, u)
d2 = descendant_span_tree(ts, u)
np.testing.assert_array_equal(d1, d2)I think this is basically what you were looking for at some point @hyanwong? This would be substantially faster than using set-ops, I expect (you just filter by descendant_span > 0 to find nodes that descend somewhere) I think the ARG algorithm could be made quicker with an index and some more logic, but it would probably be good to run on something larger before making any solid judgements. |
Beta Was this translation helpful? Give feedback.
-
|
Just tried a quick numba'd version of this: @numba.njit
def _descendant_span(u, num_nodes, sequence_length, arg_child_range,
edges_left, edges_right, edges_child):
total_descending = np.zeros(num_nodes)
stack = [(u, 0, sequence_length)]
while len(stack) > 0:
u, left, right = stack.pop()
for e in range(arg_child_range[u, 0], arg_child_range[u, 1]):
e_left = edges_left[e]
e_right = edges_right[e]
if e_right > left and right > e_left:
inter_left = max(e_left, left)
inter_right = min(e_right, right)
e_child = edges_child[e]
total_descending[e_child] += inter_right - inter_left
stack.append((e_child, inter_left, inter_right))
return total_descending
def descendant_span_numba(ts, arg, u):
"""
Return an array giving the total sequence lenght over which
each node in the tree sequence descends from the specified node.
"""
return _descendant_span(u, ts.num_nodes, ts.sequence_length,
arg.child_range, ts.edges_left, ts.edges_right,
ts.edges_child)On a SARS-CoV-2 ARG (780k nodes) the Python version takes 2.6 seconds to run for the "root" node (should be worst case scenario) and the numba version takes 33ms (after warming the jit). The (fairly fast) tree-by-tree method above takes 14 seconds. |
Beta Was this translation helpful? Give feedback.
-
|
I don't really understand how this works without maintaining an interval library, or assuming some sort of nestedness within the inherited intervals, but I think I'm missing something. Maybe you could explain it to me on the board? |
Beta Was this translation helpful? Give feedback.
-
|
No need for an interval library, we're doing it ourselves (there's tons of interval stuff in, e.g., simplify). We're just computing interval intersections one-by-one here. |
Beta Was this translation helpful? Give feedback.
-
The main tsdate algorithm just requires running through the all the edges for a parent, ordering parents by time (easy), or all the edges for a child (ordering children by time, then by parent id: requires a sort like the one here). To traverse the DAG efficiently, visiting children before parents (or vice versa), might we require the new tables to be ordered by node time? This structure might make it slightly easier, but I don't think there's much in it? We would require the secondary sort key to be the node ID, not the left coordinate, though. |
Beta Was this translation helpful? Give feedback.
-
|
Indexing into the In tsdate we rely, roughly, on the fact that edges are sorted first by parent (time), then by child ID, then by left coord. I think that the list of sequential parent edges for a child should therefore be sorted by parent ID, and only after by left coord? That we we can do a You could imagine wrapping this in something so that you could easily traverse through the parent-child (or child-parent) links without caring about the intervals. |
Beta Was this translation helpful? Give feedback.
-
I thought hard about this, and I don't see a way of separating the graph edges from the inheritance intervals while still maintaining an efficient, numba-able structure. |
Beta Was this translation helpful? Give feedback.
-
|
I'm glad you thought about it: I can well believe it would kill efficiency. If the edges within the parent_range are sorted by child id first (i.e. child ID is your secondary sort key), then at least you don't need to sort the edge-intervals to get the list of child IDs: you just skip until the child ID changes (equivalent to wrapping in We could probably make a wrapper function that does this skipping and which is numba-able? |
Beta Was this translation helpful? Give feedback.
-
|
For the list of descendants of a node u, you could imagine wanting either the list of samples that have inherited anything from u (which is what you have), or the list of descendants of u regardless of whether they have inherited anything from it. This latter list is actually easier to construct, as you don't have any intervals to keep track of. Another interesting (related) one is, for a set of samples s , to find the most recent common ancestor anywhere in the genome, or the most recent common ancestor in the ARG regardless of genomic inheritance. |
Beta Was this translation helpful? Give feedback.
-
Here's the graph descendants (although I wonder what it would actually be useful for?) def graph_descendants(ts, arg, u):
"""
Return boolean array marking whether a node is a graph descendant
of the specified node u. Note that this does not require that the
node inherited any genetic material.
"""
is_descendant = np.zeros(ts.num_nodes, dtype=bool)
stack = [u]
while len(stack) > 0:
u = stack.pop()
for e in range(*arg.child_range[u]):
e_child = ts.edges_child[e]
if not is_descendant[e_child]:
# Note: setting is_descendant here because we can
# push the same node on the stack multiple times otherwise
is_descendant[e_child] = True
stack.append(e_child)
return is_descendant |
Beta Was this translation helpful? Give feedback.
-
|
Yes, thanks, that's a nice simple one. Re usage, we might want to find approximations to pedigree relations from an ARG, without caring about genetic inheritance, right? Like we might want to know the biparental MRCA from a non-sample-resolved forward simulation. |
Beta Was this translation helpful? Give feedback.
-
|
FYI, another way of getting the graph descendants is to make the adjacency matrix (as a sparse matrix) and then do matrix multiplication: if |
Beta Was this translation helpful? Give feedback.
-
|
Unlikely to be faster I would imagine, but it's a great way to test! |
Beta Was this translation helpful? Give feedback.
-
|
Here's some corresponding operations for ancestors: def ancestor_span(ts, arg, u):
"""
Return an array giving the total sequence lenght over which
each node in the tree sequence is an ancestor to the specified node.
"""
total_ancestral = np.zeros(ts.num_nodes)
stack = [(u, 0, ts.sequence_length)]
total_ancestral[u] = ts.sequence_length # NOTE questionable quick hack
while len(stack) > 0:
u, left, right = stack.pop()
# NOTE: if we had an index here sorted by left coord
# we could binary search to first match, and could
# break once e_right > left (I think?)
for j in range(*arg.parent_range[u]):
e = arg.parent_index[j]
e_left = ts.edges_left[e]
e_right = ts.edges_right[e]
if e_right > left and right > e_left:
inter_left = max(e_left, left)
inter_right = min(e_right, right)
e_parent = ts.edges_parent[e]
total_ancestral[e_parent] += inter_right - inter_left
stack.append((e_parent, inter_left, inter_right))
return total_ancestral
def ancestor_span_tree(ts, u):
total_ancestral = np.zeros(ts.num_nodes)
for tree in ts.trees():
v = u
while v != -1:
total_ancestral[v] += tree.span
v = tree.parent(v)
return total_ancestral
def graph_ancestors(ts, arg, u):
"""
Return boolean array marking whether a node is a graph ancestor
of the specified node u. Note that this does not require that the
node inherited any genetic material.
"""
is_ancestor = np.zeros(ts.num_nodes, dtype=bool)
is_ancestor[u] = True
stack = [u]
while len(stack) > 0:
u = stack.pop()
for j in range(*arg.parent_range[u]):
e = arg.parent_index[j]
e_parent = ts.edges_parent[e]
if not is_ancestor[e_parent]:
# Note: setting is_ancestor here because we can
# push the same node on the stack multiple times otherwise
is_ancestor[e_parent] = True
stack.append(e_parent)
return is_ancestor
for u in range(ts.num_nodes):
d1 = ancestor_span(ts, arg, u)
d2 = ancestor_span_tree(ts, u)
np.testing.assert_array_equal(d1, d2)Note these are basically symmetric with the |
Beta Was this translation helpful? Give feedback.
-
|
We can use the def graph_mrca(ts, arg, u, v):
u_ancestors = graph_ancestors(ts, arg, u)
v_ancestors = graph_ancestors(ts, arg, v)
common = np.logical_and(u_ancestors, v_ancestors)
min_index = np.argmin(ts.nodes_time[common])
return np.where(common)[0][min_index]You could implement it a bit better, but the time difference wouldn't be that much. The "real" MRCA is much harder because you have to propagate the combined ancestral material through the graph for u and v. What you're doing, essentially, is simplifying the ARG wrt to the samples u and v, so implementing in terms of simplify would be a reasonable thing to do (if not optimal). |
Beta Was this translation helpful? Give feedback.
-
|
I assume most of these algorithms would work fine if we had a "non-sample-resolved" ARG, e.g. from an unsimplified forward simulation? We'd have lots of "hanging topology", but the interval tracking stuff should still work OK. |
Beta Was this translation helpful? Give feedback.
-
|
These algorithms should be fully general, but I haven't thought about the specifics. There may be some corner cases lurking if you ran them on an unsimplified prospective ARG, but generally I'd expect them to work. |
Beta Was this translation helpful? Give feedback.
-
|
Yep, that was my thought too. If we do implement anything, we should run tests on both simplified and unsimplified / prospective genealogies, to check we get the same answers. |
Beta Was this translation helpful? Give feedback.
-
|
So, the interesting algorithmic question that's emerging here is whether sorting parent/child intervals by left coordinate for would actually make the fundamental Are there graphs in which we have enough intervals per parent/child pair to make sorting worthwhile? It's not clear to me. |
Beta Was this translation helpful? Give feedback.
-
|
I'm not sure how much efficiency we would gain. I think the main gain is in syntactic simplicity, so we need to think very hard about how we present the API to the user. For instance, I see something like the following as canonical: for node_obj in ts.arg.nodes(time="asc"):
for arg_edge in node_obj.parent_arg_edges():
parent = arg_edge.parent
child = node_obj.id
for interval in arg_edge.intervals():
# Do something with the intervals(and similarly with |
Beta Was this translation helpful? Give feedback.
-
|
That's a separate issue - we're not talking about high-level APIs yet: this is about how we support efficient ARG operations. Python loops are by definition not efficient. |
Beta Was this translation helpful? Give feedback.
-
|
Right, but accessing the intervals for a given parent/child combination is not efficient if the indexes have the secondary sort key of |
Beta Was this translation helpful? Give feedback.
-
|
I don't understand - all of the above examples are doing this efficiently? It's just the ordering of the intervals within a parent-child combination that's in question. |
Beta Was this translation helpful? Give feedback.
-
|
I thought that if the first sort key was |
Beta Was this translation helpful? Give feedback.
-
|
Here's a version in which we capture the sub-ARG descending from a given individual as an edge table: def descendant_intervals(ts, arg, u):
descending = tskit.EdgeTable()
stack = [(u, 0, ts.sequence_length)]
while len(stack) > 0:
u, left, right = stack.pop()
for e in range(*arg.child_range[u]):
e_left = ts.edges_left[e]
e_right = ts.edges_right[e]
if e_right > left and right > e_left:
inter_left = max(e_left, left)
inter_right = min(e_right, right)
e_child = ts.edges_child[e]
descending.add_row(inter_left, inter_right, u, e_child)
stack.append((e_child, inter_left, inter_right))
return descendingThis is basically the same as |
Beta Was this translation helpful? Give feedback.
-
|
The main use for these ARG algorithms are when we want to perform node-focussed analysis (e.g. we have a big ARG, and want to know specific things about a node or set of nodes and their ancestors/descendants). We discussed a few interesting use cases:
|
Beta Was this translation helpful? Give feedback.
-
Interesting one. Easy way to do it (not checking actual syntax): descending = descendant_intervals(ts, arg, u)
ancestral = ancestral_intervals(ts, arg, u)
tables = ts.tables.copy()
tables.edges = ancestral + descending # This doesn't work, but you get the idea
tables.delete_older(ts.nodes_time[u] + 20)
tables.delete_younger(ts.nodes_time[u] - 20) # method doesn't exist
tables.sort() # Needed?
subts = tables.tree_sequence() |
Beta Was this translation helpful? Give feedback.
-
|
I just thought of another case where vertical traversal might make more sense that left-right traversing. Imagine a large number of |
Beta Was this translation helpful? Give feedback.
-
|
This is very nice! Tell me if this is right - the point is essentially providing efficient ways to look up "who are the parents" and "who are the children" of a given node(and on which intervals)? Thus making bottom-up or top-down iteration easier? And so it's 'just' indexing into the same data structure, but by saying "ARG API" you're saying "ARG" (as opposed to "tree sequence") to connote thinking in the up/down direction, rather than left/right? I like it, although I worry a bit that the ARG/tree sequence terminology doesn't overlap terribly well with our definitions of those? Another set of operations are IBD-related ones. |
Beta Was this translation helpful? Give feedback.
-
|
Yes, that's basically it I think. The current TreeSequence APIs are very good at doing tree-by-tree stuff, but not good at all at answering questions about particular nodes. Viewing the structure as a graph rather than a sequence of trees makes these time-order operations simpler and more efficient. @gtsambos has already done a bunch of pioneering work on bottom-to-top algorithms (including IBD), so she might have some thoughts? The downside is that we would need to build some additional indexes to fully support the forward-time-order operations (see the initial post above). We might be able to make some backwards time operations a bit more efficient by sorting intervals within a parent node by left coordinate, however. The nomeclature is really hard here. Perhaps it would be better to reserve the term "ARG" as the general, slightly vague thing, and refer to this specific structure here (the view on the edge table, plus indexes) as an "InheritanceGraph" or something? It's more descriptive and may be less confusing in the long run. So, the ARG is the general thing, and a succinct tree sequence (nodes + edge table + other tables) is the specific encoding of an ARG. A TreeSequence is a view on the node and edge tables with two additional left-to-right sort-indexes on the edges which facilitates fast tree-by-tree operations. An InheritanceGraph is also a view on the edge table, with two additional forward-and-backwards time sorting-indexes into the edges. This probably would be less confusing than "An ARG is the general thing, and succinct... A TreeSequence is a view... An ARG is another view ..." We don't have to, and probably shouldn't try, sort out nomenclature here, I'm just trying to clarify here for the sake of understanding what the structure above is for and why we might want to do it. |
Beta Was this translation helpful? Give feedback.
-
|
+1 for Inheritance Graph. Maybe even +2. Agreed we don't have to sort that out here - I just needed to get straight in my head what the big picture goal here was. |
Beta Was this translation helpful? Give feedback.
-
|
Also - I am still impressed by how clever and simple this idea is. |
Beta Was this translation helpful? Give feedback.
-
|
It's neat isn't it? I feel like |
Beta Was this translation helpful? Give feedback.
-
|
Totally. I guess that simplify is already reasonably natural because the ordering in the tables derives from the natural ordering of coalescent simulations, which is back-in-time. Pushing things forwards in time, on the other hand... |
Beta Was this translation helpful? Give feedback.
-
|
Here's another useful node-focussed algorithm: give 2 samples, which part(s) of the genome have the most recent MRCA. I.e. which pieces of the genome are the closest, and what ancestral node(s) do those correspond to? |
Beta Was this translation helpful? Give feedback.
-
|
That's the same question as the genetic mrca isn't it? With the addition of returning the actual intervals, which you would have to track anyway. |
Beta Was this translation helpful? Give feedback.
-
|
Oh yes, sorry. I had lost track of the discussion on that point. The complications are (a) there could be multiple MRCAs at the same time and (b) I guess it would be useful to do a trick like the one in the normal MRCA function which doesn't require us to find all the ancestors but use the time ordering of parents to work out which interval we should look at next, to have a chance of finding the most recent common ancestor, without having to go all the way back to the root. For a pair of nodes (especially if they are closely related), I think this would be more efficient than simplifying wrt the two nodes? Although I'm not sure of the internal details of simplify. |
Beta Was this translation helpful? Give feedback.
-
|
Yes, I'm sure it could be implemented more efficiently, but it's basically the same algorithm |
Beta Was this translation helpful? Give feedback.
-
|
I'm looking at shared recombinations at the moment, and, for a specific node which has multiple ARG parents, where the parent changes from A to B at a specific breakpoint x, I want to find the "MRCA" shared between parent A in the left tree and parent B in the right tree. I also needed to do exactly this for the sc2ts tree, to find the "age" of the recombinant parents. So I suspect this might be a commonly-used ARG operation? Basically: how distantly related are two parents at a breakpoint position? At the moment I have to jump to the tree at the breakpoint, find the chain of parents from B, then switch to the previous tree and find the other chain from A. There's probably a neater way to do this on a per-node basis using the ARG structures than by building the entire tree every time, right? |
Beta Was this translation helpful? Give feedback.
-
|
You could use the A minor tweak would give you all graph most recent common ancestors, if multiple exist at the same time. That's probably what you want? If you want the genetic MRCAs, then you'd need to use |
Beta Was this translation helpful? Give feedback.
-
|
I doubt either of those is much quicker than just looking at the tree on the left and right of the breakpoint, though. |
Beta Was this translation helpful? Give feedback.
-
|
Hey folks, just getting to this now, it's been a hectic few weeks here.
Will weigh in once I've parsed all you've already done here!
…On Wed, Nov 15, 2023, 5:49 AM Jerome Kelleher ***@***.***> wrote:
I doubt either of those is much quicker than just looking at the tree on
the left and right of the breakpoint, though.
—
Reply to this email directly, view it on GitHub
<https://urldefense.com/v3/__https://github.com/tskit-dev/tskit/discussions/2869*discussioncomment-7577258__;Iw!!K-Hz7m0Vt54!hpem6iY0PcerQ7cRAYZnZ6wARUKq0BMvFI2cSrPJwfd8jqOiul_Yf5uKTc0CnAl0f4XWy50nu6QSyePGcJ9GmFn8$>,
or unsubscribe
<https://urldefense.com/v3/__https://github.com/notifications/unsubscribe-auth/AEHOXQTK4QZABDWIT6GX4H3YETCADAVCNFSM6AAAAAA7GDTAA6VHI2DSMVQWIX3LMV43SRDJONRXK43TNFXW4Q3PNVWWK3TUHM3TKNZXGI2TQ__;!!K-Hz7m0Vt54!hpem6iY0PcerQ7cRAYZnZ6wARUKq0BMvFI2cSrPJwfd8jqOiul_Yf5uKTc0CnAl0f4XWy50nu6QSyePGcJx3k62Y$>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
Beta Was this translation helpful? Give feedback.
-
|
I just noticed that in the original definition of I am coming round to the idea that the next sort key after this should be the left position, as this allows more efficient processing of intervals. We should probably sort by both right coordinate and parent ID after than, so that we enforce a canonical ordering? |
Beta Was this translation helpful? Give feedback.
-
|
Here's another nice algorithm to count the number of samples reachable from each node in the graph: #2882 (comment) |
Beta Was this translation helpful? Give feedback.
-
|
Another easy thing to write using parent indexes is a |
Beta Was this translation helpful? Give feedback.
-
We don't have much support at the moment for ways to view the whole tree sequence as a graph - everything is very much focused on the "sequence of trees" view. We should add some graph-level functionality. One way we can do this is to add an
ARGclass (and correspondingtsk_arg_tC struct) which is analogous to the Tree class: a view on the underlying data model with some additional data structures to facilitate computations.Before getting into details, it's important to set out some requirements for the low-level implementation. We can certainly build higher-level Python stuff on top of this (like export to networkx, etc), but we have to do it in a way that's efficient at the low-level first.
The first two requirements rule out quite a lot of possibilities when representing a graph structure. An adjacency matrix would work, but would need a lot of memory and this wouldn't solve the problem of how to annotate the edges with inheritance intervals.
After some thought, I think something like this is as good as we're going to do:
(The eagle-eyed will notice that the
lexsortin here violates property 3, but we can get back to that later).The idea is that we use the existing edge table to represent the inheritance intervals. Ultimately, we have to attach O(num_edges) information to the graph, so we may as well use what we already have. Each node is associated with a set of indexes into the edge table, those which define either the parents or children of that node. Because the edge table is already sorted by
parentID (through the sortedness requirements) we can represent all the child edges of a node by two numbers: thestartandstopvalues for the range of indexes.To do the same for finding the parents of a node, we have to build an additional index into the edge table (
parent_indexabove).Let's see an example of this working:
gives
So far so good. The most basic application I can think of for an ARG data structure is to recover the local tree at a given position, so:
Note that here we've make the left coordinate secondary sorting criterion in the
parent_indexwhich makes finding the edge for a given position potentially binary searchable.Additional indexes
The only real problem with this approach I think is that it requires us to do an expensive sort on the edge table. I don't see any way around this, if we want efficient access to the parent edges for a given node. In principle, we can build this extra index in the build_index function, and store on disk. We deliberately left the indexing open-ended to facilitate adding additional indexes later. Questions about whether to build this index by default etc would need to answered.
Secondary sort key?
Assuming we are building an extra index for ARG traversal, we could need to decide what the secondary sort key is. Should we group the parent edges for a given node by ID or left coordinate? Left coordinate seems more useful to me (see above) but I guess it would be useful to keep all the edges for a given parent node adjacent when doing a graph traversal?
If we decide to sort by left coordinate for the required new
parent_index, we would have to consider also computing and storing the equivalentchild_index, which has the left coordinate as the secondary sort key, rather thanchildnode. This is in principle much cheaper though, as the edge table is already sorted by parent ID and so you would just need to do lots of much smaller sorts within the edges for a given parent.Applications
There's no point in building these APIs if we don't have some applications. The only things I can think of (besides tree construction, which may potentially be faster than what we have) are:
@hyanwong @nspope - you're in the weeds with tsdate right now, would the structures above make your lives any easier?
Beta Was this translation helpful? Give feedback.
All reactions