added NPRK trees functionality in src/NPRKTrees.jl and tests in test/…

src/NPRKTrees.jl

@@ -0,0 +1,173 @@
+module NPRKTrees
+
+# Requires RootedTrees.jl package
+using RootedTrees
+
+export ARKTreeIterator
+export NPRKTreeIterator
+export isMultiColored
+export GetParentIndex
+export isColorBranching
+export elementaryWeightNPRK2
+export residualNPRK2
+
+# see tests for these functions in the script ../test/nprktrees_test.jl
+
+# Generates all non-isomorphic ARK trees with num_partitions (i.e., node colors) of a given order
+# 	returned as a Vector of ColoredRootedTree's
+function ARKTreeIterator(num_partitions::Int64, order::Int64)::Vector{<:ColoredRootedTree}
+	toReturn = ColoredRootedTree[]
+
+	if num_partitions < 1
+		throw(DomainError(num_partitions,"num_partitions argument must be integer >= 1"))
+	elseif order < 1
+		throw(DomainError(order,"order argument must be integer >= 1"))
+	elseif order == 1
+		# if order is 1, generate trivial trees with root color 0<=i<=num_partitions-1
+		for i in 0:num_partitions-1
+			trivialtree = rootedtree([1],[i])
+			push!(toReturn,trivialtree)
+		end
+	else
+		# else order > 1
+		# Generate all possible color sequences with colorings 0,1,...,num_partitions-1
+		# 	of length order
+		colors = reverse.(Iterators.product(fill(0:num_partitions-1,order)...))[:]
+
+		# Loop over all uncolored rooted trees and append the above color sequences
+		# 	to define colored rooted trees
+		# Note that this will generate duplicate trees that are isomorphic
+		for tree in RootedTreeIterator(order)
+			for c in colors
+				colorseq = collect(c)
+				push!(toReturn, rootedtree(tree.level_sequence,colorseq))
+			end
+		end
+	end
+	# remove isomorphic trees
+	toReturn = unique(toReturn,dims=1)
+	return toReturn
+end;
+
+# Generates all non-isomorphic NPRK trees with num_partitions (i.e., edge colors) of a given order
+#   equivalent to all ARK trees with num_partition colors of that order, with root node coloring fixed = 0
+# 	returned as a Vector of ColoredRootedTree's
+function NPRKTreeIterator(num_partitions::Int64, order::Int64)::Vector{<:ColoredRootedTree}
+	toReturn = ColoredRootedTree[]
+
+	if num_partitions < 1
+		throw(DomainError(num_partitions,"num_partitions argument must be integer >= 1"))
+	elseif order < 1
+		throw(DomainError(order,"order argument must be integer >= 1"))
+	elseif order == 1
+		# if order is 1, generate trivial tree with fixed root color 0
+		trivialtree = rootedtree([1],[0])
+		push!(toReturn,trivialtree)
+	else
+		# else order > 1:
+		# Generate all possible color sequences with colorings 0,1,...,partitions-1
+		# 	of length order-1 (minus one because the root coloring is fixed to 0)
+		colors = reverse.(Iterators.product(fill(0:num_partitions-1,order-1)...))[:]
+
+		# Loop over all uncolored rooted trees and append the above color sequences
+		# 	to define colored rooted trees with fixed root color 0
+		# Note that this will generate duplicate trees that are isomorphic
+		for tree in RootedTreeIterator(order)
+			for c in colors
+				colorseq = collect(c)
+				pushfirst!(colorseq,0) # push 0 to start of color sequence
+				push!(toReturn, rootedtree(tree.level_sequence,colorseq))
+			end
+		end
+	end
+	# remove isomorphic trees
+	toReturn = unique(toReturn,dims=1)
+
+	return toReturn
+end;
+
+# Determines whether a ColoredRootedTree is multicolored:
+# 	argument root_color_fixed = true for NPRK trees, i.e., checks if edges of the NPRK tree are multicolored
+# 	argument root_color_fixed = false for ARK trees, i.e., checks if nodes of the NPRK tree are multicolored
+# 		corresponds to additive coupling conditions
+function isMultiColored(colored_tree::ColoredRootedTree, root_color_fixed::Bool)::Bool
+	if root_color_fixed
+		colorseq = copy(colored_tree.color_sequence)
+		popfirst!(colorseq)
+		return ~allequal(colorseq)
+	else
+		return ~allequal(colored_tree.color_sequence)
+	end
+end
+
+# For a given tree's level sequence and a given node in that level sequence specified by position nodeindex in the level sequence
+# 	return the index in the level sequence of the parent node
+#   note that if nodeindex == 1, i.e., the node specified is the root node, this returns 0 since the root has no parent
+function GetParentIndex(levelseq::Vector, nodeindex::Int64)::Int64
+	toReturn = 0
+	for i in reverse(1:length(levelseq))
+		if levelseq[i] == levelseq[nodeindex]-1
+			toReturn = i
+			break
+		end
+	end
+	return toReturn
+end
+
+# Determines whether an NPRK tree is color branching, i.e.,
+# 	it contains a node with out-degree at least two such that at least two of its outward (or upward) edges have different colors
+#   such color branching trees corresponding to nonlinear NPRK order conditions
+# Assumes order of the tree >= 3 (trees of order 1 and 2 cannot be color branching)
+function isColorBranching(NPRK_tree::ColoredRootedTree)::Bool
+	toReturn = false
+	for i in 2:length(NPRK_tree.level_sequence)-1
+		for j in i+1:length(NPRK_tree.level_sequence)
+			if GetParentIndex(NPRK_tree.level_sequence,i) == GetParentIndex(NPRK_tree.level_sequence,j) && NPRK_tree.color_sequence[i] != NPRK_tree.color_sequence[j]
+				toReturn = true
+				break
+			end
+		end
+	end
+	return toReturn
+end
+
+# Computes the elementary weight of an NPRK tree for a given 
+#	  NPRK tableau specified by: A cubic 3-tensor and b square matrix
+#     corresponds to a 2-partition F(y,y)
+function elementaryWeightNPRK2(Atens::Array, bmatrix::Array, NPRK_tree::ColoredRootedTree)::Float64
+	toReturn = 0.0
+	if length(size(Atens)) != 3 | length(size(bmatrix)) != 2
+		throw(ArgumentError("This not a tableau for an NPRK method with 2 partitions"))
+	elseif ~allequal(size(Atens)) | ~allequal(size(bmatrix))
+		throw(ArgumentError("This tableau is not cubic/square"))	
+	elseif length(NPRK_tree.level_sequence) == 1
+		toReturn = sum(bmatrix)
+	else
+		numstages = copy(size(Atens)[1])
+		order = length(NPRK_tree.level_sequence)
+		singleitervec = collect(1:numstages)
+		indexset = Iterators.product(ntuple(i->singleitervec, 2*order)...)
+
+		for ikjk in indexset # represents indices i1, j1, i2, j2, ..., iN, jN where N=order
+			prod = 1.0
+			for k in reverse(2:order)
+				if NPRK_tree.color_sequence[k] == 0
+					prod *= Atens[ikjk[2*GetParentIndex(NPRK_tree.level_sequence, k)-1], ikjk[2*k-1], ikjk[2*k]]
+				elseif NPRK_tree.color_sequence[k] == 1
+					prod *= Atens[ikjk[2*GetParentIndex(NPRK_tree.level_sequence, k)], ikjk[2*k-1], ikjk[2*k]]
+				else
+					throw(ArgumentError("The edge colors must be 0 or 1"))
+				end
+			end
+			toReturn += bmatrix[ikjk[1],ikjk[2]]*prod
+		end
+	end
+	return toReturn
+end
+
+# computes the residual |elementaryweight - 1/density(tree)| for a given tree 
+function residualNPRK2(Atens::Array, bmatrix::Array, NPRK_tree::ColoredRootedTree)::Float64
+	return abs(elementaryWeightNPRK2(Atens, bmatrix, NPRK_tree)-1/density(NPRK_tree))
+end
+
+end

test/nprktrees_test.jl

@@ -0,0 +1,65 @@
+using RootedTrees
+include("../src/NPRKTrees.jl")
+using .NPRKTrees
+
+### TESTS BELOW ###
+# Some tests, to delete later
+
+# 3 stage Lobatto IIIA-IIIB pair produces an embedded NPRK method
+# test the Order conditions code:
+# 	Method 1 defined below should satisfy all 3rd order conditions
+# 	Method 2 satisfies all additive 3rd order conditions but fails the nonlinear order condition
+A1 = [[0 5//24 1//6]' [0 1//3 2//3]' [0 -1//24 1//6]']
+A2 = [[1//6 1//6 1//6]' [-1//6 1//3 5//6]' [0 0 0]']
+bvec = [1//6, 2//3, 1//6]
+cvec = [0, 1//2, 1]
+s = size(A1)[1]
+Atensor = zeros(s, s, s)
+bdiagmatrix = zeros(s,s) # method 1
+bdensematrix = zeros(s,s) # method 2
+
+for i in 1:s
+	bdiagmatrix[i,i] = bvec[i]
+	for j in 1:s
+		bdensematrix[i,j] = bvec[i]//s + bvec[j]//s - 1//s^2
+		for k in 1:s
+			Atensor[i,j,k] = A1[i,j]//s + A2[i,k]//s - cvec[i]//s^2
+		end
+	end
+end
+
+tol = 10^-7 # tolerance for difference between lhs and rhs of order conditions |sum[b*a*a*(...)] - 1/density(t)| < tol
+for order in 1:4
+	for t in NPRKTreeIterator(2,order)
+		# print("Order $order $(t.level_sequence) $(t.color_sequence) LHS $(NPRKOrderCondition2Partitions(Atensor, bdensematrix, t)) vs RHS $(1/density(t)), Is Tree Nonlinear: $(isColorBranching(t)) \n\n")
+		print("Method 1: Order $order $(t.level_sequence) $(t.color_sequence), residual<tol: $(residualNPRK2(Atensor, bdiagmatrix, t) < tol), Is Tree Nonlinear: $(isColorBranching(t)) \n\n")
+	end
+end
+for order in 1:4
+	for t in NPRKTreeIterator(2,order)
+		# print("Order $order $(t.level_sequence) $(t.color_sequence) LHS $(NPRKOrderCondition2Partitions(Atensor, bdensematrix, t)) vs RHS $(1/density(t)), Is Tree Nonlinear: $(isColorBranching(t)) \n\n")
+		print("Method 2: Order $order $(t.level_sequence) $(t.color_sequence), residual<tol: $(residualNPRK2(Atensor, bdensematrix, t) < tol), Is Tree Nonlinear: $(isColorBranching(t)) \n\n")
+	end
+end
+
+# Verify isMultiColored and isColorBranching work properly:
+for t in NPRKTreeIterator(3,3)
+	print(t.level_sequence)
+	print(t.color_sequence)
+	print(" isMultiColored: $(isMultiColored(t,true)), isColorBranching: $(isColorBranching(t)) \n\n")
+end
+
+# Verify the correct number of ARK trees are generated for various number of partitions and orders
+for partitions in 1:4
+	for order in 1:5
+		print("ARK $partitions $order $(length(ARKTreeIterator(partitions,order))) \n\n")
+	end
+end
+
+# Verify the correct number of NPRK trees are generated for various number of partitions and orders
+for partitions in 1:4
+	for order in 1:5
+		print("NPRK $partitions $order $(length(NPRKTreeIterator(partitions,order))) \n\n")
+	end
+end
+