diff --git a/QtSLiM/recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py b/QtSLiM/recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py
index a56d3c0a..51fd497f 100644
--- a/QtSLiM/recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py	
+++ b/QtSLiM/recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py	
@@ -9,23 +9,23 @@
 
 # Calculate tree heights, giving uncoalesced sites the maximum time
 def tree_heights(ts):
-    uncoalesced_height = ts.metadata['SLiM']['tick']
     heights = np.zeros(ts.num_trees + 1)
     for tree in ts.trees():
-        if tree.num_roots > 1:  # not fully coalesced
-            heights[tree.index] = uncoalesced_height
-        else:
-            children = tree.children(tree.root)
-            real_root = tree.root if len(children) > 1 else children[0]
-            heights[tree.index] = tree.time(real_root)
+        roots = tree.roots
+        if len(roots) == 1:
+            children = tree.children(roots[0])
+            if len(children) == 1:
+                roots = children
+        root_heights = [tree.time(r) for r in roots]
+        assert len(set(root_heights)) == 1
+        heights[tree.index] = root_heights[0]
     heights[-1] = heights[-2]  # repeat the last entry for plotting with step
     return heights
 
 # Plot tree heights before recapitation
 breakpoints = list(ts.breakpoints())
 heights = tree_heights(ts)
-plt.step(breakpoints, heights, where='post')
-plt.show()
+plt.step(breakpoints, heights, where='post', label="SLiM")
 
 # Recapitate!
 recap = pyslim.recapitate(ts, ancestral_Ne=1e5, recombination_rate=3e-10, random_seed=1)
@@ -34,7 +34,9 @@ def tree_heights(ts):
 # Plot the tree heights after recapitation
 breakpoints = list(recap.breakpoints())
 heights = tree_heights(recap)
-plt.step(breakpoints, heights, where='post')
-plt.show()
+plt.step(breakpoints, heights, where='post', label="recapitated")
+plt.xlabel("genome"); plt.ylabel("time ago (generations)")
+plt.savefig("tree_heights.png")
+
 
 
diff --git a/SLiMgui/Recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py b/SLiMgui/Recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py
index a56d3c0a..51fd497f 100644
--- a/SLiMgui/Recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py	
+++ b/SLiMgui/Recipes/Recipe 18.10 - Adding a neutral burn-in after simulation with recapitation II.py	
@@ -9,23 +9,23 @@
 
 # Calculate tree heights, giving uncoalesced sites the maximum time
 def tree_heights(ts):
-    uncoalesced_height = ts.metadata['SLiM']['tick']
     heights = np.zeros(ts.num_trees + 1)
     for tree in ts.trees():
-        if tree.num_roots > 1:  # not fully coalesced
-            heights[tree.index] = uncoalesced_height
-        else:
-            children = tree.children(tree.root)
-            real_root = tree.root if len(children) > 1 else children[0]
-            heights[tree.index] = tree.time(real_root)
+        roots = tree.roots
+        if len(roots) == 1:
+            children = tree.children(roots[0])
+            if len(children) == 1:
+                roots = children
+        root_heights = [tree.time(r) for r in roots]
+        assert len(set(root_heights)) == 1
+        heights[tree.index] = root_heights[0]
     heights[-1] = heights[-2]  # repeat the last entry for plotting with step
     return heights
 
 # Plot tree heights before recapitation
 breakpoints = list(ts.breakpoints())
 heights = tree_heights(ts)
-plt.step(breakpoints, heights, where='post')
-plt.show()
+plt.step(breakpoints, heights, where='post', label="SLiM")
 
 # Recapitate!
 recap = pyslim.recapitate(ts, ancestral_Ne=1e5, recombination_rate=3e-10, random_seed=1)
@@ -34,7 +34,9 @@ def tree_heights(ts):
 # Plot the tree heights after recapitation
 breakpoints = list(recap.breakpoints())
 heights = tree_heights(recap)
-plt.step(breakpoints, heights, where='post')
-plt.show()
+plt.step(breakpoints, heights, where='post', label="recapitated")
+plt.xlabel("genome"); plt.ylabel("time ago (generations)")
+plt.savefig("tree_heights.png")
+