Merge pull request #448 from SciML/myb/skipzeros

Add skipzeros

Author: YingboMa

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "3.8.1"
+version = "3.9.0"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/build_function.jl

@@ -167,7 +167,7 @@ function _build_function(target::JuliaTarget, rhss, args...;
                          checkbounds = false, constructor=nothing,
                          linenumbers = false, multithread=nothing,
                          headerfun=addheader, outputidxs=nothing,
-						 parallel=SerialForm())
+                         skipzeros = false, parallel=SerialForm())
 
 	if multithread isa Bool
 		@warn("multithraded is deprecated for the parallel argument. See the documentation.")
@@ -202,18 +202,55 @@ function _build_function(target::JuliaTarget, rhss, args...;
 		_rhss = rhss
 	end
 
-	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 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 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[])
+    ip_sys_exprs = Expr[]
+    if is_array_array_sparse_matrix(rhss) # Array of arrays of sparse matrices
+        for (i, rhsel) ∈ enumerate(_rhss)
+            for (j, rhsel2) ∈ enumerate(rhsel)
+                for (k, rhs) ∈ enumerate(rhsel2.nzval)
+                    rhs′ = conv(rhs)
+                    (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+                    push!(ip_sys_exprs, :($X[$i][$j].nzval[$k] = $rhs′))
+                end
+            end
+        end
+    elseif is_array_array_matrix(rhss) # Array of arrays of arrays
+        for (i, rhsel) ∈ enumerate(_rhss)
+            for (j, rhsel2) ∈ enumerate(rhsel)
+                for (k, rhs) ∈ enumerate(rhsel2)
+                    rhs′ = conv(rhs)
+                    (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+                    push!(ip_sys_exprs, :($X[$i][$j][$k] = $rhs′))
+                end
+            end
+        end
+    elseif is_array_sparse_matrix(rhss) # Array of sparse matrices
+        for (i, rhsel) ∈ enumerate(_rhss)
+            for (j, rhs) ∈ enumerate(rhsel.nzval)
+                rhs′ = conv(rhs)
+                (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+                push!(ip_sys_exprs, :($X[$i].nzval[$j] = $rhs′))
+            end
+        end
     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[])
+        for (i, rhsel) ∈ enumerate(_rhss)
+            for (j, rhs) ∈ enumerate(rhsel)
+                rhs′ = conv(rhs)
+                (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+                push!(ip_sys_exprs, :($X[$i][$j] = $rhs′))
+            end
+        end
     elseif rhss isa SparseMatrixCSC
-        ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(_rhss)]
+        for (i, rhs) ∈ enumerate(_rhss)
+            rhs′ = conv(rhs)
+            (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+            push!(ip_sys_exprs, :($X.nzval[$i] = $rhs′))
+        end
     else
-        ip_sys_exprs = [:($X[$(oidx(i))] = $(conv(rhs))) for (i, rhs) ∈ enumerate(_rhss)]
+        for (i, rhs) ∈ enumerate(_rhss)
+            rhs′ = conv(rhs)
+            (skipzeros && rhs′ isa Number && iszero(rhs′)) && continue
+            push!(ip_sys_exprs, :($X[$(oidx(i))] = $rhs′))
+        end
     end
 
     ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))

src/utils.jl

@@ -133,4 +133,3 @@ end
 _substitute(expr, ks, vs) = substituter(ks, vs)(expr)
 
 @deprecate substitute_expr!(expr,s) substitute(expr,s)
-

test/build_function.jl

@@ -2,23 +2,33 @@ using ModelingToolkit, Test
 @variables a b c1 c2 c3 d e g
 
 # Multiple argument matrix
-h = [a + b + c1 + c2; c3 + d + e + g] # uses the same number of arguments as our application
-h_julia(a, b, c, d, e, g) = [a[1] + b[1] + c[1] + c[2]; c[3] + d[1] + e[1] + g[1]]
+h = [a + b + c1 + c2,
+     c3 + d + e + g,
+     0] # uses the same number of arguments as our application
+h_julia(a, b, c, d, e, g) = [a[1] + b[1] + c[1] + c[2],
+                             c[3] + d[1] + e[1] + g[1],
+                             0]
 function h_julia!(out, a, b, c, d, e, g)
-    out .= [a[1] + b[1] + c[1] + c[2]; c[3] + d[1] + e[1] + g[1]]
+    out .= [a[1] + b[1] + c[1] + c[2], c[3] + d[1] + e[1] + g[1], 0]
 end
 
 h_str = ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g])
 h_oop = eval(h_str[1])
 h_ip! = eval(h_str[2])
+h_ip_skip! = eval(ModelingToolkit.build_function(h, [a], [b], [c1, c2, c3], [d], [e], [g], skipzeros=true)[2])
 inputs = ([1], [2], [3, 4, 5], [6], [7], [8])
 
 @test h_oop(inputs...) == h_julia(inputs...)
-out_1 = Array{Int64}(undef, 2)
+out_1 = similar(h, Int)
 out_2 = similar(out_1)
 h_ip!(out_1, inputs...)
 h_julia!(out_2, inputs...)
 @test out_1 == out_2
+fill!(out_1, 10)
+h_ip_skip!(out_1, inputs...)
+@test out_1[3] == 10
+out_1[3] = 0
+@test out_1 == out_2
 
 # Multiple input matrix, some unused arguments
 h_skip = [a + b + c1; c2 + c3 + g] # skip d, e

test/build_function_arrayofarray.jl

@@ -18,11 +18,16 @@ function h_dense_arraymat_julia!(out, x)
     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_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]
+h_dense_arraymat_ip! = eval(ModelingToolkit.build_function(h_dense_arraymat, [a, b, c])[2])
+h_dense_arraymat_ip_skip! = eval(ModelingToolkit.build_function(h_dense_arraymat, [a, b, c], skipzeros=true)[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