Update lower_varname as well

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -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/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

@@ -860,12 +860,15 @@ function eval_or_rgf(expr::Expr; eval_expression = false, eval_module = @__MODUL
     end
 end
 
-function diff2term_with_unit(x, t)
-    x = diff2term(x)
+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,29 @@ let
     @test MT.get_unit(x_vec) == u"1"
     @test MT.get_unit(x_mat) == u"1"
 end
+
+module UnitTD
+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