butcher_product! for colored trees

Project.toml

@@ -1,7 +1,7 @@
 name = "RootedTrees"
 uuid = "47965b36-3f3e-11e9-0dcf-4570dfd42a8c"
 authors = ["Hendrik Ranocha <[email protected]> and contributors"]
-version = "2.22.0"
+version = "2.23.0"
 
 [deps]
 LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"

src/RootedTrees.jl

@@ -1358,7 +1358,7 @@ function Base.:∘(t1::RootedTree, t2::RootedTree)
 end
 
 """
-    butcher_product!(t::RootedTree, t1::RootedTree, t2::RootedTree)
+    butcher_product!(t, t1, t2)
 
 Compute the non-associative Butcher product `t = t1 ∘ t2` of rooted trees
 in-place. It is formed by adding an edge from the root of `t1` to the root
@@ -1371,7 +1371,9 @@ Reference: Section 301 of
   Numerical methods for ordinary differential equations.
   John Wiley & Sons, 2016.
 """
-function butcher_product!(t::RootedTree, t1::RootedTree, t2::RootedTree)
+function butcher_product!(t::RootedTree,
+                          t1::RootedTree, t2::RootedTree;
+                          buffer = similar(t.level_sequence))
     offset = first(t1.level_sequence) - first(t2.level_sequence) + 1
 
     unsafe_resize!(t, order(t1) + order(t2))
@@ -1389,7 +1391,7 @@ function butcher_product!(t::RootedTree, t1::RootedTree, t2::RootedTree)
         j += 1
     end
 
-    canonical_representation!(t)
+    canonical_representation!(t; buffer)
     return t
 end
 

src/colored_trees.jl

@@ -568,6 +568,33 @@ function Base.:∘(t1::ColoredRootedTree, t2::ColoredRootedTree)
     rootedtree(level_sequence, color_sequence)
 end
 
+function butcher_product!(t::ColoredRootedTree,
+                          t1::ColoredRootedTree, t2::ColoredRootedTree,
+                          buffer_level = similar(t.level_sequence),
+                          buffer_color = similar(t.color_sequence))
+    offset = first(t1.level_sequence) - first(t2.level_sequence) + 1
+
+    unsafe_resize!(t, order(t1) + order(t2))
+    i = firstindex(t.level_sequence)
+    j = firstindex(t1.level_sequence)
+    while j <= lastindex(t1.level_sequence)
+        t.level_sequence[i] = t1.level_sequence[j]
+        t.color_sequence[i] = t1.color_sequence[j]
+        i += 1
+        j += 1
+    end
+    j = firstindex(t2.level_sequence)
+    while j <= lastindex(t2.level_sequence)
+        t.level_sequence[i] = t2.level_sequence[j] + offset
+        t.color_sequence[i] = t2.color_sequence[j]
+        i += 1
+        j += 1
+    end
+
+    canonical_representation!(t; buffer_level, buffer_color)
+    return t
+end
+
 function butcher_representation(t::ColoredRootedTree, normalize::Bool = true;
                                 colormap = _colormap_butcher_representation(t))
     if order(t) == 0