Unityper update

YingboMa · YingboMa · commit 0681e7517442 · 2023-01-09T12:22:25.000-05:00
diff --git a/Project.toml b/Project.toml
@@ -78,7 +78,7 @@ SimpleNonlinearSolve = "0.1.0"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 SymbolicIndexingInterface = "0.1, 0.2"
-SymbolicUtils = "0.19"
+SymbolicUtils = "0.19, 0.20"
 Symbolics = "4.9"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
diff --git a/src/ModelingToolkit.jl b/src/ModelingToolkit.jl
@@ -39,8 +39,9 @@ import SymbolicIndexingInterface: independent_variables, states, parameters
 export independent_variables, states, parameters
 import SymbolicUtils
 import SymbolicUtils: istree, arguments, operation, similarterm, promote_symtype,
-                      Symbolic, Term, Add, Mul, Pow, Sym, FnType,
-                      @rule, Rewriters, substitute, metadata
+                      Symbolic, isadd, ismul, ispow, issym, FnType,
+                      @rule, Rewriters, substitute, metadata, BasicSymbolic,
+                      Sym, Term
 using SymbolicUtils.Code
 import SymbolicUtils.Code: toexpr
 import SymbolicUtils.Rewriters: Chain, Postwalk, Prewalk, Fixpoint
diff --git a/src/clock.jl b/src/clock.jl
@@ -13,21 +13,14 @@ const InferredDomain = Union{Inferred, InferredDiscrete}
 Symbolics.option_to_metadata_type(::Val{:timedomain}) = TimeDomain
 
 """
-    is_continuous_domain(x::Sym)
-Determine if variable `x` is a continuous-time variable.
-"""
-is_continuous_domain(x::Sym) = getmetadata(x, TimeDomain, false) isa Continuous
-
-"""
-    is_discrete_domain(x::Sym)
-Determine if variable `x` is a discrete-time variable.
+    is_continuous_domain(x)
+true if `x` contains only continuous-domain signals.
+See also [`has_continuous_domain`](@ref)
 """
-is_discrete_domain(x::Sym) = getmetadata(x, TimeDomain, false) isa Discrete
-
-# is_discrete_domain(x::Sym) = isvarkind(Discrete, x)
-
-has_continuous_domain(x::Sym) = is_continuous_domain(x)
-has_discrete_domain(x::Sym) = is_discrete_domain(x)
+function is_continuous_domain(x)
+    issym(x) && return getmetadata(x, TimeDomain, false) isa Continuous
+    !has_discrete_domain(x) && has_continuous_domain(x)
+end
 
 function get_time_domain(x)
     if istree(x) && operation(x) isa Operator
@@ -39,10 +32,10 @@ end
 get_time_domain(x::Num) = get_time_domain(value(x))
 
 """
-    has_time_domain(x::Sym)
+    has_time_domain(x)
 Determine if variable `x` has a time-domain attributed to it.
 """
-function has_time_domain(x::Union{Sym, Term})
+function has_time_domain(x::Symbolic)
     # getmetadata(x, Continuous, nothing) !== nothing ||
     # getmetadata(x, Discrete,   nothing) !== nothing
     getmetadata(x, TimeDomain, nothing) !== nothing
@@ -60,14 +53,20 @@ end
 true if `x` contains discrete signals (`x` may or may not contain continuous-domain signals). `x` may be an expression or equation.
 See also [`is_discrete_domain`](@ref)
 """
-has_discrete_domain(x) = hasshift(x) || hassample(x) || hashold(x)
+function has_discrete_domain(x)
+    issym(x) && return is_discrete_domain(x)
+    hasshift(x) || hassample(x) || hashold(x)
+end
 
 """
     has_continuous_domain(x)
 true if `x` contains continuous signals (`x` may or may not contain discrete-domain signals). `x` may be an expression or equation.
 See also [`is_continuous_domain`](@ref)
 """
-has_continuous_domain(x) = hasderiv(x) || hasdiff(x) || hassample(x) || hashold(x)
+function has_continuous_domain(x)
+    issym(x) && return is_continuous_domain(x)
+    hasderiv(x) || hasdiff(x) || hassample(x) || hashold(x)
+end
 
 """
     is_hybrid_domain(x)
@@ -80,14 +79,10 @@ is_hybrid_domain(x) = has_discrete_domain(x) && has_continuous_domain(x)
 true if `x` contains only discrete-domain signals.
 See also [`has_discrete_domain`](@ref)
 """
-is_discrete_domain(x) = has_discrete_domain(x) && !has_continuous_domain(x)
-
-"""
-    is_continuous_domain(x)
-true if `x` contains only continuous-domain signals.
-See also [`has_continuous_domain`](@ref)
-"""
-is_continuous_domain(x) = !has_discrete_domain(x) && has_continuous_domain(x)
+function is_discrete_domain(x)
+    issym(x) && return getmetadata(x, TimeDomain, false) isa Discrete
+    has_discrete_domain(x) && !has_continuous_domain(x)
+end
 
 struct ClockInferenceException <: Exception
     msg::Any
diff --git a/src/constants.jl b/src/constants.jl
@@ -9,11 +9,11 @@ function isconstant(x)
 end
 
 """
-    toconstant(s::Sym)
+    toconstant(s)
 
 Maps the parameter to a constant. The parameter must have a default.
 """
-function toconstant(s::Sym)
+function toconstant(s)
     hasmetadata(s, Symbolics.VariableDefaultValue) ||
         throw(ArgumentError("Constant `$(s)` must be assigned a default value."))
     setmetadata(s, MTKConstantCtx, true)
diff --git a/src/inputoutput.jl b/src/inputoutput.jl
@@ -19,8 +19,8 @@ function outputs(sys)
     lhss = [eq.lhs for eq in o]
     unique([filter(isoutput, states(sys))
             filter(isoutput, parameters(sys))
-            filter(x -> x isa Term && isoutput(x), rhss) # observed can return equations with complicated expressions, we are only looking for single Terms
-            filter(x -> x isa Term && isoutput(x), lhss)])
+            filter(x -> istree(x) && isoutput(x), rhss) # observed can return equations with complicated expressions, we are only looking for single Terms
+            filter(x -> istree(x) && isoutput(x), lhss)])
 end
 
 """
diff --git a/src/parameters.jl b/src/parameters.jl
@@ -25,7 +25,7 @@ function isparameter(x)
 end
 
 """
-    toparam(s::Sym)
+    toparam(s)
 
 Maps the variable to a paramter.
 """
@@ -41,7 +41,7 @@ end
 toparam(s::Num) = wrap(toparam(value(s)))
 
 """
-    tovar(s::Sym)
+    tovar(s)
 
 Maps the variable to a state.
 """
diff --git a/src/systems/abstractsystem.jl b/src/systems/abstractsystem.jl
@@ -572,7 +572,7 @@ function time_varying_as_func(x, sys::AbstractTimeDependentSystem)
     #
     # This is done by just making `x` the argument of the function.
     if istree(x) &&
-       operation(x) isa Sym &&
+       issym(operation(x)) &&
        !(length(arguments(x)) == 1 && isequal(arguments(x)[1], get_iv(sys)))
         return operation(x)
     end
@@ -616,7 +616,7 @@ AbstractSysToExpr(sys) = AbstractSysToExpr(sys, states(sys))
 function (f::AbstractSysToExpr)(O)
     !istree(O) && return toexpr(O)
     any(isequal(O), f.states) && return nameof(operation(O))  # variables
-    if isa(operation(O), Sym)
+    if issym(operation(O))
         return build_expr(:call, Any[nameof(operation(O)); f.(arguments(O))])
     end
     return build_expr(:call, Any[operation(O); f.(arguments(O))])
@@ -645,7 +645,7 @@ end
 function round_trip_expr(t, var2name)
     name = get(var2name, t, nothing)
     name !== nothing && return name
-    t isa Sym && return nameof(t)
+    issym(t) && return nameof(t)
     istree(t) || return t
     f = round_trip_expr(operation(t), var2name)
     args = map(Base.Fix2(round_trip_expr, var2name), arguments(t))
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -200,7 +200,7 @@ function calculate_massmatrix(sys::AbstractODESystem; simplify = false)
     M = zeros(length(eqs), length(eqs))
     state2idx = Dict(s => i for (i, s) in enumerate(dvs))
     for (i, eq) in enumerate(eqs)
-        if eq.lhs isa Term && operation(eq.lhs) isa Differential
+        if istree(eq.lhs) && operation(eq.lhs) isa Differential
             st = var_from_nested_derivative(eq.lhs)[1]
             j = state2idx[st]
             M[i, j] = 1
diff --git a/src/systems/diffeqs/odesystem.jl b/src/systems/diffeqs/odesystem.jl
@@ -31,7 +31,7 @@ struct ODESystem <: AbstractODESystem
     """The ODEs defining the system."""
     eqs::Vector{Equation}
     """Independent variable."""
-    iv::Sym
+    iv::BasicSymbolic{Real}
     """
     Dependent (state) variables. Must not contain the independent variable.
 
diff --git a/src/systems/diffeqs/sdesystem.jl b/src/systems/diffeqs/sdesystem.jl
@@ -37,7 +37,7 @@ struct SDESystem <: AbstractODESystem
     """The expressions defining the diffusion term."""
     noiseeqs::AbstractArray
     """Independent variable."""
-    iv::Sym
+    iv::BasicSymbolic{Real}
     """Dependent (state) variables. Must not contain the independent variable."""
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
diff --git a/src/systems/discrete_system/discrete_system.jl b/src/systems/discrete_system/discrete_system.jl
@@ -32,7 +32,7 @@ struct DiscreteSystem <: AbstractTimeDependentSystem
     """The differential equations defining the discrete system."""
     eqs::Vector{Equation}
     """Independent variable."""
-    iv::Sym
+    iv::BasicSymbolic{Real}
     """Dependent (state) variables. Must not contain the independent variable."""
     states::Vector
     """Parameter variables. Must not contain the independent variable."""
diff --git a/src/systems/validation.jl b/src/systems/validation.jl
@@ -32,7 +32,6 @@ equivalent(x, y) = isequal(1 * x, 1 * y)
 const unitless = Unitful.unit(1)
 
 #For dispatching get_unit
-const Literal = Union{Sym, Symbolics.ArrayOp, Symbolics.Arr, Symbolics.CallWithMetadata}
 const Conditional = Union{typeof(ifelse), typeof(IfElse.ifelse)}
 const Comparison = Union{typeof.([==, !=, ≠, <, <=, ≤, >, >=, ≥])...}
 
@@ -41,13 +40,15 @@ get_unit(x::Real) = unitless
 get_unit(x::Unitful.Quantity) = screen_unit(Unitful.unit(x))
 get_unit(x::AbstractArray) = map(get_unit, x)
 get_unit(x::Num) = get_unit(value(x))
-get_unit(x::Literal) = screen_unit(getmetadata(x, VariableUnit, unitless))
+get_unit(x::Union{Symbolics.ArrayOp, Symbolics.Arr, Symbolics.CallWithMetadata}) = get_literal_unit(x)
 get_unit(op::Differential, args) = get_unit(args[1]) / get_unit(op.x)
 get_unit(op::Difference, args) = get_unit(args[1]) / get_unit(op.t)
 get_unit(op::typeof(getindex), args) = get_unit(args[1])
 get_unit(x::SciMLBase.NullParameters) = unitless
 get_unit(op::typeof(instream), args) = get_unit(args[1])
 
+get_literal_unit(x) = screen_unit(getmetadata(x, VariableUnit, unitless))
+
 function get_unit(op, args) # Fallback
     result = op(1 .* get_unit.(args)...)
     try
@@ -69,28 +70,6 @@ function get_unit(op::Integral, args)
     return get_unit(args[1]) * unit
 end
 
-function get_unit(x::Pow)
-    pargs = arguments(x)
-    base, expon = get_unit.(pargs)
-    @assert expon isa Unitful.DimensionlessUnits
-    if base == unitless
-        unitless
-    else
-        pargs[2] isa Number ? base^pargs[2] : (1 * base)^pargs[2]
-    end
-end
-
-function get_unit(x::Add)
-    terms = get_unit.(arguments(x))
-    firstunit = terms[1]
-    for other in terms[2:end]
-        termlist = join(map(repr, terms), ", ")
-        equivalent(other, firstunit) ||
-            throw(ValidationError(", in sum $x, units [$termlist] do not match."))
-    end
-    return firstunit
-end
-
 function get_unit(op::Conditional, args)
     terms = get_unit.(args)
     terms[1] == unitless ||
@@ -116,11 +95,31 @@ function get_unit(op::Comparison, args)
 end
 
 function get_unit(x::Symbolic)
-    if SymbolicUtils.istree(x)
+    if issym(x)
+        get_literal_unit(x)
+    elseif isadd(x)
+        terms = get_unit.(arguments(x))
+        firstunit = terms[1]
+        for other in terms[2:end]
+            termlist = join(map(repr, terms), ", ")
+            equivalent(other, firstunit) ||
+                throw(ValidationError(", in sum $x, units [$termlist] do not match."))
+        end
+        return firstunit
+    elseif ispow(x)
+        pargs = arguments(x)
+        base, expon = get_unit.(pargs)
+        @assert expon isa Unitful.DimensionlessUnits
+        if base == unitless
+            unitless
+        else
+            pargs[2] isa Number ? base^pargs[2] : (1 * base)^pargs[2]
+        end
+    elseif istree(x)
         op = operation(x)
-        if op isa Sym || (op isa Term && operation(op) isa Term) # Dependent variables, not function calls
+        if issym(op) || (istree(op) && istree(operation(op))) # Dependent variables, not function calls
             return screen_unit(getmetadata(x, VariableUnit, unitless)) # Like x(t) or x[i]
-        elseif op isa Term && !(operation(op) isa Term)
+        elseif istree(op) && !istree(operation(op))
             gp = getmetadata(x, Symbolics.GetindexParent, nothing) # Like x[1](t)
             return screen_unit(getmetadata(gp, VariableUnit, unitless))
         end  # Actual function calls:
diff --git a/src/utils.jl b/src/utils.jl
@@ -26,11 +26,8 @@ function detime_dvs(op)
     end
 end
 
-function retime_dvs(op::Sym, dvs, iv)
-    Sym{FnType{Tuple{symtype(iv)}, Real}}(nameof(op))(iv)
-end
-
 function retime_dvs(op, dvs, iv)
+    issym(op) && return Sym{FnType{Tuple{symtype(iv)}, Real}}(nameof(op))(iv)
     istree(op) ?
     similarterm(op, operation(op), retime_dvs.(arguments(op), (dvs,), (iv,))) :
     op
@@ -337,19 +334,18 @@ isdiffeq(eq) = isdifferential(eq.lhs)
 isdifference(expr) = isoperator(expr, Difference)
 isdifferenceeq(eq) = isdifference(eq.lhs)
 
-function iv_from_nested_difference(x::Term)
+function iv_from_nested_difference(x::Symbolic)
+    istree(x) || return x
     operation(x) isa Difference ? iv_from_nested_difference(arguments(x)[1]) :
     arguments(x)[1]
 end
-iv_from_nested_difference(x::Sym) = x
 iv_from_nested_difference(x) = nothing
 
 var_from_nested_difference(x, i = 0) = (nothing, nothing)
-function var_from_nested_difference(x::Term, i = 0)
-    operation(x) isa Difference ? var_from_nested_difference(arguments(x)[1], i + 1) :
+function var_from_nested_difference(x::Symbolic, i = 0)
+    istree(x) && operation(x) isa Difference ? var_from_nested_difference(arguments(x)[1], i + 1) :
     (x, i)
 end
-var_from_nested_difference(x::Sym, i = 0) = (x, i)
 
 isvariable(x::Num) = isvariable(value(x))
 function isvariable(x)
@@ -373,8 +369,7 @@ v  = ModelingToolkit.vars(D(y) ~ u)
 v == Set([D(y), u])
 ```
 """
-vars(x::Sym; op = Differential) = Set([x])
-vars(exprs::Symbolic; op = Differential) = vars([exprs]; op = op)
+vars(exprs::Symbolic; op = Differential) = istree(exprs) ? vars([exprs]; op = op) : Set([exprs])
 vars(exprs; op = Differential) = foldl((x, y) -> vars!(x, y; op = op), exprs; init = Set())
 vars(eq::Equation; op = Differential) = vars!(Set(), eq; op = op)
 function vars!(vars, eq::Equation; op = Differential)
@@ -448,7 +443,7 @@ The difference compared to `collect_operator_variables` is that `collect_operato
 function collect_applied_operators(x, op)
     v = vars(x, op = op)
     filter(v) do x
-        x isa Sym && return false
+        issym(x) && return false
         istree(x) && return operation(x) isa op
         false
     end
@@ -467,7 +462,7 @@ function find_derivatives!(vars, expr, f)
 end
 
 function collect_vars!(states, parameters, expr, iv)
-    if expr isa Sym
+    if issym(expr)
         collect_var!(states, parameters, expr, iv)
     else
         for var in vars(expr)
@@ -481,7 +476,7 @@ function collect_vars!(states, parameters, expr, iv)
 end
 
 function collect_vars_difference!(states, parameters, expr, iv)
-    if expr isa Sym
+    if issym(expr)
         collect_var!(states, parameters, expr, iv)
     else
         for var in vars(expr)
@@ -506,7 +501,7 @@ end
 
 """ Find all the symbolic constants of some equations or terms and return them as a vector. """
 function collect_constants(x)
-    constants = Symbolics.Sym[]
+    constants = BasicSymbolic[]
     collect_constants!(constants, x)
     return constants
 end
@@ -529,9 +524,9 @@ end
 
 collect_constants!(constants, x::Num) = collect_constants!(constants, unwrap(x))
 collect_constants!(constants, x::Real) = nothing
-collect_constants(n::Nothing) = Symbolics.Sym[]
+collect_constants(n::Nothing) = BasicSymbolic[]
 
-function collect_constants!(constants, expr::Symbolics.Symbolic)
+function collect_constants!(constants, expr::Symbolic)
     if issym(expr) && isconstant(expr)
         push!(constants, expr)
     else
diff --git a/test/odesystem.jl b/test/odesystem.jl
