Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Document numpy traversals #1952

Merged
merged 1 commit into from
Dec 1, 2021
Merged

Document numpy traversals #1952

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 50 additions & 2 deletions python/tskit/trees.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2102,15 +2102,51 @@ def num_tracked_samples(self, u=None):
u = self.virtual_root if u is None else u
return self._ll_tree.get_num_tracked_samples(u)

# TODO document these traversal arrays
# https://github.com/tskit-dev/tskit/issues/1788
def preorder(self, u=NULL):
"""
Returns a numpy array of node ids in preorder
<https://en.wikipedia.org/wiki/Tree_traversal#Pre-order_(NLR)>. If the node u
is specified the traversal is rooted at this node (and it will be the first
element in the returned array). Otherwise, all nodes reachable from the tree
roots will be returned. See :ref:`tutorials:sec_analysing_trees_traversals` for
examples.

:param int u: If specified, return all nodes in the subtree rooted at u
(including u) in traversal order.
:return: Array of node ids
:rtype: numpy.ndarray (dtype=np.int32)
"""
return self._ll_tree.get_preorder(u)

def postorder(self, u=NULL):
"""
Returns a numpy array of node ids in postorder
<https://en.wikipedia.org/wiki/Tree_traversal##Post-order_(LRN)>. If the node u
is specified the traversal is rooted at this node (and it will be the last
element in the returned array). Otherwise, all nodes reachable from the tree
roots will be returned. See :ref:`tutorials:sec_analysing_trees_traversals` for
examples.

:param int u: If specified, return all nodes in the subtree rooted at u
(including u) in traversal order.
:return: Array of node ids
:rtype: numpy.ndarray (dtype=np.int32)
"""
return self._ll_tree.get_postorder(u)

def timeasc(self, u=NULL):
"""
Returns a numpy array of node ids. Starting at `u`, returns the reachable
descendant nodes in order of increasing time (most recent first), falling back
to increasing ID if times are equal. Also see
:ref:`tutorials:sec_analysing_trees_traversals` for examples of how to use
traversals.

:param int u: If specified, return all nodes in the subtree rooted at u
(including u) in traversal order.
:return: Array of node ids
:rtype: numpy.ndarray (dtype=np.int32)
"""
nodes = self.preorder(u)
is_virtual_root = u == self.virtual_root
time = self.tree_sequence.tables.nodes.time
Expand All @@ -2124,6 +2160,18 @@ def timeasc(self, u=NULL):
return nodes[order]

def timedesc(self, u=NULL):
"""
Returns a numpy array of node ids. Starting at `u`, returns the reachable
descendant nodes in order of decreasing time (least recent first), falling back
to decreasing ID if times are equal. Also see
:ref:`tutorials:sec_analysing_trees_traversals` for examples of how to use
traversals.

:param int u: If specified, return all nodes in the subtree rooted at u
(including u) in traversal order.
:return: Array of node ids
:rtype: numpy.ndarray (dtype=np.int32)
"""
return self.timeasc(u)[::-1]

def _preorder_traversal(self, root):
Expand Down