Hello!
I’m new to Julia, and I’m working on a constrained optimization problem using JuMP.
I’m having a performance issue when calculating a kronecker product in a function that receives the type ForwardDiff.Dual. In particular, I think I might have a type-stability issue since the kronecker product seems to use a large number of allocations, and amount of memory. However, when using @code_warntype either on the optimization() function or on the functions inside, I wasn’t able to find the problem.
I would greatly appreciate if someone could take a look to see if I’m doing something wrong or if there are tweaks that could improve performance. Below you can see a simple example of my code:
using Pkg
Pkg.add("DataFrames")
Pkg.add("StatFiles")
Pkg.add("Queryverse")
Pkg.add("LinearAlgebra")
Pkg.add("CSV")
Pkg.add("Distributions")
Pkg.add("JuMP")
Pkg.add("Ipopt")
Pkg.add("ForwardDiff")
using DataFrames, StatFiles, Queryverse, LinearAlgebra, CSV, Distributions, JuMP, Ipopt, ForwardDiff
####################################
############ Data ##############
####################################
const s_t = rand(Float64, (20,8000))
const d_t = rand(Float64, (20,8000))
const d_t_lag = rand(Float64, (20,8000))
const i_t = rand(Float64, (120,8000))
const N = size(s_t, 1)
const K = size(s_t, 2)
const L = floor(Int, size(i_t, 1)/N)
####################################
########## Optimization ############
####################################
# We define a simple macro to time the process
macro time_calc_e(expr)
quote
print("calc_e:")
@time $(esc(expr))
end
end
macro time_kronecker(expr)
quote
print("kronecker:")
@time $(esc(expr))
end
end
# Set up JuMP problem
# We use the solver 'Ipopt'
function optimization()
# We calculate the error evaluating the parameters in the calc_e function
function calc_e(β::T, γ::T, W::T...) where {T <: Real}
M = Matrix{T}(reshape([i for i in W], N, N))
e_t = zeros(T, (N, K))
# Note: the for loop to calculate e is faster than the direct matrix multiplication.
@time_calc_e for t in 1:K
e_t[:,t] = s_t[:,t] - β*d_t[:,t] - γ*M*d_t_lag[:,t]
end
g = zeros(T, (L*(N^2), 1))
@time_kronecker for t = 1:K
g .+= @views kron(i_t[:,t], e_t[:,t])
end
G = g'*I*g
return G[1]
end
# We calculate the penalization term ∑|W_i,j|
function penalization_1(W::T...) where {T <: Real}
pen1 = Vector{T}(undef, 1)
pen1 = sum(abs.([i for i in W]))
return pen1[1]
end
# We calculate the penalization term ∑|W_i,j|^2
function penalization_2(W::T...) where {T <: Real}
pen2 = Vector{T}(undef, 1)
pen2 = sum((abs.([i for i in W])).^2)
return pen2[1]
end
# We compute the value of the (penalized) GMM objective function
function gmm_pen(G::T, pen_1::T, pen_2::T) where {T <: Real}
p1 = 0.5
p2 = 0.5
G_pen = G + p1*pen_1 + p2*pen_2
return G_pen
end
# We create the start values for the matrix of interactions
W0 = rand(N,N)
# We create a model by passing the optimizer 'Ipopt' to the Model function
m = Model(Ipopt.Optimizer; add_bridges = false)
# We set some of the attributes of the optimizer 'Ipopt'
set_optimizer_attribute(m, "max_iter", 500)
# We register the function that returns the matrix of errors
JuMP.register(m, :calc_e, 402, calc_e; autodiff = true)
# We register the function that returns the first penalization term
JuMP.register(m, :penalization_1, 400, penalization_1; autodiff = true)
# We register the function that returns the first penalization term
JuMP.register(m, :penalization_2, 400, penalization_2; autodiff = true)
# We register the function that returns the value of the objective function
JuMP.register(m, :gmm_pen, 3, gmm_pen; autodiff = true)
# We create the variables (the objects that change in the optimization)
@variable(m, β, start = -1.0)
@variable(m, γ, start = -1.0)
@variable(m, W[i = 1:N, j = 1:N], start = W0[i,j])
# We create constraints on the variables
@constraint(m, [i=1:N], W[i,i] == 0.0) # The diagonal of W is zero.
@constraint(m, [i=1:N], sum(W[i,:]) == 1.0) # The row-sum of W is one.
# We want to minimize the objective function value
@NLobjective(m, Min, gmm_pen(calc_e(β, γ, W...), penalization_1(W...), penalization_2(W...)))
# We can call print to display the model
print(m)
# To find the solution of the (penalized) GMM we use optimize!
println("Solving...")
optimize!(m)
return m
end
model = optimization()
Thank you very much in advance!