Speed up JuMP.jl Estimation When Jacobian and Hessian of Constraints are Sparse

[Strange_Xue]

Hi guys, in last questions (link is shown below), I asked how to supply gradient of objective and jacobian of constraints in nonlinear optimization. But @odow told me I dont have to do that because JuMP.jl will automatically calculate them and feed them to Ipopt.jl. However, my model still runs slowly. So I show my code to find any opportunities to speed up.

Thanks for your help.

# ******************************************************** #
#     MATHEMATICAL PROGRAM WITH EQUILIBRIUM CONSTRAINTS    #
# ******************************************************** #
using Base.Threads
using Random, Distributions, Statistics
using LinearAlgebra, LogExpFunctions
using JuMP, Ipopt, ForwardDiff
const seed = Random.seed!(2023 - 11 - 24)

# ******************************************************** #
#                      DATA GENERATION                     #
# ******************************************************** #

Ps = 25 # number of products
Mkts = 40 # number of markets
Inds = 100 # number of individuals in each market
vcovXc = [
	[1, -0.8, 0.3] [-0.8, 1, 0.3] [0.3, 0.3, 1]
]
E𝛽 = [-1.0, 1.5, 1.5, 0.5, -3.0] # 𝜃1
𝜎𝛽 = sqrt.([0.5, 0.5, 0.5, 0.5, 0.2])  # 𝜃2
𝛽i = rand(seed, MvNormal(𝜎𝛽), Inds)

struct Market
	J::Int64         # number of products
	𝑃::AbstractVector{<:Real} # price vector, [J, 1]
	Xc::AbstractMatrix{<:Real} # observed product characteristics, [J, B]
	share::AbstractVector{<:Real} # market share exluding outshare, [J, 1]
	out_share::Float64
	v::AbstractMatrix{<:Real} # simulated individuals, [B, Inds]
	iv::AbstractMatrix{<:Real} # instrumental variables
end

d = Vector{Market}(undef, Mkts);
Xc = rand(seed, MvNormal(vcovXc), Ps)'
let
	for mkt ∈ eachindex(d)
		𝜉, e = randn(seed, Ps), randn(seed, Ps)

		price = abs.(0.5 * 𝜉 + e + 1.1 * sum(Xc; dims = 2)[:])
		Z = rand(seed, Ps, 6) .+ 0.25 * (e + 1.1 * sum(Xc; dims = 2)[:])
		IV = hcat(ones(Ps), Z, Xc, Z .^ 2, Z .^ 3, Xc .^ 2, Xc .^ 3, sum(Z; dims = 2), sum(Xc; dims = 2), kron(Xc[:, 1], ones(size(Z, 1))') * Z, kron(Xc[:, 2], ones(size(Z, 1))') * Z)

		X = hcat(ones(Ps), Xc, price)
		U = X * E𝛽 .+ X * 𝛽i
		denom = 1 .+ sum(exp.(U); dims = 1)
		pchoice = exp.(U) ./ denom
		share = mean(pchoice; dims = 2)[:]
		out_share = 1 - sum(share)
		v = rand(seed, MvNormal(ones(size(X, 2))), Inds)
		d[mkt] = Market(Ps, price, Xc, share, out_share, v, IV)
	end
end
# ======================================================== #
#                        ESTIMATION                        #
# ======================================================== #
model = Model(Ipopt.Optimizer)
@variable(model, theta2[1:length(𝜎𝛽)], start = 1.0)
@variable(model, theta1[1:length(E𝛽)])
𝜉 = Vector{Vector{VariableRef}}(undef, Mkts)
for mkt in 1:Mkts
	𝜉[mkt] =  @variable(model, [1:Ps], base_name = "𝜉$(mkt)")
end

pXc = Vector{Matrix{Float64}}(undef, Mkts)
for m in 1:Mkts
	pXc[m] = hcat(ones(d[m].J), d[m].Xc, d[m].𝑃)
end

𝜇 = @expression(model, [m = 1:Mkts], pXc[m] * Diagonal(theta2) * d[m].v);
𝛿 = @expression(model, [m = 1:Mkts], pXc[m] * theta1 .+ 𝜉[m]);
𝑈 = @expression(model, [m = 1:Mkts], 𝛿[m] .+ 𝜇[m]);
exp𝑈 = @expression(model, [m = 1:Mkts], exp.(𝑈[m]));
denom = @expression(model, [m = 1:Mkts], mapreduce(x -> 1 + sum(x), hcat, eachcol(exp𝑈[m])));
pchoice = @expression(model, [m = 1:Mkts], exp𝑈[m] ./ denom[m]);
share_eq = @constraint(model, [m = 1:Mkts, p = 1:Ps], mean(pchoice[m][p, :]) == d[m].share[p]);
𝜉_vec = @expression(model, reduce(vcat, 𝜉))
IV = reduce(vcat, getproperty.(d, :iv))
𝛷 = IV' * IV

EM = @expression(model, mean(IV .* 𝜉_vec; dims = 1));
obj = @expression(model, EM * 𝛷 * EM');
@objective(model, Min,  obj[1, 1]);
optimize!(model)
value.(theta1)
value.(theta2)

JuMP.jl: Supplying gradient of objective and jacobian of constraints in Nonlinear Optimization

Hi, I am faced with a troublesome optimization problem. As shown below, the non-linear optimization problem is easy to model and solve via Ipopt with JuMP.jl as interface. However, the objective function and constraints are time-consuming to compute, even with the assistance of auto-differenciation embedded in JuMP.jl. \text{min} \; f(x)\\ \text{s.t.}\; g(x) = a \\ \quad \; s(x) = b\\ However, the jacobian and hessian of constraints is very sparse. Someone has provided the analytical jacobia…

[odow]

This is a rather nasty example that stresses a lot of JuMP. Your constraints are pretty dense!

I’ll take a look, but you might want to investigate alternative modeling packages, like those mentioned in your other thread.

[Strange_Xue]

many thanks to your attention! I will try use NLPModels.jl.

[odow]

The issue is this function in the constraints:

function foo(X::Matrix)
    U = exp.(X)
    d = 1 .+ sum(U; dims = 1)
    return Statistics.mean(U ./ d; dims = 2)[:]
end

JuMP doesn’t have vector-valued nonlinear functions, so we scalarize, and then you end up with Mkts nonlinear equality constraints, each of which is dense.

I tried a few things without success. Here’s one approach:

using JuMP
import Distributions
import Ipopt
import LinearAlgebra
import Random
import Statistics

const seed = Random.seed!(2023 - 11 - 24)

# ******************************************************** #
#                      DATA GENERATION                     #
# ******************************************************** #

Ps = 25
Mkts = 40
Inds = 100
vcovXc = [[1, -0.8, 0.3] [-0.8, 1, 0.3] [0.3, 0.3, 1]]
E𝛽 = [-1.0, 1.5, 1.5, 0.5, -3.0]
𝜎𝛽 = sqrt.([0.5, 0.5, 0.5, 0.5, 0.2])
𝛽i = rand(seed, Distributions.MvNormal(𝜎𝛽), Inds)
Xc = rand(seed, Distributions.MvNormal(vcovXc), Ps)'

struct Market
	J::Int64
	𝑃::Vector{Float64}
	Xc::Matrix{Float64}
	share::Vector{Float64}
	out_share::Float64
	v::Matrix{Float64}
	iv::Matrix{Float64}
end

function compute_pchoice(U::Matrix)
    expU = exp.(U)
    denom = mapreduce(x -> 1 + sum(x), hcat, eachcol(expU))
    return expU ./ denom
end

d = Vector{Market}(undef, Mkts);
let
    for mkt in 1:Mkts
        𝜉, e = randn(seed, Ps), randn(seed, Ps)
        price = abs.(0.5 * 𝜉 + e + 1.1 * sum(Xc; dims = 2)[:])
        Z = rand(seed, Ps, 6) .+ 0.25 * (e + 1.1 * sum(Xc; dims = 2)[:])
        IV = hcat(
            ones(Ps),
            Z,
            Xc,
            Z .^ 2,
            Z .^ 3,
            Xc .^ 2,
            Xc .^ 3,
            sum(Z; dims = 2),
            sum(Xc; dims = 2),
            kron(Xc[:, 1], ones(size(Z, 1))') * Z,
            kron(Xc[:, 2], ones(size(Z, 1))') * Z,
        )
        X = hcat(ones(Ps), Xc, price)
        U = X * E𝛽 .+ X * 𝛽i
        pchoice = compute_pchoice(U)
        share = Statistics.mean(pchoice; dims = 2)[:]
        out_share = 1 - sum(share)
        v = rand(seed, Distributions.MvNormal(ones(size(X, 2))), Inds)
        d[mkt] = Market(Ps, price, Xc, share, out_share, v, IV)
    end
end
pXc = [hcat(ones(d[m].J), d[m].Xc, d[m].𝑃) for m in 1:Mkts]

IV = reduce(vcat, getproperty.(d, :iv))
𝛷 = IV' * IV

# ======================================================== #
#                        ESTIMATION                        #
# ======================================================== #

@time begin
    model = direct_model(Ipopt.Optimizer())
    @variable(model, theta2[1:length(𝜎𝛽)], start = 1.0)
    @variable(model, theta1[1:length(E𝛽)], start = 0.0)
    @variable(model, 𝜉[1:Ps, 1:Mkts], start = 0.0)
    function foo(X::Matrix)
        U = exp.(X)
        d = 1 .+ sum(U; dims = 1)
        return Statistics.mean(U ./ d; dims = 2)[:]
    end
    for m in 1:Mkts
        𝜇 = pXc[m] * LinearAlgebra.Diagonal(theta2) * d[m].v
        𝛿 = @expression(model, pXc[m] * theta1 .+ 𝜉[:, m])
        @constraint(model, foo(𝛿 .+ 𝜇) .== d[m].share)
    end
    # @expression(model, EM, Statistics.mean(IV .* vec(𝜉); dims = 1)[:])
    M, N = size(IV)
    @variable(model, EM[1:N])
    @constraint(model, [i in 1:N], M * EM[i] == IV[:, i]' * vec(𝜉))
    # obj = @expression(model, EM * 𝛷 * EM');
    # @objective(model, Min,  obj[1, 1]);
    @variable(model, residuals[1:M])
    @constraint(model, residuals .== IV * EM)
    @objective(model, Min, sum(residuals.^2));
    @info "Calling optimize"
    optimize!(model)
    model
end

Ipopt runs into a bunch of The equality constraints contain an invalid number errors. You should add appropriate bounds to theta2 and theta1.

[Strange_Xue]

Thanks. Its OK. I will check the mathematical model and the code. If I find the answer some day, I will post it in this thread. Thank you.

[gdalle]

Just FYI, if you need to construct sparse Jacobians and Hessians to feed to JuMP because its own internal AD struggles, you can use SparseDiffTools.jl or the (still experimental) DifferentiationInterface.jl.

[tmigot]

Hi @Strange_Xue !

If you plan on passing manually the Jacobian and Hessian derivatives, then you could use ManualNLPModels.

using ManualNLPModels, NLPModelsIpopt
nlp = NLPModel(x, lvar, uvar, f; kwargs...) # to be completed, https://jso.dev/ManualNLPModels.jl/dev/reference/#ManualNLPModels.NLPModel
stats = ipopt(nlp)

You could use AD indeed, ADNLPModels makes the bridge for you to use ipopt, but I am suspecting that if it is difficult for JuMP this might not improve much.