feat: add support for bounds of array variables

Author: AayushSabharwal

src/variables.jl

@@ -254,7 +254,6 @@ end
 ## Bounds ======================================================================
 struct VariableBounds end
 Symbolics.option_to_metadata_type(::Val{:bounds}) = VariableBounds
-getbounds(x::Num) = getbounds(Symbolics.unwrap(x))
 
 """
     getbounds(x)
@@ -266,10 +265,35 @@ Create parameters with bounds like this
 @parameters p [bounds=(-1, 1)]
 ```
 """
-function getbounds(x)
+function getbounds(x::Union{Num, Symbolics.Arr, SymbolicUtils.Symbolic})
+    x = unwrap(x)
     p = Symbolics.getparent(x, nothing)
-    p === nothing || (x = p)
-    Symbolics.getmetadata(x, VariableBounds, (-Inf, Inf))
+    if p === nothing
+        bounds = Symbolics.getmetadata(x, VariableBounds, (-Inf, Inf))
+        if symbolic_type(x) == ArraySymbolic() && Symbolics.shape(x) != Symbolics.Unknown()
+            bounds = map(bounds) do b
+                b isa AbstractArray && return b
+                return fill(b, size(x))
+            end
+        end
+    else
+        # if we reached here, `x` is the result of calling `getindex`
+        bounds = @something Symbolics.getmetadata(x, VariableBounds, nothing) getbounds(p)
+        idxs = arguments(x)[2:end]
+        bounds = map(bounds) do b
+            if b isa AbstractArray
+                if Symbolics.shape(p) != Symbolics.Unknown() && size(p) != size(b)
+                    throw(DimensionMismatch("Expected array variable $p with shape $(size(p)) to have bounds of identical size. Found $bounds of size $(size(bounds))."))
+                end
+                return b[idxs...]
+            elseif symbolic_type(x) == ArraySymbolic()
+                return fill(b, size(x))
+            else
+                return b
+            end
+        end
+    end
+    return bounds
 end
 
 """
@@ -280,7 +304,7 @@ See also [`getbounds`](@ref).
 """
 function hasbounds(x)
     b = getbounds(x)
-    isfinite(b[1]) || isfinite(b[2])
+    any(isfinite.(b[1]) .|| isfinite.(b[2]))
 end
 
 ## Disturbance =================================================================

test/test_variable_metadata.jl

@@ -10,6 +10,47 @@ using ModelingToolkit
 @test !hasbounds(y)
 @test !haskey(ModelingToolkit.dump_variable_metadata(y), :bounds)
 
+@variables y[1:3]
+@test !hasbounds(y)
+@test getbounds(y)[1] == [-Inf, -Inf, -Inf]
+@test getbounds(y)[2] == [Inf, Inf, Inf]
+for i in eachindex(y)
+    @test !hasbounds(y[i])
+    b = getbounds(y[i])
+    @test b[1] == -Inf && b[2] == Inf
+end
+
+@variables y[1:3] [bounds = (-1, 1)]
+@test hasbounds(y)
+@test getbounds(y)[1] == -ones(3)
+@test getbounds(y)[2] == ones(3)
+for i in eachindex(y)
+    @test hasbounds(y[i])
+    b = getbounds(y[i])
+    @test b[1] == -1.0 && b[2] == 1.0
+end
+@test getbounds(y[1:2])[1] == -ones(2)
+@test getbounds(y[1:2])[2] == ones(2)
+
+@variables y[1:2, 1:2] [bounds = (-1, [1.0 Inf; 2.0 3.0])]
+@test hasbounds(y)
+@test getbounds(y)[1] == [-1 -1; -1 -1]
+@test getbounds(y)[2] == [1.0 Inf; 2.0 3.0]
+for i in eachindex(y)
+    @test hasbounds(y[i])
+    b = getbounds(y[i])
+    @test b[1] == -1 && b[2] == [1.0 Inf; 2.0 3.0][i]
+end
+
+@variables y[1:2] [bounds = (-Inf, [1.0, Inf])]
+@test hasbounds(y)
+@test getbounds(y)[1] == [-Inf, -Inf]
+@test getbounds(y)[2] == [1.0, Inf]
+@test hasbounds(y[1])
+@test getbounds(y[1]) == (-Inf, 1.0)
+@test !hasbounds(y[2])
+@test getbounds(y[2]) == (-Inf, Inf)
+
 # Guess
 @variables y [guess = 0]
 @test getguess(y) === 0