Merge pull request #434 from wupeifan/arrays-of-arrays-heterogeneity

Fix issues on heterogeneous element types inside an array

Author: ChrisRackauckas

src/build_function.jl

@@ -131,6 +131,21 @@ function _build_function(target::JuliaTarget, op::Operation, args...;
     end
 end
 
+# Detect heterogeneous element types of "arrays of matrices/sparce matrices"
+function is_array_matrix(F)
+    return isa(F, AbstractVector) && all(x->isa(x, AbstractArray), F)
+end
+function is_array_sparse_matrix(F)
+    return isa(F, AbstractVector) && all(x->isa(x, AbstractSparseMatrix), F)
+end
+# Detect heterogeneous element types of "arrays of arrays of matrices/sparce matrices"
+function is_array_array_matrix(F)
+    return isa(F, AbstractVector) && all(x->isa(x, AbstractArray{<:AbstractMatrix}), F)
+end
+function is_array_array_sparse_matrix(F)
+    return isa(F, AbstractVector) && all(x->isa(x, AbstractArray{<:AbstractSparseMatrix}), F)
+end
+
 function _build_function(target::JuliaTarget, rhss, args...;
                          conv = simplified_expr, expression = Val{true},
                          checkbounds = false, constructor=nothing,
@@ -171,13 +186,13 @@ function _build_function(target::JuliaTarget, rhss, args...;
 		_rhss = rhss
 	end
 
-	if eltype(eltype(rhss)) <: SparseMatrixCSC # Array of arrays of sparse matrices
+	if is_array_array_sparse_matrix(rhss) # Array of arrays of sparse matrices
 		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j].nzval[$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2.nzval)]) for (j, rhsel2) ∈ enumerate(rhsel)], init=Expr[])) for (i,rhsel) ∈ enumerate(_rhss)],init=Expr[])
-	elseif eltype(eltype(rhss)) <: AbstractArray # Array of arrays of arrays
+	elseif is_array_array_matrix(rhss) # 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(rhss) <: SparseMatrixCSC # Array of sparse matrices
+	elseif is_array_sparse_matrix(rhss) # Array of sparse matrices
 		ip_sys_exprs = reduce(vcat,[vec([:($X[$i].nzval[$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel.nzval)]) for (i,rhsel) ∈ enumerate(_rhss)], init=Expr[])
-    elseif eltype(rhss) <: AbstractArray # Array of arrays
+    elseif is_array_matrix(rhss) # Array of arrays
 		ip_sys_exprs = reduce(vcat,[vec([:($X[$i][$j] = $(conv(rhs))) for (j, rhs) ∈ enumerate(rhsel)]) for (i,rhsel) ∈ enumerate(_rhss)], init=Expr[])
     elseif rhss isa SparseMatrixCSC
         ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(_rhss)]

test/build_function_arrayofarray.jl

@@ -11,10 +11,6 @@ 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]
@@ -23,23 +19,32 @@ function h_dense_arraymat_julia!(out, x)
 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)
+# Arrays of Matrices, Heterogeneous element types
+test_exp = [exp(a) * exp(b), a]
+h_dense_arraymat_het = [ModelingToolkit.hessian(t, [a, b]) for t in test_exp]
+function h_dense_arraymat_het_julia!(out, x)
     a, b, c = x
-    return [[a[1], 0, c[1]], [0, 0, 0], [1, a[1], b[1]]]
+    out[1] .= [exp(a[1]) * exp(b[1]) exp(a[1]) * exp(b[1]); exp(a[1]) * exp(b[1]) exp(a[1]) * exp(b[1])]
+    out[2] .= [0 0; 0 0]
 end
+
+h_dense_arraymat_het_str = ModelingToolkit.build_function(h_dense_arraymat_het, [a, b, c])
+h_dense_arraymat_het_ip! = eval(h_dense_arraymat_het_str[2])
+out_1_arraymat_het = [Array{Float64}(undef, 2, 2) for i in 1:2]
+out_2_arraymat_het = [similar(x) for x in out_1_arraymat_het]
+h_dense_arraymat_het_julia!(out_1_arraymat_het, input)
+h_dense_arraymat_het_ip!(out_2_arraymat_het, input)
+@test out_1_arraymat_het == out_2_arraymat_het
+
+# 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!(out, x)
     a, b, c = x
     out[1] .= [a[1], 0, c[1]]
@@ -48,24 +53,15 @@ function h_dense_arrayvec_julia!(out, x)
 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]
@@ -75,24 +71,34 @@ function h_dense_arrayNestedMat_julia!(out, x)
 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
 
+# Arrays of Arrays of Matrices, Heterogeneous element types
+test_exp = [exp(a) * exp(b), a]
+h_dense_arrayNestedMat_het = [[ModelingToolkit.hessian(t, [a, b]) for t in test_exp], [[ModelingToolkit.Constant(0) ModelingToolkit.Constant(0); ModelingToolkit.Constant(0) ModelingToolkit.Constant(0)], [ModelingToolkit.Constant(0) ModelingToolkit.Constant(0); ModelingToolkit.Constant(0) ModelingToolkit.Constant(0)]]]
+function h_dense_arrayNestedMat_het_julia!(out, x)
+    a, b, c = x
+    out[1][1] .= [exp(a[1]) * exp(b[1]) exp(a[1]) * exp(b[1]); exp(a[1]) * exp(b[1]) exp(a[1]) * exp(b[1])]
+    out[1][2] .= [0 0; 0 0]
+    out[2][1] .= [0 0; 0 0]
+    out[2][2] .= [0 0; 0 0]
+end
+h_dense_arrayNestedMat_het_str = ModelingToolkit.build_function(h_dense_arrayNestedMat_het, [a, b, c])
+h_dense_arrayNestedMat_het_ip! = eval(h_dense_arrayNestedMat_het_str[2])
+out_1_arrayNestedMat_het = [[rand(Int64, 2, 2), rand(Int64, 2, 2)], [rand(Int64, 2, 2), rand(Int64, 2, 2)]] # avoid undef broadcasting issue
+out_2_arrayNestedMat_het = [[rand(Int64, 2, 2), rand(Int64, 2, 2)], [rand(Int64, 2, 2), rand(Int64, 2, 2)]]
+h_dense_arrayNestedMat_julia!(out_1_arrayNestedMat_het, input)
+h_dense_arrayNestedMat_ip!(out_2_arrayNestedMat_het, input)
+@test out_1_arrayNestedMat_het == out_2_arrayNestedMat_het
+
 # ===== 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]
@@ -105,24 +111,16 @@ function h_sparse_arraymat_julia!(out, x)
 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]
@@ -134,24 +132,16 @@ function h_sparse_arrayvec_julia!(out, x)
 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 = [sparse.([[a 1; b 0], [0 0; 0 0]]), sparse.([[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]
@@ -167,14 +157,10 @@ function h_sparse_arrayNestedMat_julia!(out, x)
 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
@@ -184,26 +170,20 @@ h_sparse_arrayNestedMat_ip!(out_2_arrayNestedMat, input)
 # 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