allow arrays of arrays of stuff in build_function iip

Fixes #262

Author: ChrisRackauckas

src/build_function.jl

@@ -107,7 +107,15 @@ function _build_function(target::JuliaTarget, rhss, args...;
     fargs = Expr(:tuple,argnames...)
 
     X = gensym(:MTIIPVar)
-    if rhss isa SparseMatrixCSC
+	if eltype(eltype(rhss)) <: AbstractArray # Array of arrays of arrays
+		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j][$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)],init=Expr[])) for (i,rhsel) ∈ enumerate(rhss)],init=Expr[])
+	elseif eltype(eltype(rhss)) <: SparseMatrixCSC # Array of arrays of arrays
+		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j].nzval[$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)])) for (i,rhsel) ∈ enumerate(rhss)])
+	elseif eltype(rhss) <: SparseMatrixCSC # Array of sparse matrices
+		ip_sys_exprs = reduce(vcat,[vec([:($X[$i].nzval[$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(rhss)])
+    elseif eltype(rhss) <: AbstractArray # Array of arrays
+		ip_sys_exprs = reduce(vcat,[vec([:($X[$i][$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(rhss)])
+    elseif rhss isa SparseMatrixCSC
         ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss.nzval)]
     else
         ip_sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]

src/direct.jl

@@ -54,3 +54,5 @@ simplified_expr(c::Constant) = c.value
 function simplified_expr(eq::Equation)
     Expr(:(=), simplified_expr(eq.lhs), simplified_expr(eq.rhs))
 end
+
+simplified_expr(eq::AbstractArray) = simplified_expr.(eq)

test/build_function_arrayofarray.jl

@@ -0,0 +1,209 @@
+using ModelingToolkit, Test, SparseArrays
+@variables a b c
+
+# Auxiliary Functions and Constants
+get_sparsity_pattern(h::Array{Expression}) = sparse(Int64.(map(~, h .=== ModelingToolkit.Constant(0))))
+get_sparsity_pattern(h::SparseMatrixCSC{Expression,Int64}) = sparse(Int64.(map(~, h .=== ModelingToolkit.Constant(0))))
+get_sparsity_pattern(h::SparseVector{Expression,Int64}) = sparse(Int64.(map(~, h .=== ModelingToolkit.Constant(0))))
+
+input = [1, 2, 3]
+
+# ===== Dense tests =====
+# Arrays of Matrices
+h_dense_arraymat = [[a 1; b 0], [0 0; 0 0], [a c; 1 0]] # empty array support required
+function h_dense_arraymat_julia(x)
+    a, b, c = x
+    return [[a[1] 1; b[1] 0], [0 0; 0 0], [a[1] c[1]; 1 0]]
+end
+function h_dense_arraymat_julia!(out, x)
+    a, b, c = x
+    out[1] .= [a[1] 1; b[1] 0]
+    out[2] .= [0 0; 0 0]
+    out[3] .= [a[1] c[1]; 1 0]
+end
+
+h_dense_arraymat_str = ModelingToolkit.build_function(h_dense_arraymat, [a, b, c])
+h_dense_arraymat_oop = eval(h_dense_arraymat_str[1])
+h_dense_arraymat_ip! = eval(h_dense_arraymat_str[2])
+out_1_arraymat = [Array{Int64}(undef, 2, 2) for i in 1:3]
+out_2_arraymat = [similar(x) for x in out_1_arraymat]
+julia_dense_arraymat = h_dense_arraymat_julia(input)
+mtk_dense_arraymat = h_dense_arraymat_oop(input)
+@test_broken julia_dense_arraymat == mtk_dense_arraymat
+h_dense_arraymat_julia!(out_1_arraymat, input)
+h_dense_arraymat_ip!(out_2_arraymat, input)
+@test out_1_arraymat == out_2_arraymat
+
+# Arrays of 1D Vectors
+h_dense_arrayvec = [[a, 0, c], [0, 0, 0], [1, a, b]] # same for empty vectors, etc.
+function h_dense_arrayvec_julia(x)
+    a, b, c = x
+    return [[a[1], 0, c[1]], [0, 0, 0], [1, a[1], b[1]]]
+end
+function h_dense_arrayvec_julia!(out, x)
+    a, b, c = x
+    out[1] .= [a[1], 0, c[1]]
+    out[2] .= [0, 0, 0]
+    out[3] .= [1, a[1], b[1]]
+end
+
+h_dense_arrayvec_str = ModelingToolkit.build_function(h_dense_arrayvec, [a, b, c])
+h_dense_arrayvec_oop = eval(h_dense_arrayvec_str[1])
+h_dense_arrayvec_ip! = eval(h_dense_arrayvec_str[2])
+out_1_arrayvec = [Vector{Int64}(undef, 3) for i in 1:3]
+out_2_arrayvec = [Vector{Int64}(undef, 3) for i in 1:3]
+julia_dense_arrayvec = h_dense_arrayvec_julia(input)
+mtk_dense_arrayvec = h_dense_arrayvec_oop(input)
+@test_broken julia_dense_arrayvec == mtk_dense_arrayvec
+mtk_dense_arrayvec = @test_broken h_dense_arrayvec_oop(input)
+h_dense_arrayvec_julia!(out_1_arrayvec, input)
+h_dense_arrayvec_ip!(out_2_arrayvec, input)
+@test out_1_arrayvec == out_2_arrayvec
+
+# Arrays of Arrays of Matrices
+h_dense_arrayNestedMat = [[[a 1; b 0], [0 0; 0 0]], [[b 1; a 0], [b c; 0 1]]]
+function h_dense_arrayNestedMat_julia(x)
+    a, b, c = x
+    return [[[a[1] 1; b[1] 0], [0 0; 0 0]], [[b[1] 1; a[1] 0], [b[1] c[1]; 0 1]]]
+end
+function h_dense_arrayNestedMat_julia!(out, x)
+    a, b, c = x
+    out[1][1] .= [a[1] 1; b[1] 0]
+    out[1][2] .= [0 0; 0 0]
+    out[2][1] .= [b[1] 1; a[1] 0]
+    out[2][2] .= [b[1] c[1]; 0 1]
+end
+
+h_dense_arrayNestedMat_str = ModelingToolkit.build_function(h_dense_arrayNestedMat, [a, b, c])
+h_dense_arrayNestedMat_oop = eval(h_dense_arrayNestedMat_str[1])
+h_dense_arrayNestedMat_ip! = eval(h_dense_arrayNestedMat_str[2])
+out_1_arrayNestedMat = [[rand(Int64, 2, 2), rand(Int64, 2, 2)], [rand(Int64, 2, 2), rand(Int64, 2, 2)]] # avoid undef broadcasting issue
+out_2_arrayNestedMat = [[rand(Int64, 2, 2), rand(Int64, 2, 2)], [rand(Int64, 2, 2), rand(Int64, 2, 2)]]
+julia_dense_arrayNestedMat = h_dense_arrayNestedMat_julia(input)
+mtk_dense_arrayNestedMat = h_dense_arrayNestedMat_oop(input)
+@test_broken julia_dense_arrayNestedMat == mtk_dense_arrayNestedMat
+h_dense_arrayNestedMat_julia!(out_1_arrayNestedMat, input)
+h_dense_arrayNestedMat_ip!(out_2_arrayNestedMat, input)
+@test out_1_arrayNestedMat == out_2_arrayNestedMat
+
+# ===== Sparse tests =====
+# Array of Matrices
+h_sparse_arraymat = sparse.([[a 1; b 0], [0 0; 0 0], [a c; 1 0]])
+function h_sparse_arraymat_julia(x)
+    a, b, c = x
+    return [sparse([a[1] 1; b[1] 0]), sparse([0 0; 0 0]), sparse([a[1] c[1]; 1 0])] # necessary because sparse([]) is a SparseVector, not SparseMatrix
+end
+function h_sparse_arraymat_julia!(out, x)
+    a, b, c = x
+    out[1][1, 1] = a[1]
+    out[1][1, 2] = 1
+    out[1][2, 1] = b[1]
+    out[2] = sparse([0 0; 0 0]) # no undef constructor for SparseMatrixCSC
+    out[3][1, 1] = a[1]
+    out[3][1, 2] = c[1]
+    out[3][2, 1] = 1
+end
+
+h_sparse_arraymat_str = ModelingToolkit.build_function(h_sparse_arraymat, [a, b, c])
+h_sparse_arraymat_oop = eval(h_sparse_arraymat_str[1])
+h_sparse_arraymat_ip! = eval(h_sparse_arraymat_str[2])
+h_sparse_arraymat_sparsity_patterns = map(get_sparsity_pattern, h_sparse_arraymat)
+out_1_arraymat = [similar(h) for h in h_sparse_arraymat_sparsity_patterns]
+out_2_arraymat = [similar(h) for h in h_sparse_arraymat_sparsity_patterns] # can't do similar() because it will just be #undef, with the wrong sparsity pattern
+julia_sparse_arraymat = h_sparse_arraymat_julia(input)
+mtk_sparse_arraymat = h_sparse_arraymat_oop(input)
+@test_broken julia_sparse_arraymat == mtk_sparse_arraymat
+h_sparse_arraymat_julia!(out_1_arraymat, input)
+h_sparse_arraymat_ip!(out_2_arraymat, input)
+@test out_1_arraymat == out_2_arraymat
+
+# Array of 1D Vectors
+h_sparse_arrayvec = sparse.([[a, 0, c], [0, 0, 0], [1, a, b]])
+function h_sparse_arrayvec_julia(x)
+    a, b, c = x
+    return sparse.([[a[1], 0, c[1]], [0, 0, 0], [1, a[1], b[1]]])
+end
+function h_sparse_arrayvec_julia!(out, x)
+    a, b, c = x
+    out[1][1] = a[1]
+    out[1][3] = c[1]
+    out[2] = sparse([0, 0, 0]) # necessary because sparsity pattern is 3 elements with 0 stored, not 0 elements
+    out[3][1] = 1
+    out[3][2] = a[1]
+    out[3][3] = b[1]
+end
+
+h_sparse_arrayvec_str = ModelingToolkit.build_function(h_sparse_arrayvec, [a, b, c])
+h_sparse_arrayvec_oop = eval(h_sparse_arrayvec_str[1])
+h_sparse_arrayvec_ip! = eval(h_sparse_arrayvec_str[2])
+h_sparse_arrayvec_sparsity_patterns = map(get_sparsity_pattern, h_sparse_arrayvec)
+out_1_arrayvec = [similar(h) for h in h_sparse_arrayvec_sparsity_patterns]
+out_2_arrayvec = [similar(h) for h in h_sparse_arrayvec_sparsity_patterns]
+julia_sparse_arrayvec = h_sparse_arrayvec_julia(input)
+mtk_sparse_arrayvec = h_sparse_arrayvec_oop(input)
+@test_broken julia_sparse_arrayvec == mtk_sparse_arrayvec
+h_sparse_arrayvec_julia!(out_1_arrayvec, input)
+h_sparse_arrayvec_ip!(out_2_arrayvec, input)
+@test out_1_arrayvec == out_2_arrayvec
+
+# Arrays of Arrays of Matrices
+h_sparse_arrayNestedMat = [[[a 1; b 0], [0 0; 0 0]], [[b 1; a 0], [b c; 0 1]]]
+function h_sparse_arrayNestedMat_julia(x)
+    a, b, c = x
+    return [sparse.([[a[1] 1; b[1] 0], [0 0; 0 0]]), sparse.([[b[1] 1; a[1] 0], [b[1] c[1]; 0 1]])]
+end
+function h_sparse_arrayNestedMat_julia!(out, x)
+    a, b, c = x
+    out[1][1][1, 1] = a[1]
+    out[1][1][1, 2] = 1
+    out[1][1][2, 1] = b[1]
+    out[1][2] = sparse([0 0; 0 0])
+    out[2][1][1, 1] = b[1]
+    out[2][1][1, 2] = 1
+    out[2][1][2, 1] = a[1]
+    out[2][2][1, 1] = b[1]
+    out[2][2][1, 2] = c[1]
+    out[2][2][2, 2] = 1
+end
+
+h_sparse_arrayNestedMat_str = ModelingToolkit.build_function(h_sparse_arrayNestedMat, [a, b, c])
+h_sparse_arrayNestedMat_oop = eval(h_sparse_arrayNestedMat_str[1])
+h_sparse_arrayNestedMat_ip! = eval(h_sparse_arrayNestedMat_str[2])
+h_sparse_arrayNestedMat_sparsity_patterns = [map(get_sparsity_pattern, h) for h in h_sparse_arrayNestedMat]
+out_1_arrayNestedMat = [[similar(h_sub) for h_sub in h] for h in h_sparse_arrayNestedMat_sparsity_patterns]
+out_2_arrayNestedMat = [[similar(h_sub) for h_sub in h] for h in h_sparse_arrayNestedMat_sparsity_patterns]
+julia_sparse_arrayNestedMat = h_sparse_arrayNestedMat_julia(input)
+mtk_sparse_arrayNestedMat = h_sparse_arrayNestedMat_oop(input)
+@test_broken julia_sparse_arrayNestedMat == mtk_sparse_arrayNestedMat
+h_sparse_arrayNestedMat_julia!(out_1_arrayNestedMat, input)
+h_sparse_arrayNestedMat_ip!(out_2_arrayNestedMat, input)
+@test out_1_arrayNestedMat == out_2_arrayNestedMat
+
+# Additional Tests
+# Returning 0-element structures (corresponding to empty Jacobians)
+# Arrays of Matrices
+h_empty = [[a b; c 0], Array{Expression,2}(undef, 0,0)]
+h_empty_str = ModelingToolkit.build_function(h_empty, [a, b, c])
+h_empty_oop = eval(h_empty_str[1])
+h_empty_ip! = eval(h_empty_str[2])
+@test_broken h_empty_oop(input) == [[1 2; 3 0], Array{Int64,2}(undef,0,0)]
+out = [Matrix{Int64}(undef, 2, 2), Matrix{Int64}(undef, 0, 0)]
+h_empty_ip!(out, input) # should just not fail
+
+# Array of Vectors
+h_empty_vec = [[a, b, c, 0], Vector{Expression}(undef,0)]
+h_empty_vec_str = ModelingToolkit.build_function(h_empty_vec, [a, b, c])
+h_empty_vec_oop = eval(h_empty_vec_str[1])
+h_empty_vec_ip! = eval(h_empty_vec_str[2])
+@test_broken h_empty_vec_oop(input) == [[1, 2, 3, 0], Vector{Int64}(undef,0)]
+out = [Vector{Int64}(undef, 4), Vector{Int64}(undef, 0)]
+h_empty_vec_ip!(out, input) # should just not fail
+
+# Arrays of Arrays of Matrices
+h_emptyNested = [[[a b; c 0]], Array{Array{Expression, 2}}(undef, 0)] # emptyNested array of arrays
+h_emptyNested_str = ModelingToolkit.build_function(h_emptyNested, [a, b, c])
+h_emptyNested_oop = eval(h_emptyNested_str[1])
+h_emptyNested_ip! = eval(h_emptyNested_str[2])
+@test_broken h_emptyNested_oop(input) == [[[1 2; 3 0]], Array{Array{Int64, 2}}(undef, 0)]
+out = [[[1 2;3 4]], Array{Array{Int64,2},1}(undef, 0)]
+h_emptyNested_ip!(out, input) # should just not fail