User defined operator in JuMP when one of the arguments is data

Hi there,

I need to solve for a MLE problem where the likelihood function I defined reads in an Array{Obs{Float64}} where Obs{Float64} is an immutable struct containing a bunch of numbers.

Since my problem involves a constraint, I am trying to solve it via JuMP but its seems like non-linear user defined operators can only be defined with the args being the optimizing variables.

ll(x)  = ll(x, data)

Causes type-instability and therefore makes the function evaluation much slower. Is there an easy way around this? It seems like user defined functions in general might have the same issue.

1 Like

Cann you provide a reproducible example of what you are doing?

Sorry for the delayed response. It did not allow me to edit the post at this point so my reply will contain the reproducible example.

Consider the problem of finding the inverse matrix of another matrix M. A simple implementation, friendly with JuMP syntax, would be:

M = [1.0  3.0; 3.0  4.0]
n = size(M,1)
I_n = Matrix{eltype(M)}(I, n, n)
function norm_diff(args...)
    # Collect the scalar inputs into a vector
    A_vec = collect(args)
    # Reshape the vector into an n×n matrix
    A = reshape(A_vec, n, n)
    # Compute the difference A*M - I
    diff = A * M - I_n
    return mapreduce(x -> x^2, +, diff)
end
model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
register(model, :norm_diff, n*n, norm_diff; autodiff=true)
@objective(model, Min, norm_diff(A...))
optimize!(model)

this works fine but in my application I am looking at optimizing a user-defined function that uses some data to compute the scalar objective value. A modification of the above in lines with the feature described would be:

parameters = (
            n = n,
)
function norm_diff_data(args...; M, I_n, params)

    (; n) = params 
    A_temp = zeros(typeof(args[1]), n, n)
    diff = similar(A_temp)
    # Collect the scalar inputs into a vector
    for i in eachindex(A_temp)
        A_temp[i] = args[i]
    end
    # Compute the difference A*M - I
    mul!(diff, A_temp, M)
    @. diff -= I_n
    return mapreduce(x -> x^2, +, diff)
end

However, the problem is that I am unsure about how to register such function given that now, arguments are a combination variables over which to optimize and data. A work-around this is to wrap the above as:

f(args...) = norm_diff_data(args...; M = M, I_n = I_n, params = parameters)
model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
register(model, :f, n*n, f; autodiff=true)
@NLobjective(model, Min, f(A...))
optimize!(model)

This works but clearly f is type-unstable which makes its much slower than needed and this could be a big issue when the wrapped around function (in this case norm_diff_data) is expensive.

So my question is, is there a way around this? It seems like the need to optimize a user-defined function that uses data is not niche.

1 Like

Your first example works because you have mixed the legacy nonlinear syntax register with the new function tracing syntax. So when you did:

julia> @objective(model, Min, norm_diff(A...))
10 A[1,1]² + 30 A[1,2]*A[1,1] + 25 A[1,2]² + 10 A[2,1]² + 30 A[2,2]*A[2,1] + 25 A[2,2]² - 2 A[1,1] - 6 A[1,2] - 6 A[2,1] - 8 A[2,2] + 2

we actually evaluated norm_diff with JuMP VariableRef as the input and got back an algebraic expression. We did not use the registered user-defined operator.

The correct syntax is:

model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
@operator(model, op_norm_diff, n^2, norm_diff)
@objective(model, Min, op_norm_diff(A...))
optimize!(model)

For your second function, it looks pretty good. But the JuMP code should be:

model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
@operator(model, op_f, n^2, f)
@objective(model, Min, op_f(A...))
optimize!(model)

This works but clearly f is type-unstable which makes its much slower than needed

Do you have evidence that the type stability makes Ipopt run slower? It probably doesn’t really matter.

I would tweak your example to be:

norm_diff_data2(A, M, I_n) = sum((A * M - I_n).^2)
g(args...) = norm_diff_data2(reshape(collect(args), n, n), M, I_n)
model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
@operator(model, op_g, n^2, g)
@objective(model, Min, op_g(A...))
optimize!(model)

but I might write it as

norm_diff_data2(A, M, I_n) = sum((A * M - I_n).^2)
model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
@objective(model, Min, norm_diff_data2(A, M, I_n))
optimize!(model)

or even

model = Model(Ipopt.Optimizer)
@variable(model, A[1:n,1:n])
@variable(model, diff[1:n,1:n])
@constraint(model, diff .== A * M .- I_n)
@objective(model, Min, sum(diff.^2))
optimize!(model)

Sorry for the delayed response. It did not allow me to edit the post at this point so my reply will contain the reproducible example.

No worries :smile: we’re not StackOverflow, so it’s even encouraged to post replies rather than editing old posts.

1 Like

Thanks a lot for the quick and detailed reply (as usual), Oscar!

Actually, I do not. That I took for granted since I though that the solving time would be an increasing function (significantly sensitive too) of the objective function evaluation time. But most likely I’m wrong. In the example I used here it does not matter. I will see whether it does in my bigger application.

Interesting. I’m guessing that this only works because norm_diff_data2 has a simple algebraic representation? In the case of a more complicated user defined function I need a functions thats takes f(args...) right?

1 Like

Interesting. I’m guessing that this only works because norm_diff_data2 has a simple algebraic representation?

Correct

I will see whether it does in my bigger application.

:+1: I wouldn’t spend effort worrying until you have something that works end-to-end and the runtime is too slow to be useful.

1 Like