function calls with splatted variables not working within constraint()

[jsphchoi]

Can not use splatted ExaModel variables over colon inside function calls within constraint()

using ExaModels, MadNLP

V = 5
N = 10
K = 3

c = ExaCore()
z = variable(c, 1:V, 1:N, 0:K; start = 1.0)
h = exp10.(range(0,1,N))

# CASE 1: this works (expected)
itr_coll = [(i,k,h[i]) for i in 1:N, k in 1:K]
ftest(x...) = -x[1]*2 + x[2]*x[3] - x[4] + x[5]
constraint(c, -hi*ftest(z[1,i,k],z[2,i,k],z[3,i,k],z[4,i,k],z[5,i,k]) for (i,k,hi) in itr_coll)

# CASE 2: this works (expected)
itr_coll = [(i,k,h[i]) for i in 1:N, k in 1:K]
ftest2(x1,x2,x3,x4,x5) = -x1*2 + x2*x3 - x4 + x5
constraint(c, -hi*ftest2(z[1,i,k],z[2,i,k],z[3,i,k],z[4,i,k],z[5,i,k]) for (i,k,hi) in itr_coll)

# (IDEAL WORKING CASE) CASE 3: this doesn't work (not expected)
itr_coll = [(i,k,h[i]) for i in 1:N, k in 1:K]
ftest(x...) = -x[1]*2 + x[2]*x[3] - x[4] + x[5]
constraint(c, -hi*ftest(z[:,i,k]...) for (i,k,hi) in itr_coll)

# CASE 4: this doesn't work (not expected)
itr_coll = [(i,k,h[i]) for i in 1:N, k in 1:K]
ftest2(x1,x2,x3,x4,x5) = -x1*2 + x2*x3 - x4 + x5
constraint(c, -hi*ftest2(z[:,i,k]...) for (i,k,hi) in itr_coll)

######################################################
objective(c, (z[1,i,0] - 1)^2 for i in 1:N)
m = ExaModel(c)
madnlp(m)

[sshin23]

Hey @jsphchoi 3 is indeed nice, and we can support it

[jsphchoi]

function _create_lagrange_collocation(c::ExaCore, PEprob::PEtabODEProblem, p, Np, z, Nz, N, K, Nc)
    # Lagrange basis polynomial collocation equations (10.7)
    fs = _get_rhs_funcs(PEprob)
    fexprs = [
        (i,k,cidx) -> f(ntuple(j -> z[j,i,k,cidx], Nz)..., ntuple(j -> p[j], Np)...)
        for f in fs
    ]
    itr_coll = [(i,k,cidx,h[i]) for i in 1:N, k in 1:K, cidx in 1:Nc]
    c_coll = [
        constraint(c, -hi*f(i,k,cidx) for (i,k,cidx,hi) in itr_coll)
        for f in fexprs
    ]
    DLDTAU = [_eval_dldtau(j,k,taus) for j in 0:K, k in 1:K]
    itr_coll! = [(i,j,k,cidx,DLDTAU[j+1,k]) for i in 1:N, j in 0:K, k in 1:K, cidx in 1:Nc]
    for v in 1:Nz
        constraint!(c, c_coll[v],
            (i,k,cidx) => z[v,i,j,cidx]*DLDTAU
            for (i,j,k,cidx,DLDTAU) in itr_coll!
        )
    end
end

This seems to be a fine workaround for now

[sshin23]

I think we should support the following:

c1 = constraint(c, 1:9,1:4)
constraint!(c, g(c1[i,:], x[i,:], x[i+1,:]) for i=1:9)

where

function g(y,x,xp) 
y[1] += x[1] + x[2]
y[2] += x[2] * xp[3]
...
return y
end