Merge pull request #643 from SciML/s/fix-buildfn

fix build_function when args are terms e.g. x(t)

Author: ChrisRackauckas

src/ModelingToolkit.jl

@@ -110,8 +110,8 @@ tosymbolic(a) = a
 @num_method Base.:(<=) (tosymbolic(a) <= tosymbolic(b)) (Real,)
 @num_method Base.:(>) (tosymbolic(a) > tosymbolic(b)) (Real,)
 @num_method Base.:(>=) (tosymbolic(a) >= tosymbolic(b)) (Real,)
-@num_method Base.isequal isequal(tosymbolic(a), tosymbolic(b)) (Number, Symbolic)
-@num_method Base.:(==) tosymbolic(a) == tosymbolic(b) (Number,)
+@num_method Base.isequal isequal(tosymbolic(a), tosymbolic(b)) (AbstractFloat, Number, Symbolic)
+@num_method Base.:(==) tosymbolic(a) == tosymbolic(b) (AbstractFloat,Number)
 
 Base.hash(x::Num, h::UInt) = hash(value(x), h)
 

src/build_function.jl

@@ -104,14 +104,16 @@ function _build_function(target::JuliaTarget, op, args...;
                          linenumbers = true, headerfun=addheader)
 
     argnames = [gensym(:MTKArg) for i in 1:length(args)]
-    arg_pairs = map(vars_to_pairs,zip(argnames,args))
+    symsdict = Dict()
+    arg_pairs = map((x,y)->vars_to_pairs(x,y, symsdict), argnames, args)
+    process = unflatten_long_ops∘(x->substitute(x, symsdict, fold=false))
     ls = reduce(vcat,first.(arg_pairs))
     rs = reduce(vcat,last.(arg_pairs))
-    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, unflatten_long_ops.(rs)))
+    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, process.(rs)))
 
     fname = gensym(:ModelingToolkitFunction)
-    op = unflatten_long_ops(op)
-    out_expr = conv(op)
+    op = process(op)
+    out_expr = conv(substitute(op, symsdict, fold=false))
     let_expr = Expr(:let, var_eqs, Expr(:block, out_expr))
     bounds_block = checkbounds ? let_expr : :(@inbounds begin $let_expr end)
 
@@ -229,7 +231,8 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
 	end
 
     argnames = [gensym(:MTKArg) for i in 1:length(args)]
-    arg_pairs = map(vars_to_pairs,zip(argnames,args))
+    symsdict = Dict()
+    arg_pairs = map((x,y)->vars_to_pairs(x,y, symsdict), argnames, args)
     ls = reduce(vcat,first.(arg_pairs))
     rs = reduce(vcat,last.(arg_pairs))
     var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, rs))
@@ -241,12 +244,14 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     oidx = isnothing(outputidxs) ? (i -> i) : (i -> outputidxs[i])
     X = gensym(:MTIIPVar)
 
+    process = unflatten_long_ops∘(x->substitute(x, symsdict, fold=false))
+
     if rhss isa SparseMatrixCSC
         rhs_length = length(rhss.nzval)
-        rhss = SparseMatrixCSC(rhss.m, rhss.m, rhss.colptr, rhss.rowval, map(unflatten_long_ops, rhss.nzval))
+        rhss = SparseMatrixCSC(rhss.m, rhss.m, rhss.colptr, rhss.rowval, map(process, rhss.nzval))
     else
         rhs_length = length(rhss)
-        rhss = [unflatten_long_ops(r) for r in rhss]
+        rhss = [process(r) for r in rhss]
     end
 
 	if parallel isa DistributedForm
@@ -388,9 +393,9 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     end
 
     if rhss isa SparseMatrixCSC
-        rhss′ = map(conv∘unflatten_long_ops, rhss.nzval)
+        rhss′ = map(conv∘process, rhss.nzval)
     else
-        rhss′ = [conv(unflatten_long_ops(r)) for r in rhss]
+        rhss′ = [conv(process(r)) for r in rhss]
     end
 
     tuple_sys_expr = build_expr(:tuple, rhss′)
@@ -456,13 +461,16 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     end
 end
 
-vars_to_pairs(args) = vars_to_pairs(args[1],args[2])
-function vars_to_pairs(name,vs::AbstractArray)
-    vs_names = [tosymbol(u) for u ∈ vs]
+function vars_to_pairs(name,vs::Union{Tuple, AbstractArray}, symsdict=Dict())
+    vs_names = tosymbol.(vs)
+    for (v,k) in zip(vs_names, vs)
+        symsdict[k] = v
+    end
     exs = [:($name[$i]) for (i, u) ∈ enumerate(vs)]
     vs_names,exs
 end
-function vars_to_pairs(name,vs)
+function vars_to_pairs(name,vs, symsdict)
+    symsdict[vs] = tosymbol(vs)
     [tosymbol(vs)], [name]
 end
 

test/build_function.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, Test
+using ModelingToolkit, SparseArrays, Test
 @variables a b c1 c2 c3 d e g
 
 # Multiple argument matrix
@@ -71,3 +71,21 @@ h_oop_scalar = eval(h_str_scalar)
 
 @test isequal(simplify(ModelingToolkit.unflatten_long_ops(prod(z))),
               simplify(prod(z)))
+
+@variables t x(t) y(t) k
+f = eval(build_function((x+y)/k, [x,y,k]))
+@test f([1,1,2]) == 1
+
+f = eval(build_function([(x+y)/k], [x,y,k])[1])
+@test f([1,1,2]) == [1]
+
+f = eval(build_function([(x+y)/k], [x,y,k])[2])
+z = [0.0]
+f(z, [1,1,2])
+@test z == [1]
+
+f = eval(build_function(sparse([1],[1], [(x+y)/k], 10,10), [x,y,k])[1])
+
+@test size(f([1.,1.,2])) == (10,10)
+@test f([1.,1.,2])[1,1] == 1.0
+@test sum(f([1.,1.,2])) == 1.0

test/latexify.jl

@@ -42,9 +42,9 @@ eqs = [D(u[1]) ~ p[3]*(u[2]-u[1]),
 
 @test latexify(eqs) ==
 raw"\begin{align}
-\frac{du_{1(t)}}{dt} =& p_{3} \left( \mathrm{u_2}\left( t \right) - \mathrm{u_1}\left( t \right) \right) \\
-0 =& \frac{p_{2} p_{3} \mathrm{u_1}\left( t \right) \left( p_{1} - \mathrm{u_1}\left( t \right) \right)}{10} - \mathrm{u_2}\left( t \right) \\
-\frac{du_{3(t)}}{dt} =& \mathrm{u_1}\left( t \right) \left( \mathrm{u_2}\left( t \right) \right)^{\frac{2}{3}} - p_{3} \mathrm{u_3}\left( t \right)
+\frac{du{_1}(t)}{dt} =& p{_3} \left( \mathrm{u{_2}}\left( t \right) - \mathrm{u{_1}}\left( t \right) \right) \\
+0 =& \frac{p{_2} p{_3} \mathrm{u{_1}}\left( t \right) \left( p{_1} - \mathrm{u{_1}}\left( t \right) \right)}{10} - \mathrm{u{_2}}\left( t \right) \\
+\frac{du{_3}(t)}{dt} =& \mathrm{u{_1}}\left( t \right) \left( \mathrm{u{_2}}\left( t \right) \right)^{\frac{2}{3}} - p{_3} \mathrm{u{_3}}\left( t \right)
 \end{align}
 "
 
@@ -56,8 +56,8 @@ sys = ODESystem(eqs)
 
 @test latexify(eqs) ==
 raw"\begin{align}
-\frac{du_{1(t)}}{dt} =& p_{3} \left( \mathrm{u_2}\left( t \right) - \mathrm{u_1}\left( t \right) \right) \\
-\frac{du_{2(t)}}{dt} =& \frac{p_{2} p_{3} \mathrm{u_1}\left( t \right) \left( p_{1} - \mathrm{u_1}\left( t \right) \right)}{10} - \mathrm{u_2}\left( t \right) \\
-\frac{du_{3(t)}}{dt} =& \mathrm{u_1}\left( t \right) \left( \mathrm{u_2}\left( t \right) \right)^{\frac{2}{3}} - p_{3} \mathrm{u_3}\left( t \right)
+\frac{du{_1}(t)}{dt} =& p{_3} \left( \mathrm{u{_2}}\left( t \right) - \mathrm{u{_1}}\left( t \right) \right) \\
+\frac{du{_2}(t)}{dt} =& \frac{p{_2} p{_3} \mathrm{u{_1}}\left( t \right) \left( p{_1} - \mathrm{u{_1}}\left( t \right) \right)}{10} - \mathrm{u{_2}}\left( t \right) \\
+\frac{du{_3}(t)}{dt} =& \mathrm{u{_1}}\left( t \right) \left( \mathrm{u{_2}}\left( t \right) \right)^{\frac{2}{3}} - p{_3} \mathrm{u{_3}}\left( t \right)
 \end{align}
 "

test/runtests.jl

@@ -28,8 +28,8 @@ using SafeTestsets, Test
 @safetestset "Depdendency Graph Test" begin include("dep_graphs.jl") end
 @safetestset "Function Registration Test" begin include("function_registration.jl") end
 @safetestset "Array of Array Test" begin include("build_function_arrayofarray.jl") end
-@safetestset "Latexify recipes Test" begin include("latexify.jl") end
 @testset "Distributed Test" begin include("distributed.jl") end
 @safetestset "Variable Utils Test" begin include("variable_utils.jl") end
 println("Last test requires gcc available in the path!")
 @safetestset "C Compilation Test" begin include("ccompile.jl") end
+@safetestset "Latexify recipes Test" begin include("latexify.jl") end