From 56f4e35e0033646f9132a1560bee12b8216e871e Mon Sep 17 00:00:00 2001
From: Venkateshprasad <32921645+ven-k@users.noreply.github.com>
Date: Thu, 28 Mar 2024 13:39:19 +0530
Subject: [PATCH] fix: pass `VariableUnit` of diff var while converting it to a
 `Term`

Units are defined only in MTK, so pass this piece of info to `Symbolics.diff2term` to add this to the returning term-var.
---
 src/variables.jl       | 2 +-
 test/variable_utils.jl | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/variables.jl b/src/variables.jl
index ed9edcffb4..4108a865d2 100644
--- a/src/variables.jl
+++ b/src/variables.jl
@@ -109,7 +109,7 @@ function default_toterm(x)
             end
             x = normalize_to_differential(op)(arguments(x)...)
         end
-        Symbolics.diff2term(x)
+        Symbolics.diff2term(x, Dict(VariableUnit => get_unit(x)))
     else
         x
     end
diff --git a/test/variable_utils.jl b/test/variable_utils.jl
index 551614038b..a49179754d 100644
--- a/test/variable_utils.jl
+++ b/test/variable_utils.jl
@@ -19,7 +19,7 @@ new = (((1 / β - 1) + δ) / γ)^(1 / (γ - 1))
 using ModelingToolkit: isdifferential, vars, collect_differential_variables,
                        collect_ivs
 @variables t u(t) y(t)
-D = Differential(t)
+using ModelingToolkit: D
 eq = D(y) ~ u
 v = vars(eq)
 @test v == Set([D(y), u])
@@ -32,3 +32,9 @@ aov = ModelingToolkit.collect_applied_operators(eq, Differential)
 
 ts = collect_ivs([eq])
 @test ts == Set([t])
+
+# Test units of diff vars of `Term` type.
+using ModelingToolkit: t, get_unit, default_toterm
+@variables k(t)
+k2 = D(D(k))
+get_unit(default_toterm(k2)) == get_unit(k2)