add dispatch to symbolic array for units

SciML · Sep 4, 2024 · 8f2d885 · 8f2d885
1 parent f1b8e72
commit 8f2d885
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/src/systems/unit_check.jl b/src/systems/unit_check.jl
@@ -69,6 +69,7 @@ get_unit(x::Real) = unitless
 get_unit(x::DQ.AbstractQuantity) = screen_unit(x)
 get_unit(x::AbstractArray) = map(get_unit, x)
 get_unit(x::Num) = get_unit(unwrap(x))
+get_unit(x::Symbolics.Arr) = get_unit(unwrap(x))
 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])

diff --git a/test/dq_units.jl b/test/dq_units.jl
@@ -239,3 +239,10 @@ let
     @test mm2units == MT.oneunit(mm2units)
     @test mmunits == mm2units
 end
+
+# test for array variable units https://github.com/SciML/ModelingToolkit.jl/issues/3009
+let
+    @variables x_vec(t)[1:3] [unit = u"1"] x_mat(t)[1:3, 1:3] [unit = u"1"]
+    @test MT.get_unit(x_vec) == u"1"
+    @test MT.get_unit(x_mat) == u"1"
+end