Merge pull request #3034 from SciML/myb/unit_fix

Add unit aware diff2term

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -11,7 +11,7 @@ using SymbolicUtils: maketerm, iscall
 
 using ModelingToolkit
 using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Differential,
-                       unknowns, equations, vars, Symbolic, diff2term, value,
+                       unknowns, equations, vars, Symbolic, diff2term_with_unit, value,
                        operation, arguments, Sym, Term, simplify, symbolic_linear_solve,
                        isdiffeq, isdifferential, isirreducible,
                        empty_substitutions, get_substitutions,
@@ -22,7 +22,7 @@ using ModelingToolkit: ODESystem, AbstractSystem, var_from_nested_derivative, Di
                        get_postprocess_fbody, vars!,
                        IncrementalCycleTracker, add_edge_checked!, topological_sort,
                        invalidate_cache!, Substitutions, get_or_construct_tearing_state,
-                       filter_kwargs, lower_varname, setio, SparseMatrixCLIL,
+                       filter_kwargs, lower_varname_with_unit, setio, SparseMatrixCLIL,
                        get_fullvars, has_equations, observed,
                        Schedule
 

src/structural_transformation/pantelides.jl

@@ -16,7 +16,8 @@ function pantelides_reassemble(state::TearingState, var_eq_matching)
     fill!(out_vars, nothing)
     out_vars[1:length(fullvars)] .= fullvars
 
-    D = Differential(get_iv(sys))
+    iv = get_iv(sys)
+    D = Differential(iv)
 
     for (varidx, diff) in edges(var_to_diff)
         # fullvars[diff] = D(fullvars[var])
@@ -25,7 +26,7 @@ function pantelides_reassemble(state::TearingState, var_eq_matching)
         # `fullvars[i]` needs to be not a `D(...)`, because we want the DAE to be
         # first-order.
         if isdifferential(vi)
-            vi = out_vars[varidx] = diff2term(vi)
+            vi = out_vars[varidx] = diff2term_with_unit(vi, iv)
         end
         out_vars[diff] = D(vi)
     end

src/structural_transformation/symbolics_tearing.jl

@@ -274,7 +274,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
         dv === nothing && continue
         if var_eq_matching[var] !== SelectedState()
             dd = fullvars[dv]
-            v_t = setio(diff2term(unwrap(dd)), false, false)
+            v_t = setio(diff2term_with_unit(unwrap(dd), unwrap(iv)), false, false)
             for eq in 𝑑neighbors(graph, dv)
                 dummy_sub[dd] = v_t
                 neweqs[eq] = fast_substitute(neweqs[eq], dd => v_t)

src/structural_transformation/utils.jl

@@ -426,7 +426,7 @@ function lower_varname_withshift(var, iv, order)
         op = operation(var)
         return Shift(op.t, order)(var)
     end
-    return lower_varname(var, iv, order)
+    return lower_varname_with_unit(var, iv, order)
 end
 
 function isdoubleshift(var)

src/systems/diffeqs/abstractodesystem.jl

@@ -456,7 +456,7 @@ end
 
 function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys),
         ps = parameters(sys), u0 = nothing;
-        ddvs = map(diff2term ∘ Differential(get_iv(sys)), dvs),
+        ddvs = map(Base.Fix2(diff2term, get_iv(sys)) ∘ Differential(get_iv(sys)), dvs),
         version = nothing, p = nothing,
         jac = false,
         eval_expression = false,

src/systems/diffeqs/first_order_transform.jl

@@ -34,7 +34,7 @@ function ode_order_lowering(eqs, iv, unknown_vars)
             var, maxorder = var_from_nested_derivative(eq.lhs)
             maxorder > get(var_order, var, 1) && (var_order[var] = maxorder)
             var′ = lower_varname(var, iv, maxorder - 1)
-            rhs′ = diff2term(eq.rhs)
+            rhs′ = diff2term_with_unit(eq.rhs, iv)
             push!(diff_vars, var′)
             push!(diff_eqs, D(var′) ~ rhs′)
         end

src/utils.jl

@@ -859,3 +859,16 @@ function eval_or_rgf(expr::Expr; eval_expression = false, eval_module = @__MODUL
         return drop_expr(RuntimeGeneratedFunction(eval_module, eval_module, expr))
     end
 end
+
+function _with_unit(f, x, t, args...)
+    x = f(x, args...)
+    if hasmetadata(x, VariableUnit) && (t isa Symbolic && hasmetadata(t, VariableUnit))
+        xu = getmetadata(x, VariableUnit)
+        tu = getmetadata(t, VariableUnit)
+        x = setmetadata(x, VariableUnit, xu / tu)
+    end
+    return x
+end
+
+diff2term_with_unit(x, t) = _with_unit(diff2term, x, t)
+lower_varname_with_unit(var, iv, order) = _with_unit(lower_varname, var, iv, iv, order)

test/dq_units.jl

@@ -246,3 +246,31 @@ let
     @test MT.get_unit(x_vec) == u"1"
     @test MT.get_unit(x_mat) == u"1"
 end
+
+module UnitTD
+using Test
+using ModelingToolkit
+using ModelingToolkit: t, D
+using DynamicQuantities
+
+@mtkmodel UnitsExample begin
+    @parameters begin
+        g, [unit = u"m/s^2"]
+        L = 1.0, [unit = u"m"]
+    end
+    @variables begin
+        x(t), [unit = u"m"]
+        y(t), [state_priority = 10, unit = u"m"]
+        λ(t), [unit = u"s^-2"]
+    end
+    @equations begin
+        D(D(x)) ~ λ * x
+        D(D(y)) ~ λ * y - g
+        x^2 + y^2 ~ L^2
+    end
+end
+
+@mtkbuild pend = UnitsExample()
+@test ModelingToolkit.get_unit.(filter(x -> occursin("ˍt", string(x)), unknowns(pend))) ==
+      [u"m/s", u"m/s"]
+end