VectorNonlinearOracle

[odow]

This is hopefully a summary of a good chat that I had with @Robbybp today.

Motivation

We have a standard Ipopt NLP:

min  f(x, y)
s.t. l <= g(x, y) <= u

that we want to add a vector-valued nonlinear function to:

min  f(x, y)
s.t. l <= g(x, y) <= u
     y = F(x)

Vector-valued user-defined functions crop up all the time on the forum, especially in the optimal control community. Currently we tell them to follow https://jump.dev/JuMP.jl/dev/tutorials/nonlinear/tips_and_tricks/, which is pretty hacky.

We currently support this in MathOptAI: https://github.com/lanl-ansi/MathOptAI.jl/blob/53121ff5fbdae73f4750272d2545000bf5c67e97/src/predictors/GrayBox.jl#L7-L140

This works, but the downside is that when computing the Hessian-of-the-Lagrangian we compute one (dense) Hessian for each row in the output of y. The Hessian evaluation is a bottleneck in @Robbybp's work.

He has a solution that hacks abuses the internal API of Ipopt.jl to take the y = F(x) constraints onto the end of the NLP evaluator, and there is a fast way in PyTorch to compute $u' F_{xx}(x)$ which is used in eval_hessian_lagrangian. But it'd be nice if there was first-class support in Ipopt.

Solution

Define a new set in Ipopt.jl

struct VectorNonlinearOracle <: MOI.AbstractVectorSet
    input_dimension::Int
    output_dimension::Int
    f::Function
    jacobian::Function
    jacobian_structure::Vector{Tuple{Int,Int}}
    hessian_of_the_lagrangian::Function
    hessian_structure::Vector{Tuple{Int,Int}}
end

Make Ipopt support

MOI.add_constraint(::Optimizer, ::MOI.VectorOfVariables, ::MOI.VectorNonlinearOracle)

This would tack on the appropriate calls to

Ipopt.jl/src/MOI_wrapper.jl

Lines 919 to 925 in 6bde96f

	function MOI.eval_hessian_lagrangian(model::Optimizer, H, x, σ, μ)
	offset = MOI.eval_hessian_lagrangian(model.qp_data, H, x, σ, μ)
	H_nlp = view(H, offset:length(H))
	μ_nlp = view(μ, (length(model.qp_data)+1):length(μ))
	MOI.eval_hessian_lagrangian(model.nlp_data.evaluator, H_nlp, x, σ, μ_nlp)
	return
	end

Here's the API I have in mind:

using JuMP, Ipopt
set = Ipopt.VectorNonlinearOracle(;
    input_dimension = 3,
    output_dimension = 2,
    f = (g, x) -> begin
        g[1] = x[1]^2
        g[2] = x[2]^2 + x[3]^3
        return
    end,
    jacobian = (J, x) -> begin
        J[1] = 2 * x[1]
        J[2] = 2 * x[2]
        J[3] = 3 * x[3]^2
        return
    end,
    jacobian_structure = [(1, 1), (2, 2), (2, 3)],
    hessian_of_the_lagrangian = (H, u, x) -> begin
        H[1] = 2 * u[1]
        H[2] = 2 * u[2]
        H[3] = 6 * x[3] * u[2]
        return
    end,
    hessian_of_the_lagrangian = [(1, 1), (2, 2), (3, 3)],
)
model = Model(Ipopt.Optimizer)
@variable(model, x[1:3])
@variable(model, y[1:2])
@constraint(model, [x; y] in set)
optimize!(model)
x_v = value.(x)
@test value.(y) == [x_v[1]^2, x_v[2]^2 + x_v[3]^3]

Related

There is an open issue in MOI to add vector-valued nonlinear expressions: jump-dev/MathOptInterface.jl#2402. The main thrust is vector-input, but vector-output is also desirable.

If this works, I think it would solve most of the complains for vector-valued nonlinear, and allow us to punt on the implementation a little longer.

Alternatives

We might consider

F(x) = 0

this would implicitly allow

g(y) = F(x)

but it would mean the standard "I have a user-defined function" in JuMP is more complicated.

[Robbybp]

Thanks for the writeup! I verified that my Ipopt hack, my "Pytorch Lagrangian", and MathOptAI all work well together, so next week I'll try to put together a GrayBoxOpt.jl package that demos my approach.

[odow]

Started in #466

I wonder if we should just make this

x in VectorNonlinearOracle()

as short for

l <= f(x) <= u

The f(x) = y oracle would be:

set = Ipopt.VectorNonlinearOracle(;
    input_dimension = 5,
    l = zeros(2),
    u = zeros(2),
    f = (g, x) -> begin
        g[1] = x[1]^2 - x[4]
        g[2] = x[2]^2 + x[3]^3 - x[5]
        return
    end,
    jacobian_structure = [(1, 1), (2, 2), (2, 3), (1, 4), (2, 5)],
    jacobian = (J, x) -> begin
        J[1] = 2 * x[1]
        J[2] = 2 * x[2]
        J[3] = 3 * x[3]^2
        J[4] = -1.0
        J[5] = -1.0
        return
    end,
    hessian_lagrangian_structure = [(1, 1), (2, 2), (3, 3)],
    hessian_lagrangian = (H, x, u) -> begin
        H[1] = 2 * u[1]
        H[2] = 2 * u[2]
        H[3] = 6 * x[3] * u[2]
        return
    end,
)