Refactor to accomodate noise equations for SDESystems. Extended to cover Control, Nonlinear, PDE, & Discrete systems. Renamed functions to reflect what they do now.

Author: lamorton

src/systems/control/controlsystem.jl

@@ -77,6 +77,7 @@ struct ControlSystem <: AbstractControlSystem
         check_parameters(ps, iv)
         check_equations(deqs, iv)
         check_equations(observed, iv)
+        check_units(deqs)
         new(loss, deqs, iv, dvs, controls, ps, observed, name, systems, defaults)
     end
 end

src/systems/diffeqs/odesystem.jl

@@ -88,7 +88,7 @@ struct ODESystem <: AbstractODESystem
         check_variables(dvs,iv)
         check_parameters(ps,iv)
         check_equations(deqs,iv)
-        validate(deqs)
+        check_units(deqs)
         new(deqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, structure, connection_type)
     end
 end

src/systems/diffeqs/sdesystem.jl

@@ -90,7 +90,7 @@ struct SDESystem <: AbstractODESystem
         check_variables(dvs,iv)
         check_parameters(ps,iv)
         check_equations(deqs,iv)
-        validate(deqs)
+        check_units(deqs,neqs)
         new(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connection_type)
     end
 end

src/systems/diffeqs/validation.jl

@@ -1,79 +1,105 @@
 Base.:*(x::Union{Num,Symbolic},y::Unitful.AbstractQuantity) = x * y
 
-
-function vartype(x::Symbolic)
-    if !(x.metadata isa Nothing)
-        return haskey(x.metadata, VariableUnit) ? x.metadata[VariableUnit] : 1
-    end
-    1
-end
-vartype(x::Num) = vartype(value(x))
-
-instantiate(x) = 1
-instantiate(x::Unitful.Quantity) = 1 * Unitful.unit(x)
-instantiate(x::Num) = instantiate(value(x))
-function instantiate(x::Symbolic)
+"Find the units of a symbolic item."
+get_units(x) = 1
+get_units(x::Unitful.Quantity) = 1 * Unitful.unit(x)
+get_units(x::Num) = get_units(value(x))
+function get_units(x::Symbolic)
     vx = value(x)
     if vx isa Sym || operation(vx) isa Sym || (operation(vx) isa Term && operation(vx).f == getindex) || vx isa Symbolics.ArrayOp
-        return oneunit(1 * vartype(vx))
+        if x.metadata !== nothing
+            symunits = haskey(x.metadata, VariableUnit) ? x.metadata[VariableUnit] : 1
+        else
+            symunits = 1
+        end
+        return oneunit(1 * symunits)
     elseif operation(vx) isa Differential || operation(vx) isa Difference
-        return instantiate(arguments(vx)[1]) / instantiate(arguments(arguments(vx)[1])[1])
+        return get_units(arguments(vx)[1]) / get_units(arguments(arguments(vx)[1])[1])
     elseif vx isa Pow
         pargs = arguments(vx)
-        base,expon = instantiate.(pargs)
+        base,expon = get_units.(pargs)
         uconvert(NoUnits, expon) # This acts as an assertion
         return base == 1 ? 1 : operation(vx)(base, pargs[2])
     elseif vx isa Add # Cannot simply add the units b/c they may differ in magnitude (eg, kg vs g)
-        terms = instantiate.(arguments(vx))
+        terms = get_units.(arguments(vx))
         firstunit = unit(terms[1])
         @assert all(map(x -> ustrip(firstunit, x) == 1, terms[2:end]))
         return 1 * firstunit
     elseif operation(vx) == Symbolics._mapreduce 
         if vx.arguments[2] == +
-            instantiate(vx.arguments[3])
+            get_units(vx.arguments[3])
         else
             throw(ArgumentError("Unknown array operation $vx"))
         end
     else
-        return oneunit(operation(vx)(instantiate.(arguments(vx))...))
+        return oneunit(operation(vx)(get_units.(arguments(vx))...))
     end
 end
 
-function validate(eq::ModelingToolkit.Equation; eqnum = 1)
-    lhs = rhs = nothing
+"Get units of term, returning nothing & showing warning instead of throwing errors."
+function safe_get_units(term, info)
+    side = nothing
     try
-        lhs = instantiate(eq.lhs)
+        side = get_units(term)
     catch err
         if err isa Unitful.DimensionError 
-            @warn("In left-hand side of eq. #$eqnum: $(eq.lhs),  $(err.x) and $(err.y) are not dimensionally compatible.")
-        elseif err isa MethodError
-            @warn("In left-hand side of eq. #$eqnum: $(err.f) doesn't accept $(err.args).")
+            @warn("$info: $(err.x) and $(err.y) are not dimensionally compatible.")
+        elseif err isa MethodError #TODO: filter for only instances where the arguments are unitful
+            @warn("$info: no method matching $(err.f) for arguments $(err.args).")
         else
             rethrow()
         end
     end
-    try
-        rhs = instantiate(eq.rhs)
-    catch err
-        if err isa Unitful.DimensionError
-            @warn("In right-hand side of eq. #$eqnum: $(eq.rhs), $(err.x) and $(err.y) are not dimensionally compatible.")
-        elseif err isa MethodError
-            @warn("In right-hand side of eq. #$eqnum: $(err.f) doesn't accept $(err.args).")
-        else
-            rethrow()
-        end
-    end
-    if (rhs !== nothing) && (lhs !== nothing)
-        if !isequal(lhs, rhs)
-            @warn("In eq. #$eqnum, left-side units ($lhs) and right-side units ($rhs) don't match.")
+    side
+end
+
+function _validate(terms::Vector,labels::Vector; info::String = "")
+    equnits = safe_get_units.(terms,info.*labels)
+    allthere = all(map(x->x!==nothing,equnits))
+    allmatching = true
+    if allthere
+        for idx in 2:length(equnits)
+            if !isequal(equnits[1],equnits[idx])
+                allmatching = false
+                @warn("$info: units $(equnits[1]) for $(labels[1]) and $(equnits[idx]) for $(labels[idx]) do not match.")
+            end
         end
     end
-    (rhs !== nothing) && (lhs !== nothing) && isequal(lhs, rhs)
+    allthere && allmatching
+end
+
+function validate(eq::ModelingToolkit.Equation; info::String = "")
+    labels = ["left-hand side", "right-hand side"]
+    terms = [eq.lhs,eq.rhs]
+    _validate(terms,labels,info = info)
+end
+
+function validate(eq::ModelingToolkit.Equation,noiseterm; info::String = "")
+    labels = ["left-hand side", "right-hand side","noise term"]
+    terms = [eq.lhs,eq.rhs,noiseterm]
+    _validate(terms,labels,info = info)
+end
+
+function validate(eq::ModelingToolkit.Equation,noisevec::Vector; info::String = "")
+    labels = vcat(["left-hand side", "right-hand side"],"noise term #".* string.(1:length(noisevec)))
+    terms = vcat([eq.lhs,eq.rhs],noisevec)
+    _validate(terms,labels,info = info)
 end
 
 function validate(eqs::Vector{ModelingToolkit.Equation})
-    correct = [validate(eqs[idx],eqnum=idx) for idx in 1:length(eqs)]
-    all(correct) || throw(ArgumentError("Invalid equations, see warnings for details."))
+    all([validate(eqs[idx],info = "In eq. #$idx") for idx in 1:length(eqs)])
+end
+
+function validate(eqs::Vector{ModelingToolkit.Equation},noise::Vector)
+    all([validate(eqs[idx],noise[idx],info = "In eq. #$idx") for idx in 1:length(eqs)])
 end
 
-validate(sys::AbstractODESystem) = validate(equations(sys))
+function validate(eqs::Vector{ModelingToolkit.Equation},noise::Matrix)
+    all([validate(eqs[idx],noise[idx,:],info = "In eq. #$idx") for idx in 1:length(eqs)])
+end
+
+"Returns true iff units of equations are valid."
+validate(eqs::Vector) = validate(convert.(ModelingToolkit.Equation,eqs))
+
+"Throws error if units of equations are invalid."
+check_units(eqs...) = validate(eqs...) || throw(ArgumentError("Some equations had invalid units. See warnings for details."))

src/systems/discrete_system/discrete_system.jl

@@ -57,6 +57,7 @@ struct DiscreteSystem <: AbstractSystem
     function DiscreteSystem(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p)
         check_variables(dvs,iv)
         check_parameters(ps,iv)
+        check_units(discreteEqs)
         new(discreteEqs, iv, dvs, ps, var_to_name, ctrls, observed, name, systems, default_u0, default_p)
     end
 end

src/systems/nonlinear/nonlinearsystem.jl

@@ -54,6 +54,10 @@ struct NonlinearSystem <: AbstractSystem
     type: type of the system
     """
     connection_type::Any
+    function NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connection_type)
+        check_units(eqs)
+        new(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, structure, connection_type)
+    end
 end
 
 function NonlinearSystem(eqs, states, ps;

src/systems/pde/pdesystem.jl

@@ -61,6 +61,7 @@ struct PDESystem <: ModelingToolkit.AbstractSystem
                                    defaults=Dict(),
                                    connection_type=nothing,
                                   )
+        check_units(eqs)
         new(eqs, bcs, domain, indvars, depvars, ps, defaults, connection_type)
     end
 end

test/units.jl

@@ -5,41 +5,40 @@ MT = ModelingToolkit
 @variables t [unit = u"ms"] E(t) [unit = u"kJ"] P(t) [unit = u"MW"]
 D = Differential(t)
 
-@test MT.vartype(t) == u"ms"
-@test MT.vartype(E) == u"kJ"
-@test MT.vartype(τ) == u"ms"
-
-
-@test MT.instantiate(0.5) == 1.0
-@test MT.instantiate(t) == 1.0u"ms"
-@test MT.instantiate(P) == 1.0u"MW"
-@test MT.instantiate(τ) == 1.0u"ms"
-
-@test MT.instantiate(τ^-1) == 1/u"ms"
-@test MT.instantiate(D(E)) == 1.0u"MW"
-@test MT.instantiate(E/τ) == 1.0u"MW"
-@test MT.instantiate(2*P) == 1.0u"MW"
-@test MT.instantiate(t/τ) == 1.0
-@test MT.instantiate(P - E/τ)/1.0u"MW" == 1.0
-
-@test MT.instantiate(1.0^(t/τ)) == 1.0
-@test MT.instantiate(exp(t/τ)) == 1.0
-@test MT.instantiate(sin(t/τ)) == 1.0
-@test MT.instantiate(sin(1u"rad")) == 1.0
-@test MT.instantiate(t^2) == 1.0u"ms"^2
+@test MT.get_units(t) == 1u"ms"
+@test MT.get_units(E) == 1u"kJ"
+@test MT.get_units(τ) == 1u"ms"
+
+@test MT.get_units(0.5) == 1.0
+@test MT.get_units(t) == 1.0u"ms"
+@test MT.get_units(P) == 1.0u"MW"
+@test MT.get_units(τ) == 1.0u"ms"
+
+@test MT.get_units(τ^-1) == 1/u"ms"
+@test MT.get_units(D(E)) == 1.0u"MW"
+@test MT.get_units(E/τ) == 1.0u"MW"
+@test MT.get_units(2*P) == 1.0u"MW"
+@test MT.get_units(t/τ) == 1.0
+@test MT.get_units(P - E/τ)/1.0u"MW" == 1.0
+
+@test MT.get_units(1.0^(t/τ)) == 1.0
+@test MT.get_units(exp(t/τ)) == 1.0
+@test MT.get_units(sin(t/τ)) == 1.0
+@test MT.get_units(sin(1u"rad")) == 1.0
+@test MT.get_units(t^2) == 1.0u"ms"^2
 
 @test !MT.validate(E^1.5 ~ E^(t/τ))
 @test MT.validate(E^(t/τ) ~ E^(t/τ))
 
 eqs = [D(E) ~ P - E/τ
         0.0u"MW" ~ P]
-@test MT.instantiate(eqs[1].lhs) == 1.0u"MW"
-@test MT.instantiate(eqs[1].rhs) == 1.0u"MW"
+@test MT.get_units(eqs[1].lhs) == 1.0u"MW"
+@test MT.get_units(eqs[1].rhs) == 1.0u"MW"
 @test MT.validate(eqs[1])
 @test MT.validate(eqs[2])
+@test MT.validate(eqs)
 sys = ODESystem(eqs)
 sys = ODESystem(eqs, t, [P, E], [τ])
-@test MT.validate(sys)
 
 @test !MT.validate(D(D(E)) ~ P)
 @test !MT.validate(0 ~ P + E*τ)
@@ -112,4 +111,31 @@ eqs = [D(y[1]) ~ -k[1]*y[1] + k[3]*y[2]*y[3],
        D(y[2]) ~  k[1]*y[1] - k[3]*y[2]*y[3] - k[2]*y[2]^2,
        0 ~  y[1] + y[2] + y[3] - 1]
 
-sys = ODESystem(eqs,t,y,k)
+sys = ODESystem(eqs,t,y,k)
+
+# Nonlinear system
+@parameters a [unit = u"kg"^-1]
+@variables x [unit = u"kg"]
+eqs = [
+    0 ~ a*x 
+]
+nls = NonlinearSystem(eqs, [x], [a])
+
+# SDE test w/ noise vector
+@parameters τ [unit = u"ms"] Q [unit = u"MW"]
+@variables t [unit = u"ms"] E(t) [unit = u"kJ"] P(t) [unit = u"MW"]
+D = Differential(t)
+eqs = [D(E) ~ P - E/τ
+      P ~ Q]
+
+noiseeqs = [0.1u"MW",
+            0.1u"MW"]
+sys = SDESystem(eqs,noiseeqs,t,[P,E],[τ,Q])
+# With noise matrix
+noiseeqs = [0.1u"MW" 0.1u"MW"
+            0.1u"MW" 0.1u"MW"]
+sys = SDESystem(eqs,noiseeqs,t,[P,E],[τ,Q])
+
+noiseeqs = [0.1u"MW" 0.1u"MW"
+            0.1u"MW" 0.1u"s"]
+@test !MT.validate(eqs,noiseeqs)