Merge pull request #342 from SciML/substitute

substitute using SymbolicUtils

Author: ChrisRackauckas

Project.toml

@@ -37,7 +37,7 @@ NaNMath = "0.3"
 SafeTestsets = "0.0.1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10"
 StaticArrays = "0.10, 0.11, 0.12"
-SymbolicUtils = "0.1.1, 0.2"
+SymbolicUtils = "0.3"
 TreeViews = "0.3"
 UnPack = "0.1"
 Unitful = "1.1"

src/ModelingToolkit.jl

@@ -13,6 +13,9 @@ import GeneralizedGenerated
 using DocStringExtensions
 using Base: RefValue
 
+import SymbolicUtils
+import SymbolicUtils: to_symbolic, FnType
+
 import TreeViews
 
 """
@@ -112,17 +115,18 @@ export Reaction, ReactionSystem
 export Differential, expand_derivatives, @derivatives
 export IntervalDomain, ProductDomain, ⊗, CircleDomain
 export Equation, ConstrainedEquation
-export simplify_constants
-
 export Operation, Expression, Variable
+export independent_variable, states, parameters, equations
+
 export calculate_jacobian, generate_jacobian, generate_function
 export calculate_tgrad, generate_tgrad
 export calculate_gradient, generate_gradient
 export calculate_factorized_W, generate_factorized_W
 export calculate_hessian, generate_hessian
 export calculate_massmatrix, generate_diffusion_function
-export independent_variable, states, parameters, equations
-export simplified_expr, rename, get_variables, substitute_expr!
+
+export simplified_expr, rename, get_variables
+export simplify, substitute
 export build_function
 export @register
 export modelingtoolkitize

src/differentials.jl

@@ -49,10 +49,10 @@ function expand_derivatives(O::Operation)
 
         return sum(1:length(o.args)) do i
             derivative(o, i) * expand_derivatives(D(o.args[i]))
-        end |> simplify_constants
+        end |> simplify
     end
 
-    return simplify_constants(O)
+    return simplify(O)
 end
 expand_derivatives(x) = x
 

src/direct.jl

@@ -8,7 +8,7 @@ an array of variable expressions.
 """
 function gradient(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
     out = [expand_derivatives(Differential(v)(O)) for v in vars]
-    simplify ? simplify_constants.(out) : out
+    simplify ? ModelingToolkit.simplify.(out) : out
 end
 
 """
@@ -21,7 +21,7 @@ an array of variable expressions.
 """
 function jacobian(ops::AbstractVector{<:Expression}, vars::AbstractVector{<:Expression}; simplify = true)
     out = [expand_derivatives(Differential(v)(O)) for O in ops, v in vars]
-    simplify ? simplify_constants.(out) : out
+    simplify ? ModelingToolkit.simplify.(out) : out
 end
 
 """
@@ -34,7 +34,7 @@ an array of variable expressions.
 """
 function hessian(O::Expression, vars::AbstractVector{<:Expression}; simplify = true)
     out = [expand_derivatives(Differential(v2)(Differential(v1)(O))) for v1 in vars, v2 in vars]
-    simplify ? simplify_constants.(out) : out
+    simplify ? ModelingToolkit.simplify.(out) : out
 end
 
 function simplified_expr(O::Operation)

src/simplify.jl

@@ -1,6 +1,3 @@
-import SymbolicUtils
-import SymbolicUtils: FnType
-
 # ModelingToolkit -> SymbolicUtils
 SymbolicUtils.istree(x::Operation) = true
 function SymbolicUtils.operation(x::Operation)
@@ -38,9 +35,10 @@ SymbolicUtils.symtype(x::Expression) = Number
 
 # SymbolicUtils -> ModelingToolkit
 
-function simplify_constants(expr)
-    SymbolicUtils.simplify(expr) |> to_mtk
-end
+simplify(expr::Expression) = SymbolicUtils.simplify(expr) |> to_mtk
+simplify(expr) = expr |> to_mtk
+
+@deprecate simplify_constants(ex) simplify(ex)
 
 to_mtk(x) = x
 to_mtk(x::Number) = Constant(x)

src/systems/diffeqs/abstractodesystem.jl

@@ -69,13 +69,13 @@ function calculate_factorized_W(sys::AbstractODESystem, simplify=true)
     Wfact = lu(W, Val(false), check=false).factors
 
     if simplify
-        Wfact = simplify_constants.(Wfact)
+        Wfact = ModelingToolkit.simplify.(Wfact)
     end
 
     W_t = - LinearAlgebra.I/gam + jac
     Wfact_t = lu(W_t, Val(false), check=false).factors
     if simplify
-        Wfact_t = simplify_constants.(Wfact_t)
+        Wfact_t = ModelingToolkit.simplify.(Wfact_t)
     end
     sys.Wfact[] = Wfact
     sys.Wfact_t[] = Wfact_t
@@ -113,7 +113,7 @@ function calculate_massmatrix(sys::AbstractODESystem, simplify=true)
             error("Only semi-explicit constant mass matrices are currently supported")
         end
     end
-    M = simplify ? simplify_constants.(M) : M
+    M = simplify ? ModelingToolkit.simplify.(M) : M
     # M should only contain concrete numbers
     M = map(x->x isa Constant ? x.value : x, M)
     M == I ? I : M

src/systems/jumps/jumpsystem.jl

@@ -9,24 +9,24 @@ struct JumpSystem <: AbstractSystem
     systems::Vector{JumpSystem}
 end
 
-function JumpSystem(eqs, iv, states, ps; systems = JumpSystem[], 
+function JumpSystem(eqs, iv, states, ps; systems = JumpSystem[],
                                           name = gensym(:JumpSystem))
     JumpSystem(eqs, iv, convert.(Variable, states), convert.(Variable, ps), name, systems)
 end
 
 
 
-generate_rate_function(js, rate) = build_function(rate, states(js), parameters(js), 
-                                                  independent_variable(js), 
+generate_rate_function(js, rate) = build_function(rate, states(js), parameters(js),
+                                                  independent_variable(js),
                                                   expression=Val{false})
 
-generate_affect_function(js, affect, outputidxs) = build_function(affect, states(js), 
-                                                      parameters(js), 
+generate_affect_function(js, affect, outputidxs) = build_function(affect, states(js),
+                                                      parameters(js),
                                                       independent_variable(js),
                                                       expression=Val{false},
-                                                      headerfun=add_integrator_header, 
+                                                      headerfun=add_integrator_header,
                                                       outputidxs=outputidxs)[2]
-function assemble_vrj(js, vrj, statetoid) 
+function assemble_vrj(js, vrj, statetoid)
     rate   = generate_rate_function(js, vrj.rate)
     outputvars = (convert(Variable,affect.lhs) for affect in vrj.affect!)
     outputidxs = ((statetoid[var] for var in outputvars)...,)
@@ -42,18 +42,20 @@ function assemble_crj(js, crj, statetoid)
     ConstantRateJump(rate, affect)
 end
 
-function assemble_maj(js, maj::MassActionJump{U,Vector{Pair{V,W}},Vector{Pair{V2,W2}}}, 
-                      statetoid, ptoid, p, pcontext) where {U,V,W,V2,W2}
+function assemble_maj(js, maj::MassActionJump{U,Vector{Pair{V,W}},Vector{Pair{V2,W2}}},
+                      statetoid, ptoid, parammap) where {U,V,W,V2,W2}
     sr = maj.scaled_rates
-    if sr isa Operation || sr isa Variable
-        pval = Base.eval(pcontext, Expr(maj.scaled_rates))
-    else   
+    if sr isa Operation
+        pval = substitute(sr,parammap)
+    elseif sr isa Variable
+        pval = Dict(parammap)[sr()]
+    else
         pval = maj.scaled_rates
     end
-    
+
     rs = Vector{Pair{Int,W}}()
     for (spec,stoich) in maj.reactant_stoch
-        if iszero(spec) 
+        if iszero(spec)
             push!(rs, 0 => stoich)
         else
             push!(rs, statetoid[convert(Variable,spec)] => stoich)
@@ -73,13 +75,13 @@ end
 
 """
 ```julia
-function DiffEqBase.DiscreteProblem(sys::AbstractSystem, u0map, tspan, 
+function DiffEqBase.DiscreteProblem(sys::AbstractSystem, u0map, tspan,
                                     parammap=DiffEqBase.NullParameters; kwargs...)
 ```
 
 Generates a DiscreteProblem from an AbstractSystem
 """
-function DiffEqBase.DiscreteProblem(sys::AbstractSystem, u0map, tspan::Tuple, 
+function DiffEqBase.DiscreteProblem(sys::AbstractSystem, u0map, tspan::Tuple,
                                     parammap=DiffEqBase.NullParameters(); kwargs...)
     u0 = varmap_to_vars(u0map, states(sys))
     p = varmap_to_vars(parammap, parameters(sys))
@@ -99,29 +101,20 @@ function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
     majs = Vector{MassActionJump}()
     pvars = parameters(js)
     statetoid = Dict(convert(Variable,state) => i for (i,state) in enumerate(states(js)))
-    ptoid     = Dict(convert(Variable,par) => i for (i,par) in enumerate(parameters(js)))
-    
-    # for mass action jumps might need to evaluate parameter expressions
-    # populate dummy module with params as local variables
-    # (for eval-ing parameter expressions)
-    param_context = Module()
-    for (i, pval) in enumerate(prob.p)        
-        psym = Symbol(pvars[i])
-        Base.eval(param_context, :($psym = $pval))
-    end
+    parammap = map(Pair,pvars,prob.p)
 
     for j in equations(js)
         if j isa ConstantRateJump
             push!(crjs, assemble_crj(js, j, statetoid))
         elseif j isa VariableRateJump
             push!(vrjs, assemble_vrj(js, j, statetoid))
         elseif j isa MassActionJump
-            push!(majs, assemble_maj(js, j, statetoid, ptoid, prob.p, param_context))
+            push!(majs, assemble_maj(js, j, statetoid, parammap))
         else
             error("JumpSystems should only contain Constant, Variable or Mass Action Jumps.")
         end
     end
-    ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")    
+    ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")
     jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, isempty(majs) ? nothing : majs)
     JumpProblem(prob, aggregator, jset)
-end
+end

src/utils.jl

@@ -100,17 +100,23 @@ end
 
 # variable substitution
 """
-substitute_expr!(expr::Operation, s::Pair{Operation, Operation})
+substitute(expr::Operation, s::Pair)
+substitute(expr::Operation, s::Dict)
+substitute(expr::Operation, s::Vector)
 
-Performs the substitution `Operation => Operation` on the `expr` Operation.
+Performs the substitution `Operation => val` on the `expr` Operation.
 """
-substitute_expr!(expr::Constant, s::Pair{Operation, Operation}) = nothing
-function substitute_expr!(expr::Operation, s::Pair{Operation, Operation})
-    if !is_singleton(expr)
-        expr.args .= replace(expr.args, s)
-        for arg in expr.args
-            substitute_expr!(arg, s)
-        end
-    end
-    return nothing
+substitute(expr::Constant, s) = expr
+substitute(expr::Operation, s::Pair) = _substitute(expr, [s[1]], [s[2]])
+substitute(expr::Operation, dict::Dict) = _substitute(expr, keys(dict), values(dict))
+substitute(expr::Operation, s::Vector) = _substitute(expr, first.(s), last.(s))
+
+function _substitute(expr, ks, vs)
+    _substitute(expr, Dict(map(Pair, map(to_symbolic, ks), map(to_symbolic, vs))))
 end
+
+function _substitute(expr, dict::Dict)
+    to_mtk(simplify(SymbolicUtils.substitute(expr, dict)))
+end
+
+@deprecate substitute_expr!(expr,s) substitute(expr,s)

test/derivatives.jl

@@ -6,7 +6,7 @@ using Test
 @variables x y z
 @derivatives D'~t D2''~t Dx'~x
 
-test_equal(a, b) = @test isequal(simplify_constants(a), simplify_constants(b))
+test_equal(a, b) = @test isequal(simplify(a), simplify(b))
 
 @test @macroexpand(@derivatives D'~t D2''~t) == @macroexpand(@derivatives (D'~t), (D2''~t))
 
@@ -17,12 +17,12 @@ dsin = D(sin(t))
 @test isequal(expand_derivatives(dsin), cos(t))
 
 dcsch = D(csch(t))
-@test isequal(expand_derivatives(dcsch), simplify_constants(-coth(t) * csch(t)))
+@test isequal(expand_derivatives(dcsch), simplify(-coth(t) * csch(t)))
 
 @test isequal(expand_derivatives(D(-7)), 0)
-@test isequal(expand_derivatives(D(sin(2t))), simplify_constants(cos(2t) * 2))
-@test isequal(expand_derivatives(D2(sin(t))), simplify_constants(-sin(t)))
-@test isequal(expand_derivatives(D2(sin(2t))), simplify_constants(-sin(2t) * 4))
+@test isequal(expand_derivatives(D(sin(2t))), simplify(cos(2t) * 2))
+@test isequal(expand_derivatives(D2(sin(t))), simplify(-sin(t)))
+@test isequal(expand_derivatives(D2(sin(2t))), simplify(-sin(2t) * 4))
 @test isequal(expand_derivatives(D2(t)), 0)
 @test isequal(expand_derivatives(D2(5)), 0)
 
@@ -32,8 +32,8 @@ dsinsin = D(sin(sin(t)))
 
 d1 = D(sin(t)*t)
 d2 = D(sin(t)*cos(t))
-@test isequal(expand_derivatives(d1), simplify_constants(t*cos(t)+sin(t)))
-@test isequal(expand_derivatives(d2), simplify_constants(cos(t)*cos(t)+(-sin(t))*sin(t)))
+@test isequal(expand_derivatives(d1), simplify(t*cos(t)+sin(t)))
+@test isequal(expand_derivatives(d2), simplify(cos(t)*cos(t)+(-sin(t))*sin(t)))
 
 eqs = [0 ~ σ*(y-x),
        0 ~ x*(ρ-z)-y,
@@ -58,12 +58,12 @@ test_equal(jac[3,3], -1β)
 
 @variables x(t) y(t) z(t)
 
-@test isequal(expand_derivatives(D(x * y)), simplify_constants(y*D(x) + x*D(y)))
-@test isequal(expand_derivatives(D(x * y)), simplify_constants(D(x)*y + x*D(y)))
+@test isequal(expand_derivatives(D(x * y)), simplify(y*D(x) + x*D(y)))
+@test isequal(expand_derivatives(D(x * y)), simplify(D(x)*y + x*D(y)))
 
 @test isequal(expand_derivatives(D(2t)), 2)
 @test isequal(expand_derivatives(D(2x)), 2D(x))
-@test isequal(expand_derivatives(D(x^2)), simplify_constants(2 * x * D(x)))
+@test isequal(expand_derivatives(D(x^2)), simplify(2 * x * D(x)))
 
 # n-ary * and +
 isequal(ModelingToolkit.derivative(Operation(*, [x, y, z*ρ]), 1), y*(z*ρ))

test/direct.jl

@@ -2,7 +2,7 @@ using ModelingToolkit, StaticArrays, LinearAlgebra, SparseArrays
 using DiffEqBase
 using Test
 
-canonequal(a, b) = isequal(simplify_constants(a), simplify_constants(b))
+canonequal(a, b) = isequal(simplify(a), simplify(b))
 
 # Calculus
 @parameters t σ ρ β