Merge pull request #3209 from AayushSabharwal/as/array-bounds

feat: add support for bounds of array variables

Author: ChrisRackauckas

docs/src/basics/Variable_metadata.md

@@ -86,6 +86,44 @@ hasbounds(u)
 getbounds(u)
 ```
 
+Bounds can also be specified for array variables. A scalar array bound is applied to each
+element of the array. A bound may also be specified as an array, in which case the size of
+the array must match the size of the symbolic variable.
+
+```@example metadata
+@variables x[1:2, 1:2] [bounds = (-1, 1)]
+hasbounds(x)
+```
+
+```@example metadata
+getbounds(x)
+```
+
+```@example metadata
+getbounds(x[1, 1])
+```
+
+```@example metadata
+getbounds(x[1:2, 1])
+```
+
+```@example metadata
+@variables x[1:2] [bounds = (-Inf, [1.0, Inf])]
+hasbounds(x)
+```
+
+```@example metadata
+getbounds(x)
+```
+
+```@example metadata
+getbounds(x[2])
+```
+
+```@example metadata
+hasbounds(x[2])
+```
+
 ## Guess
 
 Specify an initial guess for custom initial conditions of an `ODESystem`.

src/systems/optimization/optimizationsystem.jl

@@ -101,8 +101,8 @@ function OptimizationSystem(op, unknowns, ps;
         gui_metadata = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
-    constraints = value.(scalarize(constraints))
-    unknowns′ = value.(scalarize(unknowns))
+    constraints = value.(reduce(vcat, scalarize(constraints); init = []))
+    unknowns′ = value.(reduce(vcat, scalarize(unknowns); init = []))
     ps′ = value.(ps)
     op′ = value(scalarize(op))
 

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/optimizationsystem.jl

@@ -358,4 +358,13 @@ end
     @mtkbuild sys = OptimizationSystem(obj, [x, y, z], []; constraints = cons)
     @test is_variable(sys, z)
     @test !is_variable(sys, y)
+
+    @variables x[1:3] [bounds = ([-Inf, -1.0, -2.0], [Inf, 1.0, 2.0])]
+    obj = x[1]^2 + x[2]^2 + x[3]^2
+    cons = [x[2] ~ 2x[1] + 3, x[3] ~ x[1] + x[2]]
+    @mtkbuild sys = OptimizationSystem(obj, [x], []; constraints = cons)
+    @test length(unknowns(sys)) == 2
+    @test !is_variable(sys, x[1])
+    @test is_variable(sys, x[2])
+    @test is_variable(sys, x[3])
 end

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