Turn Differential into a higher-order function

Author: HarrisonGrodin

README.md

@@ -37,9 +37,9 @@ using ModelingToolkit
 Then we build the system:
 
 ```julia
-eqs = [D*x ~ σ*(y-x),
-       D*y ~ x*(ρ-z)-y,
-       D*z ~ x*y - β*z]
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
 ```
 
 Each operation builds an `Operation` type, and thus `eqs` is an array of
@@ -255,7 +255,7 @@ to better scale to larger systems. You can define derivatives for your own
 function via the dispatch:
 
 ```julia
-ModelingToolkit.Derivative(::typeof(my_function),args,::Type{Val{i}})
+ModelingToolkit.Derivative(::typeof(my_function),args,::Val{i})
 ```
 
 where `i` means that it's the derivative of the `i`th argument. `args` is the
@@ -265,7 +265,7 @@ You should return an `Operation` for the derivative of your function.
 For example, `sin(t)`'s derivative (by `t`) is given by the following:
 
 ```julia
-ModelingToolkit.Derivative(::typeof(sin),args,::Type{Val{1}}) = cos(args[1])
+ModelingToolkit.Derivative(::typeof(sin),args,::Val{1}) = cos(args[1])
 ```
 
 ### Macro-free Usage

src/operators.jl

@@ -8,15 +8,11 @@ Base.show(io::IO, D::Differential) = print(io,"($(D.x),$(D.order))")
 Base.Expr(D::Differential) = :($(Symbol("D_$(D.x.name)_$(D.order)")))
 
 function Derivative end
-Base.:*(D::Differential,x::Operation) = Operation(Derivative,Expression[x,D])
-function Base.:*(D::Differential,x::Variable)
-    if D.x === x
-        return Constant(1)
-    elseif D.x ∉ x.dependents
-        return Constant(0)
-    else
-        return Variable(x,D)
-    end
+(D::Differential)(x::Operation) = Operation(Derivative,Expression[x,D])
+function (D::Differential)(x::Variable)
+    D.x === x          && return Constant(1)
+    D.x ∉ x.dependents && return Constant(0)  # FIXME
+    return Variable(x,D)
 end
 Base.:(==)(D1::Differential, D2::Differential) = D1.order == D2.order && D1.x == D2.x
 
@@ -25,22 +21,23 @@ Variable(x::Variable,D::Differential) = Variable(x.name,x.value,x.value_type,
                 x.size,x.context)
 
 function expand_derivatives(O::Operation)
+    @. O.args = expand_derivatives(O.args)
+
     if O.op == Derivative
         D = O.args[2]
         o = O.args[1]
-        simplify_constants(sum(i->Derivative(o,i)*expand_derivatives(D*o.args[i]),1:length(o.args)))
-    else
-        for i in 1:length(O.args)
-            O.args[i] = expand_derivatives(O.args[i])
-        end
+        return simplify_constants(sum(i->Derivative(o,i)*expand_derivatives(D(o.args[i])),1:length(o.args)))
     end
+
+    return O
 end
 expand_derivatives(x::Variable) = x
+expand_derivatives(D::Differential) = D
 
 # Don't specialize on the function here
 function Derivative(O::Operation,idx)
     # This calls the Derivative dispatch from the user or pre-defined code
-    Derivative(O.op,O.args,Val{idx})
+    Derivative(O.op, O.args, Val(idx))
 end
 
 # Pre-defined derivatives
@@ -83,7 +80,7 @@ function count_order(x)
     n, x.args[1]
 end
 
-function _differetial_macro(x)
+function _differential_macro(x)
     ex = Expr(:block)
     lhss = Symbol[]
     x = flatten_expr!(x)
@@ -101,11 +98,11 @@ function _differetial_macro(x)
 end
 
 macro Deriv(x...)
-    esc(_differetial_macro(x))
+    esc(_differential_macro(x))
 end
 
 function calculate_jacobian(eqs,vars)
-    Expression[Differential(vars[j])*eqs[i] for i in 1:length(eqs), j in 1:length(vars)]
+    Expression[Differential(vars[j])(eqs[i]) for i in 1:length(eqs), j in 1:length(vars)]
 end
 
 export Differential, Derivative, expand_derivatives, @Deriv, calculate_jacobian

src/systems/diffeqs/first_order_transform.jl

@@ -36,7 +36,7 @@ function ode_order_lowering!(eqs, naming_scheme)
     for sym in keys(sym_order)
         order = sym_order[sym]
         for o in (order-1):-1:1
-            lhs = D*lower_varname(sym, idv, o-1, dv_name, naming_scheme)
+            lhs = D(lower_varname(sym, idv, o-1, dv_name, naming_scheme))
             rhs = lower_varname(sym, idv, o, dv_name, naming_scheme)
             eq = Operation(==, [lhs, rhs])
             push!(eqs, eq)
@@ -46,7 +46,7 @@ function ode_order_lowering!(eqs, naming_scheme)
 end
 
 function lhs_renaming!(eq, D, naming_scheme)
-    eq.args[1] = D*lower_varname(eq.args[1], naming_scheme, lower=true)
+    eq.args[1] = D(lower_varname(eq.args[1], naming_scheme, lower=true))
     return eq
 end
 function rhs_renaming!(eq, naming_scheme)

test/basic_variables_and_operations.jl

@@ -17,8 +17,8 @@ s = JumpVariable(:s,3,dependents=[t])
 n = NoiseVariable(:n,dependents=[t])
 
 σ*(y-x)
-D*x
-D*x ~ -σ*(y-x)
-D*y ~ x*(ρ-z)-sin(y)
+D(x)
+D(x) ~ -σ*(y-x)
+D(y) ~ x*(ρ-z)-sin(y)
 
-@test D*t == Constant(1)
+@test D(t) == Constant(1)

test/components.jl

@@ -12,9 +12,9 @@ struct Lorenz <: AbstractComponent
 end
 function generate_lorenz_eqs(t,x,y,z,σ,ρ,β)
   D = Differential(t)
-  [D*x ~ σ*(y-x)
-   D*y ~ x*(ρ-z)-y
-   D*z ~ x*y - β*z]
+  [D(x) ~ σ*(y-x)
+   D(y) ~ x*(ρ-z)-y
+   D(z) ~ x*y - β*z]
 end
 Lorenz(t) = Lorenz(first(@DVar(x(t))),first(@DVar(y(t))),first(@DVar(z(t))),first(@Param(σ)),first(@Param(ρ)),first(@Param(β)),generate_lorenz_eqs(t,x,y,z,σ,ρ,β))
 

test/derivatives.jl

@@ -6,24 +6,24 @@ using Test
 @Var x(t) y(t) z(t)
 @Param σ ρ β
 @Deriv D'~t
-dsin = D*sin(t)
+dsin = D(sin(t))
 expand_derivatives(dsin)
 
 @test expand_derivatives(dsin) == cos(t)
-dcsch = D*csch(t)
+dcsch = D(csch(t))
 @test expand_derivatives(dcsch) == simplify_constants(Operation(coth(t)*csch(t)*-1))
 
 # Chain rule
-dsinsin = D*sin(sin(t))
+dsinsin = D(sin(sin(t)))
 @test expand_derivatives(dsinsin) == cos(sin(t))*cos(t)
 # Binary
-dpow1 = Derivative(^,[x, y],Val{1})
-dpow2 = Derivative(^,[x, y],Val{2})
+dpow1 = Derivative(^,[x, y],Val(1))
+dpow2 = Derivative(^,[x, y],Val(2))
 @test dpow1 == y*x^(y-1)
 @test dpow2 == x^y*log(x)
 
-d1 = D*(sin(t)*t)
-d2 = D*(sin(t)*cos(t))
+d1 = D(sin(t)*t)
+d2 = D(sin(t)*cos(t))
 @test expand_derivatives(d1) == t*cos(t)+sin(t)
 @test expand_derivatives(d2) == simplify_constants(cos(t)*cos(t)+sin(t)*(-1*sin(t)))
 

test/runtests.jl

@@ -7,5 +7,5 @@ using ModelingToolkit, Test
 @testset "Domain Test" begin include("domains.jl") end
 @testset "Simplify Test" begin include("simplify.jl") end
 @testset "Ambiguity Test" begin include("ambiguity.jl") end
-@testset "Componets Test" begin include("components.jl") end
+@testset "Components Test" begin include("components.jl") end
 @testset "System Construction Test" begin include("system_construction.jl") end

test/system_construction.jl

@@ -10,9 +10,9 @@ using Test
 @Var a
 
 # Define a differential equation
-eqs = [D*x ~ σ*(y-x),
-       D*y ~ x*(ρ-z)-y,
-       D*z ~ x*y - β*z]
+eqs = [D(x) ~ σ*(y-x),
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
 de = DiffEqSystem(eqs,[t],[x,y,z],Variable[],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
 ModelingToolkit.generate_ode_function(de;version=ModelingToolkit.SArrayFunction)
@@ -37,15 +37,15 @@ test_vars_extraction(de, de2)
 @Deriv D3'''~t
 @Deriv D2''~t
 @DVar u(t) u_tt(t) u_t(t) x_t(t)
-eqs = [D3*u ~ 2(D2*u) + D*u + D*x + 1
-       D2*x ~ D*x + 2]
+eqs = [D3(u) ~ 2(D2(u)) + D(u) + D(x) + 1
+       D2(x) ~ D(x) + 2]
 de = DiffEqSystem(eqs, [t])
 de1 = ode_order_lowering(de)
-lowered_eqs = [D*u_tt ~ 2u_tt + u_t + x_t + 1
-               D*x_t  ~ x_t + 2
-               D*u_t  ~ u_tt
-               D*u    ~ u_t
-               D*x    ~ x_t]
+lowered_eqs = [D(u_tt) ~ 2u_tt + u_t + x_t + 1
+               D(x_t)  ~ x_t + 2
+               D(u_t)  ~ u_tt
+               D(u)    ~ u_t
+               D(x)    ~ x_t]
 function test_eqs(eqs1, eqs2)
     eq = true
     for i in eachindex(eqs1)
@@ -61,9 +61,9 @@ test_eqs(de1.eqs, lowered_eqs)
 
 # Internal calculations
 eqs = [a ~ y-x,
-       D*x ~ σ*a,
-       D*y ~ x*(ρ-z)-y,
-       D*z ~ x*y - β*z]
+       D(x) ~ σ*a,
+       D(y) ~ x*(ρ-z)-y,
+       D(z) ~ x*y - β*z]
 de = DiffEqSystem(eqs,[t],[x,y,z],[a],[σ,ρ,β])
 ModelingToolkit.generate_ode_function(de)
 jac = ModelingToolkit.calculate_jacobian(de)
@@ -87,8 +87,8 @@ ModelingToolkit.generate_nlsys_function(ns)
 @Deriv D'~t
 @Param A B C
 eqs = [_x ~ y/C,
-       D*x ~ -A*x,
-       D*y ~ A*x - B*_x]
+       D(x) ~ -A*x,
+       D(y) ~ A*x - B*_x]
 de = DiffEqSystem(eqs,[t],[x,y],Variable[_x],[A,B,C])
 @test eval(ModelingToolkit.generate_ode_function(de))([0.0,0.0],[1.0,2.0],[1,2,3],0.0) ≈ -1/3