Merge pull request #2020 from SciML/myb/uni

YingboMa · web-flow · commit fff670e888e2 · 2023-02-08T13:19:46.000-05:00
Unityper update
diff --git a/Project.toml b/Project.toml
@@ -78,8 +78,8 @@ 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"
-Symbolics = "4.9"
+SymbolicUtils = "1.0"
+Symbolics = "5.0"
 UnPack = "0.1, 1.0"
 Unitful = "1.1"
 julia = "1.6"
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,23 +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
@@ -41,11 +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
@@ -64,15 +54,21 @@ 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)
@@ -87,15 +83,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
@@ -122,3 +113,4 @@ struct Clock <: AbstractClock
 end
 
 sampletime(c) = isdefined(c, :dt) ? c.dt : nothing
+Base.:(==)(c1::Clock, c2::Clock) = isequal(c1.t, c2.t) && c1.dt == c2.dt
diff --git a/src/constants.jl b/src/constants.jl
@@ -11,11 +11,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
@@ -27,7 +27,7 @@ function isparameter(x)
 end
 
 """
-    toparam(s::Sym)
+    toparam(s)
 
 Maps the variable to a paramter.
 """
@@ -43,7 +43,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
@@ -419,25 +419,27 @@ function renamespace(sys, x)
     sys === nothing && return x
     x = unwrap(x)
     if x isa Symbolic
+        T = typeof(x)
         if istree(x) && operation(x) isa Operator
-            return similarterm(x, operation(x), Any[renamespace(sys, only(arguments(x)))])
+            return similarterm(x, operation(x),
+                               Any[renamespace(sys, only(arguments(x)))])::T
         end
         let scope = getmetadata(x, SymScope, LocalScope())
             if scope isa LocalScope
-                rename(x, renamespace(getname(sys), getname(x)))
+                rename(x, renamespace(getname(sys), getname(x)))::T
             elseif scope isa ParentScope
-                setmetadata(x, SymScope, scope.parent)
+                setmetadata(x, SymScope, scope.parent)::T
             elseif scope isa DelayParentScope
                 if scope.N > 0
                     x = setmetadata(x, SymScope,
                                     DelayParentScope(scope.parent, scope.N - 1))
-                    rename(x, renamespace(getname(sys), getname(x)))
+                    rename(x, renamespace(getname(sys), getname(x)))::T
                 else
                     #rename(x, renamespace(getname(sys), getname(x)))
-                    setmetadata(x, SymScope, scope.parent)
+                    setmetadata(x, SymScope, scope.parent)::T
                 end
             else # GlobalScope
-                x
+                x::T
             end
         end
     elseif x isa AbstractSystem
@@ -453,9 +455,8 @@ namespace_controls(sys::AbstractSystem) = controls(sys, controls(sys))
 
 function namespace_defaults(sys)
     defs = defaults(sys)
-    Dict((isparameter(k) ? parameters(sys, k) : states(sys, k)) => namespace_expr(defs[k],
-                                                                                  sys)
-         for k in keys(defs))
+    Dict((isparameter(k) ? parameters(sys, k) : states(sys, k)) => namespace_expr(v, sys)
+         for (k, v) in pairs(defs))
 end
 
 function namespace_equations(sys::AbstractSystem)
@@ -484,14 +485,19 @@ function namespace_expr(O, sys, n = nameof(sys))
     elseif isvariable(O)
         renamespace(n, O)
     elseif istree(O)
-        renamed = map(a -> namespace_expr(a, sys, n), arguments(O))
+        T = typeof(O)
         if symtype(operation(O)) <: FnType
-            renamespace(n, O)
+            renamespace(n, O)::T
         else
-            similarterm(O, operation(O), renamed)
+            renamed = let sys = sys, n = n, T = T
+                map(a -> namespace_expr(a, sys, n)::Any, arguments(O))
+            end
+            similarterm(O, operation(O), renamed)::T
         end
     elseif O isa Array
-        map(o -> namespace_expr(o, sys, n), O)
+        let sys = sys, n = n
+            map(o -> namespace_expr(o, sys, n), O)
+        end
     else
         O
     end
@@ -602,7 +608,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
@@ -646,7 +652,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))])
@@ -675,7 +681,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.
 
@@ -422,7 +422,7 @@ function convert_system(::Type{<:ODESystem}, sys, t; name = nameof(sys))
                 newsts[i] = s
                 continue
             end
-            ns = similarterm(s, operation(s), (t,); metadata = SymbolicUtils.metadata(s))
+            ns = similarterm(s, operation(s), Any[t]; metadata = SymbolicUtils.metadata(s))
             newsts[i] = ns
             varmap[s] = ns
         else
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
@@ -36,7 +36,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.([==, !=, ≠, <, <=, ≤, >, >=, ≥])...}
 
@@ -47,13 +46,17 @@ 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))
+function get_unit(x::Union{Symbolics.ArrayOp, Symbolics.Arr, Symbolics.CallWithMetadata})
+    get_literal_unit(x)
+end
 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
@@ -75,28 +78,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 ||
@@ -122,11 +103,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
diff --git a/test/odesystem.jl b/test/odesystem.jl
diff --git a/test/simplify.jl b/test/simplify.jl
