Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

We now allow src/build_function to take NamedTuples as a datastructure for building a function #929

Merged
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/build_function.jl
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,7 @@ end

unwrap_nometa(x) = unwrap(x)
unwrap_nometa(x::CallWithMetadata) = unwrap(x.f)
function destructure_arg(arg::Union{AbstractArray, Tuple}, inbounds, name)
function destructure_arg(arg::Union{AbstractArray, Tuple,NamedTuple}, inbounds, name)
if !(arg isa Arr)
DestructuredArgs(map(unwrap_nometa, arg), name, inbounds=inbounds, create_bindings=false)
else
Expand Down
197 changes: 197 additions & 0 deletions test/build_function_arrayofarray_named_tuples.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,197 @@
using Symbolics, Test, SparseArrays
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this file isn't in runtests.jl so it's not tested.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I forgot about that. Just did it.

@variables a b c

# Auxiliary Functions
get_sparsity_pattern(h::Union{SparseVector{Num}, SparseMatrixCSC{Num,Int}}) = get_sparsity_pattern(Array(h))
get_sparsity_pattern(h::Array{Num}) = sparse(Int.(.!isequal.(h, 0)))

input = (a=1, b=2, c=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!(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_ip! = eval(Symbolics.build_function(h_dense_arraymat, (a=a, b=b, c=c))[2])
h_dense_arraymat_ip_skip! = eval(Symbolics.build_function(h_dense_arraymat, (a=a, b=b, c=c), skipzeros=true, fillzeros=false)[2])
out_1_arraymat = [fill(42, 2, 2) for i in 1:3]
out_2_arraymat = deepcopy(out_1_arraymat)
h_dense_arraymat_julia!(out_1_arraymat, input)
h_dense_arraymat_ip_skip!(out_2_arraymat, input)
@test all(isequal(42), out_2_arraymat[2])
foreach(mat->fill!(mat, 0), out_2_arraymat)
h_dense_arraymat_ip_skip!(out_2_arraymat, input)
@test out_1_arraymat == out_2_arraymat
h_dense_arraymat_ip!(out_2_arraymat, input)
@test out_1_arraymat == out_2_arraymat

# Arrays of Matrices, Heterogeneous element types
test_exp = [exp(a) * exp(b), a]
h_dense_arraymat_het = [Symbolics.hessian(t, [a, b]) for t in test_exp]
function h_dense_arraymat_het_julia!(out, x)
a, b, c = x
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 = Symbolics.build_function(h_dense_arraymat_het, (a=a, b=b, c=c))
h_dense_arraymat_het_cse_str = Symbolics.build_function(h_dense_arraymat_het, (a=a, b=b, c=c), cse=true)
h_dense_arraymat_het_ip! = eval(h_dense_arraymat_het_str[2])
h_dense_arraymat_het_cse_ip! = eval(h_dense_arraymat_het_cse_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]
out_3_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)
h_dense_arraymat_het_cse_ip!(out_3_arraymat_het, input)
@test out_1_arraymat_het == out_2_arraymat_het == out_3_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]]
out[2] .= [0, 0, 0]
out[3] .= [1, a[1], b[1]]
end

h_dense_arrayvec_str = Symbolics.build_function(h_dense_arrayvec, (a=a, b=b, c=c))
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]
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!(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 = Symbolics.build_function(h_dense_arrayNestedMat, (a=a, b=b, c=c))
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)]]
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 = [[Symbolics.hessian(t, [a, b]) for t in test_exp], [Num[0 0; 0 0], Num[0 0; 0 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 = Symbolics.build_function(h_dense_arrayNestedMat_het, (a=a, b=b, c=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!(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 = Symbolics.build_function(h_sparse_arraymat, (a=a, b=b, c=c))
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
h_sparse_arraymat_ip!(out_2_arraymat, input)
h_sparse_arraymat_julia!(out_1_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!(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 = Symbolics.build_function(h_sparse_arrayvec, (a=a, b=b, c=c))
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]
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!(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 = Symbolics.build_function(h_sparse_arrayNestedMat, (a=a, b=b, c=c))
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]
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{Num,2}(undef, 0,0)]
h_empty_str = Symbolics.build_function(h_empty, (a=a, b=b, c=c))
h_empty_ip! = eval(h_empty_str[2])
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{Num}(undef,0)]
h_empty_vec_str = Symbolics.build_function(h_empty_vec, (a=a, b=b, c=c))
h_empty_vec_ip! = eval(h_empty_vec_str[2])
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{Num, 2}}(undef, 0)] # emptyNested array of arrays
h_emptyNested_str = Symbolics.build_function(h_emptyNested, (a=a, b=b, c=c))
h_emptyNested_ip! = eval(h_emptyNested_str[2])
out = [[[1 2;3 4]], Array{Array{Int64,2},1}(undef, 0)]
h_emptyNested_ip!(out, input) # should just not fail
Loading