Skip to content

Commit 4ba21a9

Browse files
committed
Simplify edges_by_child_desc
See tskit-dev/tutorials#235
1 parent 9fc9ddf commit 4ba21a9

File tree

1 file changed

+6
-13
lines changed

1 file changed

+6
-13
lines changed

tsdate/core.py

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -570,20 +570,13 @@ def edges_by_child_desc(self):
570570
Return an itertools.groupby object of edges grouped by child in descending order
571571
of the time of the child.
572572
"""
573-
wtype = np.dtype(
574-
[
575-
("child_age", self.ts.tables.nodes.time.dtype),
576-
("child_node", self.ts.tables.edges.child.dtype),
577-
]
578-
)
579-
w = np.empty(self.ts.num_edges, dtype=wtype)
580-
w["child_age"] = self.ts.tables.nodes.time[self.ts.tables.edges.child]
581-
w["child_node"] = self.ts.tables.edges.child
582-
sorted_child_parent = (
583-
self.ts.edge(i)
584-
for i in reversed(np.argsort(w, order=("child_age", "child_node")))
573+
it = (
574+
self.ts.edge(u)
575+
for u in np.lexsort(
576+
(self.ts.edges_child, -self.ts.nodes_time[self.ts.edges_child])
577+
)
585578
)
586-
return itertools.groupby(sorted_child_parent, operator.attrgetter("child"))
579+
return itertools.groupby(it, operator.attrgetter("child"))
587580

588581
def edges_by_child_then_parent_desc(self):
589582
"""

0 commit comments

Comments
 (0)