redo register macro

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

src/ModelingToolkit.jl

@@ -175,6 +175,7 @@ include("linearity.jl")
 include("solve.jl")
 include("direct.jl")
 include("domains.jl")
+include("register_function.jl")
 
 include("systems/abstractsystem.jl")
 

src/differentials.jl

@@ -191,13 +191,13 @@ for (modu, fun, arity) ∈ DiffRules.diffrules()
         else
             DiffRules.diffrule(modu, fun, ntuple(k->:(args[$k]), arity)...)[i]
         end
-        @eval derivative(::typeof($modu.$fun), args::NTuple{$arity,Any}, ::Val{$i}) = (x = $expr; !(x isa Expression || x isa Constant) ? Constant(x) : x)
+        @eval derivative(::typeof($modu.$fun), args::NTuple{$arity,Any}, ::Val{$i}) = $expr
     end
 end
 
-derivative(::typeof(+), args::NTuple{N,Any}, ::Val) where {N} = Constant(1)
+derivative(::typeof(+), args::NTuple{N,Any}, ::Val) where {N} = 1
 derivative(::typeof(*), args::NTuple{N,Any}, ::Val{i}) where {N,i} = make_operation(*, deleteat!(collect(args), i))
-derivative(::typeof(one), args::Tuple{<:Any}, ::Val) = Constant(0)
+derivative(::typeof(one), args::Tuple{<:Any}, ::Val) = 0
 
 function count_order(x)
     @assert !(x isa Symbol) "The variable $x must have an order of differentiation that is greater or equal to 1!"

src/direct.jl

@@ -215,7 +215,7 @@ function toexpr(O::Term)
           return Expr(:call, :^, Expr(:call, :inv, toexpr(O.args[1])), -(O.args[2].value))
       end
   end
-  return Expr(:call, Symbol(O.op), toexpr.(O.args)...)
+  return Expr(:call, O.op, toexpr.(O.args)...)
 end
 toexpr(s::Sym) = nameof(s)
 toexpr(s) = s

src/extra_functions.jl

@@ -1,14 +1,13 @@
-@register Base.conj(x)
+@register Base.getindex(x,i::Integer)
 @register Base.getindex(x,i)
 @register Base.binomial(n,k)
-@register Base.copysign(x,y)
 
 @register Base.signbit(x)
 ModelingToolkit.derivative(::typeof(signbit), args::NTuple{1,Any}, ::Val{1}) = 0
 
 ModelingToolkit.derivative(::typeof(abs), args::NTuple{1,Any}, ::Val{1}) = IfElse.ifelse(signbit(args[1]),-one(args[1]),one(args[1]))
 
-@register IfElse.ifelse(x,y,z)
+@register IfElse.ifelse(x,y,z::Any)
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{1}) = 0
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{2}) = IfElse.ifelse(args[1],1,0)
 ModelingToolkit.derivative(::typeof(IfElse.ifelse), args::NTuple{3,Any}, ::Val{3}) = IfElse.ifelse(args[1],0,1)

src/register_function.jl

@@ -0,0 +1,52 @@
+macro register(expr, Ts = [Num, Symbolic, Number])
+    @assert expr.head == :call
+
+    f = expr.args[1]
+    args = expr.args[2:end]
+
+    symbolic_args = findall(x->x isa Symbol, args)
+
+    types = vec(collect(Iterators.product(ntuple(_->Ts, length(symbolic_args))...)))
+
+    # remove Number Number Number methods
+    filter!(Ts->!all(T->T == Number, Ts), types)
+
+    annotype(name,T) = :($name :: $T)
+    setinds(xs, idx, vs) = (xs=copy(xs); xs[idx] .= map(annotype, xs[idx], vs); xs)
+    name(x::Symbol) = :($value($x))
+    name(x::Expr) = ((@assert x.head == :(::)); :($value($(x.args[1]))))
+
+    Expr(:block,
+         [:($f($(setinds(args, symbolic_args, ts)...)) = Term{Number}($f, [$(map(name, args)...)]))
+         for ts in types]...) |> esc
+end
+
+# Ensure that Operations that get @registered from outside the ModelingToolkit
+# module can work without having to bring in the associated function into the
+# ModelingToolkit namespace. We basically store information about functions
+# registered at runtime in a ModelingToolkit variable,
+# `registered_external_functions`. It's not pretty, but we are limited by the
+# way GeneralizedGenerated builds a function (adding "ModelingToolkit" to every
+# function call).
+# ---
+const registered_external_functions = Dict{Symbol,Module}()
+function inject_registered_module_functions(expr)
+    MacroTools.postwalk(expr) do x
+        # Find all function calls in the expression and extract the function
+        # name and calling module.
+        MacroTools.@capture(x, f_module_.f_name_(xs__))
+        if isnothing(f_module)
+            MacroTools.@capture(x, f_name_(xs__))
+        end
+
+        if !isnothing(f_name)
+            # Set the calling module to the module that registered it.
+            mod = get(registered_external_functions, f_name, f_module)
+            if !isnothing(mod) && mod != Base
+                x.args[1] = :($mod.$f_name)
+            end
+        end
+
+        return x
+    end
+end

src/systems/diffeqs/abstractodesystem.jl

@@ -44,7 +44,7 @@ function (f::ODEToExpr)(O::Term)
         any(isequal(O), f.states) && return O.op.name  # dependent variables
         return build_expr(:call, Any[O.op.name; f.(O.args)])
     end
-    return build_expr(:call, Any[Symbol(O.op); f.(O.args)])
+    return build_expr(:call, Any[O.op; f.(O.args)])
 end
 (f::ODEToExpr)(x) = toexpr(x)
 

src/utils.jl

@@ -134,7 +134,7 @@ function states_to_sym(states)
                 any(isequal(O), states) && return O.op.name  # dependent variables
                 return build_expr(:call, Any[O.op.name; _states_to_sym.(O.args)])
             else
-                return build_expr(:call, Any[Symbol(O.op); _states_to_sym.(O.args)])
+                return build_expr(:call, Any[O.op; _states_to_sym.(O.args)])
             end
         elseif O isa Num
             return _states_to_sym(value(O))

test/function_registration.jl

@@ -79,8 +79,8 @@ foo(x, y) = sin(x) * cos(y)
 using ModelingToolkit: value
 expr = value(foo(x, y))
 @test expr.op === foo
-@test expr.args[1] === x
-@test expr.args[2] === y
+@test expr.args[1] === value(x)
+@test expr.args[2] === value(y)
 ModelingToolkit.derivative(::typeof(foo), (x, y), ::Val{1}) = cos(x) * cos(y) # derivative w.r.t. the first argument
 ModelingToolkit.derivative(::typeof(foo), (x, y), ::Val{2}) = -sin(x) * sin(y) # derivative w.r.t. the second argument
 @test isequal(expand_derivatives(D(foo(x, y))), expand_derivatives(D(sin(x) * cos(y))))