Merge pull request #663 from SciML/s/fix-register

Wrap registered fn output in Num if any input is Num

Author: YingboMa

src/register_function.jl

@@ -17,7 +17,12 @@ macro register(expr, Ts = [Num, Symbolic, Real])
     name(x::Expr) = ((@assert x.head == :(::)); :($value($(x.args[1]))))
 
     Expr(:block,
-         [:($f($(setinds(args, symbolic_args, ts)...)) = Term{Real}($f, [$(map(name, args)...)]))
+         [quote
+              function $f($(setinds(args, symbolic_args, ts)...))
+                  wrap =  any(x->typeof(x) <: Num, tuple($(setinds(args, symbolic_args, ts)...),)) ? Num : identity
+                  wrap(Term{Real}($f, [$(map(name, args)...)]))
+              end
+          end
          for ts in types]...) |> esc
 end
 

test/function_registration.jl

@@ -110,3 +110,9 @@ function run_test()
     @test fun([0.5], [u0], 0.) == [do_something_4(u0) * 2]
 end
 run_test()
+
+using ModelingToolkit: arguments
+@variables a
+@register foo(x,y,z)
+@test 1 * foo(a,a,a) * Num(1) isa Num
+@test !any(x->x isa Num, arguments(value(1 * foo(a,a,a) * Num(1))))