Refactor to make it possible to specialize get_unit for custom types & registered functions.

Author: Lucas Morton

src/systems/validation.jl

@@ -9,76 +9,103 @@ 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."))
-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
+    result
 end
 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(Base.:>), typeof(Base.:<), 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) #why are these not identical?!?
+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 $x, 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 # Not a real function call, just a dependent variable. Unit is on the Sym.
+            return screen_unit(getmetadata(x, VariableUnit, unitless)) 
+        elseif op isa Term && !(operation(op) isa Term) #
+            gp = getmetadata(x,Symbolics.GetindexParent,nothing)
+            return screen_unit(getmetadata(gp, VariableUnit, unitless))
+        elseif op isa Term
+            return screen_unit(getmetadata(x, VariableUnit, unitless))
+        end        
+        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
 
@@ -92,7 +119,7 @@ function safe_get_unit(term, info)
             @warn("$info: $(err.x) and $(err.y) are not dimensionally compatible.")
         elseif err isa ValidationError
             @warn(info*err.message)
-        elseif err isa MethodError
+        elseif err isa MethodError #Warning: Unable to get unit for operation x[1] with arguments SymbolicUtils.Sym{Real, Base.ImmutableDict{DataType, Any}}[t].
             @warn("$info: no method matching $(err.f) for arguments $(typeof.(err.args)).")
         else
             rethrow()

test/units.jl

@@ -52,45 +52,13 @@ eqs = [D(E) ~ P - E/τ
 eqs = [0 ~ σ*(y - x)]
 @test MT.validate(eqs)
 
-#Array variables
-@variables t x[1:3,1:3](t)
+##Array variables
+@variables t [unit = u"s"] x[1:3](t) [unit = u"m"]
+@parameters v[1:3] = [1,2,3] [unit = u"m/s"]
 D = Differential(t)
-eqs = D.(x) .~ x
+eqs = D.(x) .~ v
 ODESystem(eqs,name=:sys)
 
-# Array ops
-using Symbolics: unwrap, wrap
-using LinearAlgebra
-@variables t
-sts = @variables x[1:3](t) y(t)
-ps = @parameters p[1:3] = [1, 2, 3]
-D = Differential(t)
-eqs = [
-       collect(D.(x) ~ x)
-       D(y) ~ norm(x)*y
-      ]
-ODESystem(eqs, t, [sts...;], [ps...;],name=:sys)
-
-#= Not supported yet b/c iterate doesn't work on unitful array
-# Array ops with units
-@variables t [unit =u"s"]
-sts = @variables x[1:3](t) [unit = u"kg"] y(t) [unit = u"kg"]
-ps = @parameters b [unit = u"s"^-1]
-D = Differential(t)
-eqs = [
-       collect(D.(x) ~ b*x)
-       D(y) ~ b*norm(x)
-      ]
-ODESystem(eqs, t, [sts...;], [ps...;])
-
-#Array variables with units
-@variables t [unit = u"s"] x[1:3,1:3](t) [unit = u"kg"] 
-@parameters a [unit = u"s"^-1]
-D = Differential(t)
-eqs = D.(x) .~ a*x
-ODESystem(eqs)
-=#
-
 #Difference equation with units
 @parameters t [unit = u"s"] a [unit = u"s"^-1]
 @variables x(t) [unit = u"kg"]
@@ -99,7 +67,7 @@ D = Difference(t; dt = 0.1u"s")
 eqs = [
     δ(x) ~ a*x 
 ]
-de = ODESystem(eqs, t, [x, y], [a],name=:sys)
+de = ODESystem(eqs, t, [x], [a],name=:sys)
 
 
 @parameters t
@@ -202,20 +170,8 @@ maj2 = MassActionJump(γ, [S => 1], [S => -1])
 @parameters t
 vars = @variables x(t)
 D = Differential(t)
-eqs = 
-[
+eqs = [
     D(x) ~ IfElse.ifelse(t>0.1,2,1)
 ]
 @named sys = ODESystem(eqs, t, vars, [])
 
-#Vectors of symbols
-@parameters t
-@register dummy(vector::Vector{Num}, scalar)
-dummy(vector, scalar) = vector[1] .- scalar
-
-@variables vec[1:2](t)
-vec = collect(vec)
-eqs = [vec .~ dummy(vec, vec[1]);]
-sts = vcat(vec)
-ODESystem(eqs, t, [sts...;], [], name=:sys)
-