How register a function in JuMP and use a subarray as args

Specific Domains Optimization (Mathematical)
1

I would register constraint somehow like c1 as a user-defined function

using Ipopt
using JuMP
using MathOptInterface
const MOI = MathOptInterface


d1 = 4
d2 = 5
d3 = 6
para_a = rand(d3)

m = Model(optimizer_with_attributes(Ipopt.Optimizer))

@variable(m, uc[i in 1:d1, j in 1:d2, k in 1:d3])
@variable(m, v[i in 1:d1, k in 1:d3])
@variable(m, w[i in 1:d1,  j in 1:d2, k in 1:d3,])

@NLconstraint(m, c1[i in 1:d1, j in 1:d2, k in 1:d3], uc[i, j, k] == v[i, k] + sum(w[i, j, s] * para_a[s] for s in 1:d3))

I tried to implement it and got errors as below.

function f(uc_,  v_, w_...)
    v_ + sum(w_[k] * para_a[k] for k in 1:length(w_)) - uc_
end
register(m, :f, 2 + d3, f,autodiff=true)
@NLconstraint(m, c2[i in 1:d1, j in 1:d2, k in 1:d3], f(uc[i, j, k], v[i, k],  w[i, j, 1:d3]...) == 0)

ERROR: Unexpected expression in f(uc[i, j, k], v[i, k], w[i, j, 1:d3]...). JuMP supports splatting only symbols. For example, x... is ok, but (x + 1)..., [x; y]... and g(f(y)...) are not.

It seems that only a symbol can be used in user-defined function.
Is there any way to use a subarray like w[i, j, :] args?

2

It doesn’t seem to me that there’s any reason to splat w. You should be able to just pass w[i, j, :] as-is. That said, I’ve never had to register a JuMP function…
Also, I personally tend work with anonymous variables and constraints more often than not because I find the [i=1:n, ...] syntax can get cumbersome.

The cleanest (also coolest) way you can do this is with Tullio:

using Tullio
@tullio ex[i,j,k] := v[i, k] + w[i,j,a]*para_a[a]
@constraint(m, c1, uc .== ex)

(I’m using @constraint because your constraints are actually linear)

3

Here’s the workaround that I’d recommend for this particular case:

for i in 1:d1, j in 1:d2, k in 1:d3
  x = w[i, j, 1:d3]
  @NLconstraint(m, f(uc[i, j, k], v[i, k],  x...) == 0)
end
2 Likes