Move Unitful checks to a separate module

Author: YingboMa

src/systems/diffeqs/sdesystem.jl

@@ -134,7 +134,8 @@ struct SDESystem <: AbstractODESystem
             check_equations(equations(cevents), iv)
         end
         if checks == true || (checks & CheckUnits) > 0
-            all_dimensionless([dvs; ps; iv]) || check_units(deqs, neqs)
+            u = __get_unit_type(dvs, ps, iv)
+            check_units(u, deqs, neqs)
         end
         new(tag, deqs, neqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac,

src/systems/jumps/jumpsystem.jl

@@ -112,7 +112,8 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
             check_parameters(ps, iv)
         end
         if checks == true || (checks & CheckUnits) > 0
-            all_dimensionless([states; ps; iv]) || check_units(ap, iv)
+            u = __get_unit_type(states, ps, iv)
+            check_units(u, ap, iv)
         end
         new{U}(tag, ap, iv, states, ps, var_to_name, observed, name, systems, defaults,
             connector_type, devents, metadata, gui_metadata, complete)

src/systems/nonlinear/nonlinearsystem.jl

@@ -88,7 +88,8 @@ struct NonlinearSystem <: AbstractTimeIndependentSystem
             tearing_state = nothing, substitutions = nothing,
             complete = false, parent = nothing; checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckUnits) > 0
-            all_dimensionless([states; ps]) || check_units(eqs)
+            u = __get_unit_type(states, ps)
+            check_units(u, eqs)
         end
         new(tag, eqs, states, ps, var_to_name, observed, jac, name, systems, defaults,
             connector_type, metadata, gui_metadata, tearing_state, substitutions, complete,

src/systems/optimization/constraints_system.jl

@@ -77,7 +77,8 @@ struct ConstraintsSystem <: AbstractTimeIndependentSystem
             tearing_state = nothing, substitutions = nothing;
             checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckUnits) > 0
-            all_dimensionless([states; ps]) || check_units(constraints)
+            u = __get_unit_type(states, ps)
+            check_units(u, constraints)
         end
         new(tag, constraints, states, ps, var_to_name, observed, jac, name, systems,
             defaults,

src/systems/optimization/optimizationsystem.jl

@@ -68,7 +68,8 @@ struct OptimizationSystem <: AbstractOptimizationSystem
         if checks == true || (checks & CheckUnits) > 0
             unwrap(op) isa Symbolic && check_units(op)
             check_units(observed)
-            all_dimensionless([states; ps]) || check_units(constraints)
+            u = __get_unit_type(states, ps)
+            check_units(u, constraints)
         end
         new(tag, op, states, ps, var_to_name, observed,
             constraints, name, systems, defaults, metadata, gui_metadata, complete,

src/systems/pde/pdesystem.jl

@@ -98,7 +98,8 @@ struct PDESystem <: ModelingToolkit.AbstractMultivariateSystem
             checks::Union{Bool, Int} = true,
             name)
         if checks == true || (checks & CheckUnits) > 0
-            all_dimensionless([dvs; ivs; ps]) || check_units(eqs)
+            u = __get_unit_type(dvs, ivs, ps)
+            check_units(u, deqs)
         end
 
         eqs = eqs isa Vector ? eqs : [eqs]

src/systems/unit_check.jl

@@ -1,4 +1,4 @@
-import DynamicQuantities
+import DynamicQuantities, Unitful
 const DQ = DynamicQuantities
 
 struct ValidationError <: Exception
@@ -8,14 +8,26 @@ end
 check_units(::Nothing, _...) = true
 
 __get_literal_unit(x) = getmetadata(x, VariableUnit, nothing)
+function __get_scalar_unit_type(v)
+    u = __get_literal_unit(v)
+    if u isa DQ.AbstractQuantity
+        return Val(:DynamicQuantities)
+    elseif u isa Unitful.Unitlike
+        return Val(:Unitful)
+    end
+    return nothing
+end
 function __get_unit_type(vs′...)
-    vs = Iterators.flatten(vs′)
-    for v in vs
-        u = __get_literal_unit(v)
-        if u isa DQ.AbstractQuantity
-            return Val(:Unitful)
+    for vs in vs′
+        if vs isa AbstractVector
+            for v in vs
+                u = __get_scalar_unit_type(v)
+                u === nothing || return u
+            end
         else
-            return Val(:DynamicQuantities)
+            v = vs
+            u = __get_scalar_unit_type(v)
+            u === nothing || return u
         end
     end
     return nothing
@@ -38,5 +50,3 @@ function screen_units(result)
         end
     end
 end
-
-

src/systems/validation.jl

@@ -1,7 +1,9 @@
 module UnitfulUnitCheck
 
-using .ModelingToolkit, Symbolics, SciMLBase
-using .ModelingToolkit: ValidationError
+using ..ModelingToolkit, Symbolics, SciMLBase, Unitful, IfElse, RecursiveArrayTools
+using ..ModelingToolkit: ValidationError,
+    ModelingToolkit, Connection, instream, JumpType, VariableUnit, get_systems
+using Symbolics: Symbolic, value, issym, isadd, ismul, ispow
 const MT = ModelingToolkit
 
 Base.:*(x::Union{Num, Symbolic}, y::Unitful.AbstractQuantity) = x * y
@@ -62,11 +64,7 @@ get_literal_unit(x) = screen_unit(getmetadata(x, VariableUnit, unitless))
 function get_unit(op, args) # Fallback
     result = op(1 .* get_unit.(args)...)
     try
-        if result isa DQ.AbstractQuantity
-            oneunit(result)
-        else
-            unit(result)
-        end
+        unit(result)
     catch
         throw(ValidationError("Unable to get unit for operation $op with arguments $args."))
     end
@@ -283,6 +281,5 @@ function MT.check_units(::Val{:Unitful}, eqs...)
     validate(eqs...) ||
         throw(ValidationError("Some equations had invalid units. See warnings for details."))
 end
-all_dimensionless(states) = all(x -> safe_get_unit(x, "") in (unitless, nothing), states)
 
 end # module

test/runtests.jl

@@ -13,10 +13,10 @@ using SafeTestsets, Test
 @safetestset "Clock Test" include("clock.jl")
 @safetestset "DiscreteSystem Test" include("discretesystem.jl")
 @safetestset "ODESystem Test" include("odesystem.jl")
-@safetestset "Dynamic Quantities Test" begin
-    using DynamicQuantities
-    include("units.jl")
-end
+#@safetestset "Dynamic Quantities Test" begin
+#    using DynamicQuantities
+#    include("units.jl")
+#end
 @safetestset "Unitful Quantities Test" begin
     using Unitful
     include("units.jl")

test/units.jl

@@ -1,51 +1,52 @@
 using ModelingToolkit, OrdinaryDiffEq, JumpProcesses, IfElse
 using Test
 MT = ModelingToolkit
+UMT = ModelingToolkit.UnitfulUnitCheck
 @parameters τ [unit = u"ms"] γ
 @variables t [unit = u"ms"] E(t) [unit = u"kJ"] P(t) [unit = u"MW"]
 D = Differential(t)
 
 #This is how equivalent works:
-@test MT.equivalent(u"MW", u"kJ/ms")
-@test !MT.equivalent(u"m", u"cm")
-@test MT.equivalent(MT.get_unit(P^γ), MT.get_unit((E / τ)^γ))
+@test UMT.equivalent(u"MW", u"kJ/ms")
+@test !UMT.equivalent(u"m", u"cm")
+@test UMT.equivalent(UMT.get_unit(P^γ), UMT.get_unit((E / τ)^γ))
 
 # Basic access
-@test MT.get_unit(t) == u"ms"
-@test MT.get_unit(E) == u"kJ"
-@test MT.get_unit(τ) == u"ms"
-@test MT.get_unit(γ) == MT.unitless
-@test MT.get_unit(0.5) == MT.unitless
-@test MT.get_unit(MT.SciMLBase.NullParameters()) == MT.unitless
+@test UMT.get_unit(t) == u"ms"
+@test UMT.get_unit(E) == u"kJ"
+@test UMT.get_unit(τ) == u"ms"
+@test UMT.get_unit(γ) == UMT.unitless
+@test UMT.get_unit(0.5) == UMT.unitless
+@test UMT.get_unit(UMT.SciMLBase.NullParameters()) == UMT.unitless
 
 # Prohibited unit types
 @parameters β [unit = u"°"] α [unit = u"°C"] γ [unit = 1u"s"]
-@test_throws MT.ValidationError MT.get_unit(β)
-@test_throws MT.ValidationError MT.get_unit(α)
-@test_throws MT.ValidationError MT.get_unit(γ)
+@test_throws UMT.ValidationError UMT.get_unit(β)
+@test_throws UMT.ValidationError UMT.get_unit(α)
+@test_throws UMT.ValidationError UMT.get_unit(γ)
 
 # Non-trivial equivalence & operators
-@test MT.get_unit(τ^-1) == u"ms^-1"
-@test MT.equivalent(MT.get_unit(D(E)), u"MW")
-@test MT.equivalent(MT.get_unit(E / τ), u"MW")
-@test MT.get_unit(2 * P) == u"MW"
-@test MT.get_unit(t / τ) == MT.unitless
-@test MT.equivalent(MT.get_unit(P - E / τ), u"MW")
-@test MT.equivalent(MT.get_unit(D(D(E))), u"MW/ms")
-@test MT.get_unit(IfElse.ifelse(t > t, P, E / τ)) == u"MW"
-@test MT.get_unit(1.0^(t / τ)) == MT.unitless
-@test MT.get_unit(exp(t / τ)) == MT.unitless
-@test MT.get_unit(sin(t / τ)) == MT.unitless
-@test MT.get_unit(sin(1u"rad")) == MT.unitless
-@test MT.get_unit(t^2) == u"ms^2"
+@test UMT.get_unit(τ^-1) == u"ms^-1"
+@test UMT.equivalent(UMT.get_unit(D(E)), u"MW")
+@test UMT.equivalent(UMT.get_unit(E / τ), u"MW")
+@test UMT.get_unit(2 * P) == u"MW"
+@test UMT.get_unit(t / τ) == UMT.unitless
+@test UMT.equivalent(UMT.get_unit(P - E / τ), u"MW")
+@test UMT.equivalent(UMT.get_unit(D(D(E))), u"MW/ms")
+@test UMT.get_unit(IfElse.ifelse(t > t, P, E / τ)) == u"MW"
+@test UMT.get_unit(1.0^(t / τ)) == UMT.unitless
+@test UMT.get_unit(exp(t / τ)) == UMT.unitless
+@test UMT.get_unit(sin(t / τ)) == UMT.unitless
+@test UMT.get_unit(sin(1u"rad")) == UMT.unitless
+@test UMT.get_unit(t^2) == u"ms^2"
 
 eqs = [D(E) ~ P - E / τ
     0 ~ P]
-@test MT.validate(eqs)
+@test UMT.validate(eqs)
 @named sys = ODESystem(eqs)
 
-@test !MT.validate(D(D(E)) ~ P)
-@test !MT.validate(0 ~ P + E * τ)
+@test !UMT.validate(D(D(E)) ~ P)
+@test !UMT.validate(0 ~ P + E * τ)
 
 # Disabling unit validation/checks selectively
 @test_throws MT.ArgumentError ODESystem(eqs, t, [E, P, t], [τ], name = :sys)
@@ -86,9 +87,9 @@ end
 good_eqs = [connect(p1, p2)]
 bad_eqs = [connect(p1, p2, op)]
 bad_length_eqs = [connect(op, lp)]
-@test MT.validate(good_eqs)
-@test !MT.validate(bad_eqs)
-@test !MT.validate(bad_length_eqs)
+@test UMT.validate(good_eqs)
+@test !UMT.validate(bad_eqs)
+@test !UMT.validate(bad_length_eqs)
 @named sys = ODESystem(good_eqs, t, [], [])
 @test_throws MT.ValidationError ODESystem(bad_eqs, t, [], []; name = :sys)
 
@@ -136,7 +137,7 @@ noiseeqs = [0.1u"MW" 0.1u"MW"
 # Invalid noise matrix
 noiseeqs = [0.1u"MW" 0.1u"MW"
     0.1u"MW" 0.1u"s"]
-@test !MT.validate(eqs, noiseeqs)
+@test !UMT.validate(eqs, noiseeqs)
 
 # Non-trivial simplifications
 @variables t [unit = u"s"] V(t) [unit = u"m"^3] L(t) [unit = u"m"]