Registering user-defined functions in JuMP for mixed-integer non-linear programming

Hello everyone,
I have been a Julia user for some years now and I really love JuMP, however I am finding particularly difficult setting up a NLobjective in a mixed-integer program with linear constraints which corresponds to a dynamic program. The part of the code causing trouble is the following

using LinearAlgebra
function ApproximateValueH(State,BasisFunctions,BasicHorizon_ind)
    n=length(BasisFunctions[1])
    f(y,u,g,z)=y+dot(g,z-u)
    Linear=[]
    for i in 1:n
        push!(Linear,f(BasisFunctions[BasicHorizon_ind][i][2],BasisFunctions[BasicHorizon_ind][i][1][:,2],BasisFunctions[BasicHorizon_ind][i][1][:,1],State))
    end
    return maximum(Linear)
end

function ApproximateValue(Weights,BasisFunctions,Horizons,State)
    Approximations=[]
    for H in Horizons
        Hix=findall(Horizons.==H)[1]
        push!(Approximations,ApproximateValueH(State,BasisFunctions,Hix))
    end
    return dot(Weights,Approximations)
end

Horizons=[2,5,15,30,45]
Weights=1/length(Horizons)*ones(length(Horizons))
BasisFunctions=[[[rand(28*33,2), rand()] for j in 1:2] for i in 1:length(Horizons)] #Made up for MWE
using JuMP
import KNITRO

model=Model(KNITRO.Optimizer)
@variable(model, NewState[l=1:28*33]>=0, Int)

# Other variables and linear constraints here

NextStateValue(X...)=ApproximateValue(Weights,BasisFunctions,Horizons,X...)
register(model,:NextValueS, length(NewState), NextStateValue; autodiff=true)

The error message is very long but the key info reads as

Unable to register the function :NextValueS because it does not support differentiation via ForwardDiff.
Common reasons for this include:
# explanation about how to debug ForwardDiff
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:33
 [2] _validate_register_assumptions(f::typeof(NextStateValue), name::Symbol, dimension::Int64)
   @ JuMP ~/.julia/packages/JuMP/0C6kd/src/nlp.jl:1979
 [3] register(m::Model, s::Symbol, dimension::Int64, f::Function; autodiff::Bool)
   @ JuMP ~/.julia/packages/JuMP/0C6kd/src/nlp.jl:2052
 [4] top-level scope
   @ REPL[25]:1

caused by: MethodError: no method matching ApproximateValue(::Vector{Float64}, ::Vector{Vector{Vector{Any}}}, ::Vector{Int64}, ::ForwardDiff.Dual{ForwardDiff.Tag{JuMP.var"#137#138"{typeof(NextStateValue)}, Float64}, Float64, 12},...

However, running

ForwardDiff.gradient(NextStateValue,rand(28*33))

returns an answer. I would appreciate any thoughts about what is causing the trouble: I suspect it has to do with the fact that ApproximateValueH take a max over a vector, but I am not sure. How does JuMP work when one tries to register a function which calls other functions? Should I register all of them? Thank you very much and I apologize for the long post!

Hi @gmantegazza,

This simplest answer is that you’re slightly off with your splatting.

# Current function
NextStateValue(X...)=ApproximateValue(Weights,BasisFunctions,Horizons,X...)

# Change it to:
NextStateValue(X...)=ApproximateValue(Weights,BasisFunctions,Horizons,X)

The first function is going to try to “splat” X into multiple arguments, where as the second is going to pass X as a tuple. This is a very subtle part of Julia:

julia> bar(x...) = length(x), x
bar (generic function with 1 method)

julia> foo1(x...) = bar(x...)
foo1 (generic function with 1 method)

julia> foo2(x...) = bar(x)
foo2 (generic function with 1 method)

julia> foo1(1, 2)
(2, (1, 2))

julia> foo2(1, 2)
(1, ((1, 2),))

foo1 is calling bar(1, 2), where as foo2 is calling bar((1, 2)). The simplest rule of thumb is either both the function definition and the call need a splat (like ApproximateValue(Weights,BasisFunctions,Horizons,State...), or none of them do.

I’m not really sure I understand your function (what is the findall(Horizons trying to do?), perhaps this rewrite will help you out:

using JuMP
using LinearAlgebra

function ApproximateValueH(State, BasisFunctions, BasicHorizon_ind)
    f(y, u, g, z) = y + sum(g[i] * (z[i] - u[i]) for i in 1:length(z))
    return maximum(
        f(b2, b1[:, 2], b1[:, 1], State) for 
        (b1, b2) in BasisFunctions[BasicHorizon_ind]
    )
end

N = 5
Horizons = [2, 5, 15, 30, 45]
Weights = ones(length(Horizons)) ./ length(Horizons)
BasisFunctions = [[[rand(N,2), rand()] for j in 1:2] for i in 1:length(Horizons)]

function NextStateValue(X...)
    return sum(
        Weights[Hix] * ApproximateValueH(State, BasisFunctions, Hix)
        for (Hix, H) in enumerate(Horizons)
    )
end

model = Model()
@variable(model, NewState[1:N] >= 0)
register(model, :NextValueS, N, NextStateValue; autodiff = true)
@NLexpression(model, my_expr, NextValueS(NewState...))

You might also run into some performance problems with a user-defined function with such a large number of arguments, but that’s a secondary consideration after you get something working first.

Hello @odow, thank you very much for your answer! I realize the subtle difference you hint at, will read more carefully the section on splatting in the docs. With your modifications the code runs smoothly, I really appreciate the time.

PS: the findall call is just bad coding for what you got to do with enumerate.

1 Like

I think, after looking more at your code, that you could simplify things down even further:

using JuMP

N = 5
horizons = [2, 5, 15, 30, 45]
H = length(horizons)
weights = fill(1 / H, N)
basis_functions = [rand() => rand(N, 2) for _ in 1:2, _ in 1:H]

model = Model()
@variable(model, x[1:N] >= 0)
@NLexpression(
    model,
    sum(
        weights[h] * maximum(
            θ₀ + sum(θ[i, 1] * (x[i] - θ[i, 2]) for i in 1:N)
            for (θ₀, θ) in basis_functions[:, h]
        ) for h in 1:H
    )
)

In general, you’re better off writing out expressions algebraically in JuMP. We can use a better form of differentiation, and we can even compute second derivatives. If you use a user-defined function, we use FowardDiff for the gradient, and we disable Hessian information.

1 Like

Thanks for showing this coincise alternative!

1 Like