diff --git a/examples/quadrotor.jl b/examples/quadrotor.jl deleted file mode 100644 index 5918f73..0000000 --- a/examples/quadrotor.jl +++ /dev/null @@ -1,127 +0,0 @@ -# Quadrotor tutorial of ExaModels translated to GenOpt -# See https://exanauts.github.io/ExaModels.jl/stable/quad/ -# This example was used for the JuMP-dev 2025 presentation -# https://jump.dev/meetings/jumpdev2025/ - -N = 3 - -n = 9 -p = 4 -function d(i, j, N) - return (j == 1 ? 1 * sin(2 * pi / N * i) : 0.0) + - (j == 3 ? 2 * sin(4 * pi / N * i) : 0.0) + - (j == 5 ? 2 * i / N : 0.0) -end -dt = 1/N -R = fill(1 / 10, 4) -Q = [1, 0, 1, 0, 1, 0, 1, 1, 1] -Qf = [1, 0, 1, 0, 1, 0, 1, 1, 1] / dt - -x0 = zeros(n) - -using JuMP -model = Model() - -@variable(model, x[1:(N+1), 1:n]) -@variable(model, u[1:N, 1:p]) - -using GenOpt -container = ParametrizedArray - -@constraint(model, [i in 1:n], x[1, i] == x0[i], container = container) -@constraint( - model, - [i in 1:N], - x[i+1, 1] == x[i, 1] + (x[i, 2]) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 2] == - x[i, 2] + - ( - u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * cos(x[i, 9]) + - u[i, 1] * sin(x[i, 7]) * sin(x[i, 9]) - ) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 3] == x[i, 3] + (x[i, 4]) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 4] == - x[i, 4] + - ( - u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * sin(x[i, 9]) - - u[i, 1] * sin(x[i, 7]) * cos(x[i, 9]) - ) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 5] == x[i, 5] + (x[i, 6]) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 6] == x[i, 6] + (u[i, 1] * cos(x[i, 7]) * cos(x[i, 8]) - 9.8) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 7] == - x[i, 7] + - ( - u[i, 2] * cos(x[i, 7]) / cos(x[i, 8]) + - u[i, 3] * sin(x[i, 7]) / cos(x[i, 8]) - ) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 8] == - x[i, 8] + (-u[i, 2] * sin(x[i, 7]) + u[i, 3] * cos(x[i, 7])) * dt, - container = container, -) -@constraint( - model, - [i in 1:N], - x[i+1, 9] == - x[i, 9] + - ( - u[i, 2] * cos(x[i, 7]) * tan(x[i, 8]) + - u[i, 3] * sin(x[i, 7]) * tan(x[i, 8]) + - u[i, 4] - ) * dt, - container = container, -) - -ss = lazy_sum(R[j] * u[1, j] for j in 1:p) -ss = lazy_sum(R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) -itr1 = [(i, j, d(i, j, N)) for i in 1:N, j in 1:n] -itr2 = [(j, d(N + 1, j, N)) for j in 1:n] -@objective( - model, - Min, - lazy_sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) + - lazy_sum(0.5 * Q[it[2]] * (x[it[1], it[2]] - it[3])^2 for it in itr1) + - lazy_sum(0.5 * Qf[it[1]] * (x[N+1, it[1]] - it[2])^2 for it in itr2), -) - -import MadNLP -import ExaModels -# Needs https://github.com/exanauts/ExaModels.jl/pull/237 -set_optimizer(model, () -> ExaModels.Optimizer(MadNLP.madnlp)) -optimize!(model) -value.(x) -value.(u) diff --git a/examples/quadrotor/main.jl b/examples/quadrotor/main.jl new file mode 100644 index 0000000..f9e55df --- /dev/null +++ b/examples/quadrotor/main.jl @@ -0,0 +1,11 @@ +import MadNLP +import ExaModels + +include(joinpath(@__DIR__, "model.jl")) +model = build_model() + +# Needs https://github.com/exanauts/ExaModels.jl/pull/237 +set_optimizer(model, () -> ExaModels.Optimizer(MadNLP.madnlp)) +optimize!(model) +value.(x) +value.(u) diff --git a/examples/quadrotor/model.jl b/examples/quadrotor/model.jl new file mode 100644 index 0000000..fbdbac6 --- /dev/null +++ b/examples/quadrotor/model.jl @@ -0,0 +1,121 @@ +# Quadrotor tutorial of ExaModels translated to GenOpt +# See https://exanauts.github.io/ExaModels.jl/stable/quad/ +# This example was used for the JuMP-dev 2025 presentation +# https://jump.dev/meetings/jumpdev2025/ + +using JuMP, GenOpt + +function d(i, j, N) + return (j == 1 ? 1 * sin(2 * pi / N * i) : 0.0) + + (j == 3 ? 2 * sin(4 * pi / N * i) : 0.0) + + (j == 5 ? 2 * i / N : 0.0) +end + +function build_model(; N = 3, n = 9, p = 4) + dt = 1/N + R = fill(1 / 10, 4) + Q = [1, 0, 1, 0, 1, 0, 1, 1, 1] + Qf = [1, 0, 1, 0, 1, 0, 1, 1, 1] / dt + + x0 = zeros(n) + + model = Model() + + @variable(model, x[1:(N+1), 1:n]) + @variable(model, u[1:N, 1:p]) + + container = ParametrizedArray + + @constraint(model, [i in 1:n], x[1, i] == x0[i], container = container) + @constraint( + model, + [i in 1:N], + x[i+1, 1] == x[i, 1] + (x[i, 2]) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 2] == + x[i, 2] + + ( + u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * cos(x[i, 9]) + + u[i, 1] * sin(x[i, 7]) * sin(x[i, 9]) + ) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 3] == x[i, 3] + (x[i, 4]) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 4] == + x[i, 4] + + ( + u[i, 1] * cos(x[i, 7]) * sin(x[i, 8]) * sin(x[i, 9]) - + u[i, 1] * sin(x[i, 7]) * cos(x[i, 9]) + ) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 5] == x[i, 5] + (x[i, 6]) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 6] == + x[i, 6] + (u[i, 1] * cos(x[i, 7]) * cos(x[i, 8]) - 9.8) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 7] == + x[i, 7] + + ( + u[i, 2] * cos(x[i, 7]) / cos(x[i, 8]) + + u[i, 3] * sin(x[i, 7]) / cos(x[i, 8]) + ) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 8] == + x[i, 8] + (-u[i, 2] * sin(x[i, 7]) + u[i, 3] * cos(x[i, 7])) * dt, + container = container, + ) + @constraint( + model, + [i in 1:N], + x[i+1, 9] == + x[i, 9] + + ( + u[i, 2] * cos(x[i, 7]) * tan(x[i, 8]) + + u[i, 3] * sin(x[i, 7]) * tan(x[i, 8]) + + u[i, 4] + ) * dt, + container = container, + ) + + ss = lazy_sum(R[j] * u[1, j] for j in 1:p) + ss = lazy_sum(R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) + itr1 = [(i, j, d(i, j, N)) for i in 1:N, j in 1:n] + itr2 = [(j, d(N + 1, j, N)) for j in 1:n] + @objective( + model, + Min, + lazy_sum(0.5 * R[j] * (u[i, j]^2) for i in 1:N, j in 1:p) + + lazy_sum(0.5 * Q[it[2]] * (x[it[1], it[2]] - it[3])^2 for it in itr1) + + lazy_sum(0.5 * Qf[it[1]] * (x[N+1, it[1]] - it[2])^2 for it in itr2), + ) + + return model +end diff --git a/examples/quadrotor/runtests.jl b/examples/quadrotor/runtests.jl new file mode 100644 index 0000000..c449daf --- /dev/null +++ b/examples/quadrotor/runtests.jl @@ -0,0 +1,86 @@ +using Test +import MathOptInterface as MOI + +include(joinpath(@__DIR__, "model.jl")) + +# The model is built with `N = 3` and `n = 9` so the 10 `@constraint` calls +# using `container = ParametrizedArray` encode `9 + 9 * 3 == 36` scalar +# equalities. The point of GenOpt is that they stay grouped as 10 generators +# instead of being scalarized, so that is what we check here. +const N, n = 3, 9 + +# `x[i+1, j] == x[i, j] + x[i, j+1] * dt` for `j in (1, 3, 5)` is affine, as is +# the initial condition `x[1, i] == x0[i]`. The 6 remaining dynamics involve +# `cos`, `sin` or `tan` and are nonlinear. +const AFFINE = ExprGenerator{JuMP.AffExpr,JuMP.VariableRef} +const NONLINEAR = ExprGenerator{JuMP.NonlinearExpr,JuMP.VariableRef} + +model = build_model(; N = N, n = n) +@test model isa JuMP.Model + +@testset "list_of_constraint_types" begin + types = JuMP.list_of_constraint_types(model) + # Both generator types are present and, importantly, the affine block was + # not widened into `NonlinearExpr` by the nonlinear dynamics. + @test Set(types) == Set([(AFFINE, MOI.Zeros), (NONLINEAR, MOI.Zeros)]) + # Nothing was scalarized: a scalarized model would report the standard + # `(AffExpr, MOI.EqualTo{Float64})` pair instead. + @test !any(S <: MOI.AbstractScalarSet for (_, S) in types) + @test all(F <: ExprGenerator for (F, _) in types) +end + +@testset "generators are not scalarized" begin + b = JuMP.backend(model) + moi_types = MOI.get(b, MOI.ListOfConstraintTypesPresent()) + # `list_of_constraint_types` is the JuMP image of the MOI list. + @test JuMP.list_of_constraint_types(model) == + [(JuMP.jump_function_type(model, F), S) for (F, S) in moi_types] + + dims = Dict{Type,Vector{Int}}() + for (F, S) in moi_types + cis = MOI.get(b, MOI.ListOfConstraintIndices{F,S}()) + dims[JuMP.jump_function_type(model, F)] = sort!([ + MOI.output_dimension(MOI.get(b, MOI.ConstraintFunction(), ci)) + for ci in cis + ]) + end + # 1 initial condition over `n` and 3 affine dynamics over `N`. + @test dims[AFFINE] == [N, N, N, n] + # The 6 remaining dynamics, each over `N`. + @test dims[NONLINEAR] == fill(N, 6) + # 10 generators encoding 36 scalar equalities. + @test sum(length, values(dims)) == 10 + @test sum(sum, values(dims)) == n + 9 * N +end + +@testset "constraint_object" begin + b = JuMP.backend(model) + for (F, S) in MOI.get(b, MOI.ListOfConstraintTypesPresent()) + E = JuMP.jump_function_type(model, F) + for ci in MOI.get(b, MOI.ListOfConstraintIndices{F,S}()) + ref = JuMP.ConstraintRef(model, ci, JuMP.VectorShape()) + con = JuMP.constraint_object(ref) + @test con isa IteratedConstraint + @test con.func isa E + @test con.set isa MOI.Zeros + # The generator expands lazily to one scalar expression per index. + func = MOI.get(b, MOI.ConstraintFunction(), ci) + @test length(con.func) == MOI.output_dimension(func) + @test MOI.dimension(con.set) == length(con.func) + @test all(e -> e isa JuMP.AbstractJuMPScalar, con.func) + end + end +end + +@testset "objective keeps its generators" begin + b = JuMP.backend(model) + F = MOI.get(b, MOI.ObjectiveFunctionType()) + @test F == MOI.ScalarNonlinearFunction + obj = MOI.get(b, MOI.ObjectiveFunction{F}()) + # The three `lazy_sum` terms are summed without being expanded. + @test obj.head == :+ + @test all(a -> a isa SumGenerator, obj.args) + terms = [prod(it -> length(it.values), a.iterators) for a in obj.args] + # `N * p` control terms, `N * n` stage terms and `n` terminal terms. + @test terms == [N * 4, N * n, n] +end diff --git a/examples/runtests.jl b/examples/runtests.jl new file mode 100644 index 0000000..f497b58 --- /dev/null +++ b/examples/runtests.jl @@ -0,0 +1,9 @@ +using Test + +for dir in readdir(@__DIR__) + if isdir(joinpath(@__DIR__, dir)) + @testset "$dir" begin + include(joinpath(@__DIR__, dir, "runtests.jl")) + end + end +end diff --git a/src/JuMP_wrapper.jl b/src/JuMP_wrapper.jl index 331b860..d0bb941 100644 --- a/src/JuMP_wrapper.jl +++ b/src/JuMP_wrapper.jl @@ -89,7 +89,35 @@ function JuMP.jump_function(model, f::FunctionGenerator{F}) where {F} ) end -_size(expr::ExprGenerator) = length.(getfield.(expr.expr.iterators, :values)) +function JuMP.jump_function_type( + model::JuMP.GenericModel{T}, + ::Type{FunctionGenerator{F}}, +) where {T,F} + return ExprGenerator{ + JuMP.jump_function_type(model, F), + JuMP.GenericVariableRef{T}, + } +end + +_size(expr::ExprGenerator) = length.(expr.expr.iterators) + +""" + _ind2sub(size::Vector{Int}, i::Integer) + +Subscripts of the `i`th entry of an array of size `size`, in column-major +order like `CartesianIndices`. `Base` only provides this for a *tuple* of +dimensions; `size` is a vector here since the number of iterators is not +known at compile time. +""" +function _ind2sub(size::Vector{Int}, i::Integer) + sub = similar(size) + rest = i - 1 + for k in eachindex(size) + rest, j = divrem(rest, size[k]) + sub[k] = j + 1 + end + return sub +end index_iterators(func, _) = func @@ -102,7 +130,10 @@ function index_iterators(func::JuMP.GenericNonlinearExpr, values) if any(JuMP._has_variable_ref_type, args) return JuMP.GenericNonlinearExpr(func.head, args) elseif func.head == :getindex - return getindex(args...) + # `JuMP.jump_function` converts every number to `Float64` when it + # converts a `MOI` function back, so the indices need to be converted + # back to `Int` here, indexing with a `Real` being deprecated. + return getindex(args[1], Int[i for i in args[2:end]]...) else registry = MOI.Nonlinear.OperatorRegistry() if length(func.args) == 1 @@ -114,13 +145,14 @@ function index_iterators(func::JuMP.GenericNonlinearExpr, values) end function Base.getindex(expr::ExprGenerator, i::Integer) - idx = CartesianIndices(Base.OneTo.(_size(expr)))[i] - values = - [expr.iterators[i].values[idx[i]] for i in eachindex(expr.iterators)] + @boundscheck checkbounds(expr, i) + its = expr.expr.iterators + sub = _ind2sub(_size(expr), i) + values = [its[k].values[sub[k]] for k in eachindex(its)] return index_iterators(expr.expr.expr, values) end -Base.length(expr::ExprGenerator) = prod(_size(expr)) +Base.size(expr::ExprGenerator) = (prod(_size(expr)),) struct ParametrizedArray constraint::Any diff --git a/test/JuMP_wrapper.jl b/test/JuMP_wrapper.jl index f512d2e..2e3f955 100644 --- a/test/JuMP_wrapper.jl +++ b/test/JuMP_wrapper.jl @@ -57,6 +57,38 @@ function test_container() @test_broken JuMP.isequal_canonical(con_expr, expr) end +function test_ind2sub() + # `bridge.jl` expands a generator with `CartesianIndices` so `_ind2sub` + # must agree with it, otherwise `getindex` and the bridge would disagree + # on which entry is the `i`th one. + for size in ([4], [2, 3], [3, 1, 2]) + indices = CartesianIndices(Tuple(size)) + for i in eachindex(IndexLinear(), indices) + @test GenOpt._ind2sub(size, i) == collect(Tuple(indices[i])) + end + end + return +end + +function test_generator_getindex() + model = Model() + @variable(model, x) + con_ref = @constraint( + model, + [i in 1:2, j in 1:3], + x >= 10 * i + j, + container = ParametrizedArray, + ) + gen = constraint_object(con_ref.constraint).func + @test gen isa ExprGenerator + @test size(gen) == (6,) + @test length(gen) == 6 + # The first iterator varies fastest, like `CartesianIndices`. + @test [sprint(show, e) for e in gen] == ["(x - $c) - 0" for c in [11, 21, 12, 22, 13, 23]] + @test_throws BoundsError gen[7] + return +end + end # module TestJuMP.runtests() diff --git a/test/runtests.jl b/test/runtests.jl index e29a3eb..4e6c475 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,3 +17,5 @@ for file in readdir(@__DIR__) include(joinpath(@__DIR__, file)) end end + +include(joinpath(dirname(@__DIR__), "examples", "runtests.jl"))