Merge pull request #1229 from lamorton/unit_convert

Refactor `get_unit`

Author: ChrisRackauckas

docs/src/basics/Validation.md

@@ -9,18 +9,24 @@ Units may assigned with the following syntax.
 ```julia
 using ModelingToolkit, Unitful
 @variables t [unit = u"s"] x(t) [unit = u"m"] g(t) w(t) [unit = "Hz"]
-#Or,
+
 @variables(t, [unit = u"s"], x(t), [unit = u"m"], g(t), w(t), [unit = "Hz"])
-#Or,
+
 @variables(begin
 t, [unit = u"s"],
 x(t), [unit = u"m"],
 g(t),
 w(t), [unit = "Hz"]
 end)
+
+# Simultaneously set default value (use plain numbers, not quantities)
+@variable x=10 [unit = u"m"]
+
+# Symbolic array: unit applies to all elements
+@variable x[1:3] [unit = u"m"]
 ```
 
-Do not use `quantities` such as `1u"s"` or `1/u"s"` or `u"1/s"` as these will result in errors; instead use `u"s"` or `u"s^1"`. 
+Do not use `quantities` such as  `1u"s"`, `1/u"s"` or `u"1/s"` as these will result in errors; instead use `u"s"`, `u"s^-1"`, or `u"s"^-1`. 
 
 ## Unit Validation & Inspection
 
@@ -62,8 +68,38 @@ eqs = eqs = [D(E) ~ P - E/τ,
                 0 ~ P       ]
 ModelingToolkit.validate(eqs) #Returns false while displaying a warning message
 ```
+## User-Defined Registered Functions and Types
+
+In order to validate user-defined types and `register`ed functions, specialize `get_unit`.  Single-parameter calls to `get_unit`
+expect an object type, while two-parameter calls expect a function type as the first argument, and a vector of arguments as the 
+second argument.
+
+```julia
+using ModelingToolkit
+# Composite type parameter in registered function
+@parameters t
+D = Differential(t)
+struct NewType
+    f
+end
+@register dummycomplex(complex::Num, scalar)
+dummycomplex(complex, scalar) = complex.f - scalar
+
+c = NewType(1)
+MT.get_unit(x::NewType) = MT.get_unit(x.f)
+function MT.get_unit(op::typeof(dummycomplex),args)
+    argunits = MT.get_unit.(args)
+    MT.get_unit(-,args)
+end
+
+sts = @variables a(t)=0 [unit = u"cm"]
+ps = @parameters s=-1 [unit = u"cm"] c=c [unit = u"cm"]
+eqs = [D(a) ~ dummycomplex(c, s);]
+sys = ODESystem(eqs, t, [sts...;], [ps...;], name=:sys)
+sys_simple = structural_simplify(sys)
+```
 
-## `Unitful` Literals & User-Defined Functions
+## `Unitful` Literals
 
 In order for a function to work correctly during both validation & execution, the function must be unit-agnostic. That is, no unitful literals may be used. Any unitful quantity must either be a `parameter` or `variable`. For example, these equations will not validate successfully. 
 

src/systems/validation.jl

@@ -5,80 +5,119 @@ struct ValidationError <: Exception
     message::String
 end
 
+"Throw exception on invalid unit types, otherwise return argument."
 function screen_unit(result)
     result isa Unitful.Unitlike || throw(ValidationError("Unit must be a subtype of Unitful.Unitlike, not $(typeof(result))."))
     result isa Unitful.ScalarUnits || throw(ValidationError("Non-scalar units such as $result are not supported. Use a scalar unit instead."))
     result == u"°" && throw(ValidationError("Degrees are not supported. Use radians instead."))
+    result
 end
-"Find the unit of a symbolic item."
-get_unit(x::Real) = unitless
-function get_unit(x::Unitful.Quantity)
-    result = Unitful.unit(x)
-    screen_unit(result)
-    return result
-end
+
+"""Test unit equivalence.
+
+Example of implemented behavior:
+```julia
+using ModelingToolkit, Unitful
+MT = ModelingToolkit
+@parameters γ P [unit = u"MW"] E [unit = u"kJ"] τ [unit = u"ms"]
+@test MT.equivalent(u"MW" ,u"kJ/ms") # Understands prefixes
+@test !MT.equivalent(u"m", u"cm") # Units must be same magnitude
+@test MT.equivalent(MT.get_unit(P^γ), MT.get_unit((E/τ)^γ)) # Handles symbolic exponents
+```
+"""
 equivalent(x,y) = isequal(1*x,1*y)
 unitless = Unitful.unit(1)
 
+#For dispatching get_unit
+Literal = Union{Sym,Symbolics.ArrayOp,Symbolics.Arr,Symbolics.CallWithMetadata}
+Conditional = Union{typeof(ifelse),typeof(IfElse.ifelse)}
+Comparison = Union{typeof.([==, !=, ≠, <, <=, ≤, >, >=, ≥])...}
+
+"Find the unit of a symbolic item."
+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))
-function get_unit(x::Symbolic)
-    if x isa Sym || operation(x) isa Sym || (operation(x) isa Term && operation(x).f == getindex) || x isa Symbolics.ArrayOp
-        if x.metadata !== nothing
-            symunits = get(x.metadata, VariableUnit, unitless)
-            screen_unit(symunits)
-        else
-            symunits = unitless
-        end
-        return symunits
-    elseif operation(x) isa Differential
-        return get_unit(arguments(x)[1]) / get_unit(operation(x).x)
-    elseif operation(x) isa Integral
-        unit = 1
-        if operation(x).x isa Vector
-            for u in operation(x).x
-                unit *= get_unit(u)
-            end
-        else
-            unit *= get_unit(operation(x).x)
-        end
-        return get_unit(arguments(x)[1]) * unit
-    elseif  operation(x) isa Difference
-        return get_unit(arguments(x)[1]) / get_unit(operation(x).t) #TODO: make this same as Differential
-    elseif x isa Pow
-        pargs = arguments(x)
-        base,expon = get_unit.(pargs)
-        @assert expon isa Unitful.DimensionlessUnits
-        if base == unitless
-            unitless
-        else
-            pargs[2] isa Number ? operation(x)(base, pargs[2]) : operation(x)(1*base, pargs[2])
-        end
-    elseif x isa Add # Cannot simply add the units b/c they may differ in magnitude (eg, kg vs g)
-        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 operation(x) in ( Base.:> ,  Base.:< , == )
-        terms = get_unit.(arguments(x))
-        equivalent(terms[1],terms[2]) || throw(ValidationError(", in comparison $x, units [$(terms[1])] and [$(terms[2])] do not match."))
-        return unitless
-    elseif operation(x) == ifelse || operation(x) == IfElse.ifelse
-         terms = get_unit.(arguments(x))
-        terms[1] == unitless || throw(ValidationError(", in $x, [$(terms[1])] is not dimensionless."))
-        equivalent(terms[2],terms[3]) || throw(ValidationError(", in $x, units [$(terms[2])] and [$(terms[3])] do not match."))
-        return terms[2]
-    elseif operation(x) == Symbolics._mapreduce
-        if x.arguments[2] == +
-            get_unit(x.arguments[3])
-        else
-            throw(ValidationError("Unsupported array operation $x"))
+get_unit(x::Literal) = screen_unit(getmetadata(x,VariableUnit, unitless))
+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]) 
+function get_unit(op,args) # Fallback
+    result = op(1 .* get_unit.(args)...)
+    try 
+        unit(result)
+    catch 
+        throw(ValidationError("Unable to get unit for operation $op with arguments $args."))
+    end
+end
+
+function get_unit(op::Integral,args)
+    unit = 1
+    if op.x isa Vector
+        for u in op.x
+            unit *= get_unit(u)
         end
     else
-        return get_unit(operation(x)(1 .* get_unit.(arguments(x))...))
+        unit *= get_unit(op.x)
+    end
+    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 || throw(ValidationError(", in $x, [$(terms[1])] is not dimensionless."))
+    equivalent(terms[2], terms[3]) || throw(ValidationError(", in $x, units [$(terms[2])] and [$(terms[3])] do not match."))
+    return terms[2]
+end
+
+function get_unit(op::typeof(Symbolics._mapreduce),args)
+    if args[2] == +
+        get_unit(args[3])
+    else
+        throw(ValidationError("Unsupported array operation $op"))
+    end
+end
+
+function get_unit(op::Comparison, args)
+    terms = get_unit.(args)
+    equivalent(terms[1], terms[2]) || throw(ValidationError(", in comparison $op, units [$(terms[1])] and [$(terms[2])] do not match."))
+    return unitless
+end
+
+function get_unit(x::Symbolic) 
+    if SymbolicUtils.istree(x)
+        op = operation(x)
+        if op isa Sym || (op isa Term && operation(op) isa Term) # 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) 
+            gp = getmetadata(x,Symbolics.GetindexParent,nothing) # Like x[1](t)
+            return screen_unit(getmetadata(gp, VariableUnit, unitless))
+        end  # Actual function calls:
+        args = arguments(x)
+        return get_unit(op, args)
+    else # This function should only be reached by Terms, for which `istree` is true
+        throw(ArgumentError("Unsupported value $x."))
     end
 end