diff --git a/src/nlp.jl b/src/nlp.jl index a7f4bce5..7270ce12 100644 --- a/src/nlp.jl +++ b/src/nlp.jl @@ -919,7 +919,7 @@ end # `Var` wrapping an arithmetic node that would be re-evaluated per row. @inline _indexed_var(i::AbstractArgNode, o) = Var(ArgLeaf(i + o)) @inline _indexed_var(i, o) = Var(i + o) -@inline function Base.getindex(v::V, is...) where {V<:AbstractVariable} +@inline function _getindex(v::V, is, ::Val{false}) where {V<:AbstractVariable} @assert(length(is) == length(v.size), "Variable index dimension error") _bound_check(v.size, is) Var(v.offset + idxx(is .- (_start.(v.size) .- 1), _length.(v.size))) @@ -944,7 +944,7 @@ end idx = idxx(is .- (_start.(s.size) .- 1), _length.(s.size)) return _reindex(s.f, s.iter[idx]) end -@inline function Base.getindex(s::Expression, is...) +@inline function _getindex(s::Expression, is, ::Val{false}) # Symbolic indices case - the symbolic indices ARE the iterator elements # No adjustment needed; the indices are used directly in expression building @assert(length(is) == length(s.size), "Expression index dimension error") @@ -955,12 +955,31 @@ end _bound_check(p.size, i) ParameterNode(i + (p.offset - _start(p.size[1]) + 1)) end -@inline function Base.getindex(p::P, is...) where {P<:Parameter} +@inline function _getindex(p::P, is, ::Val{false}) where {P<:Parameter} @assert(length(is) == length(p.size), "Parameter index dimension error") _bound_check(p.size, is) ParameterNode(p.offset + idxx(is .- (_start.(p.size) .- 1), _length.(p.size))) end +# `x[:, i]` expands each colon axis to its entries, column-major as in Base. +const _Indexable = Union{AbstractVariable,Expression,Parameter} +@inline Base.getindex(x::_Indexable, is...) = _getindex(x, is, _hascolon(is)) +@inline Base.getindex(v::V, ::Colon) where {V<:AbstractVariable} = _allentries(v) +@inline Base.getindex(s::Expression, ::Colon) = _allentries(s) +@inline Base.getindex(p::P, ::Colon) where {P<:Parameter} = _allentries(p) +@inline _allentries(x) = _getindex(x, map(_ -> :, x.size), Val(true)) + +@inline _hascolon(::Tuple{}) = Val(false) +@inline _hascolon(::Tuple{Colon,Vararg{Any}}) = Val(true) +@inline _hascolon(is::Tuple) = _hascolon(Base.tail(is)) + +@inline _axis(n, ::Colon) = _start(n):_start(n)+_length(n)-1 +@inline _axis(n, i) = (i,) +@inline function _getindex(x, is, ::Val{true}) + @assert(length(is) == length(x.size), "Colon index dimension error") + Tuple(x[I...] for I in Iterators.product(map(_axis, x.size, is)...)) +end + function _bound_check(sizes, i::I) where {I<:Integer} __bound_check(sizes[1], i) diff --git a/test/NLPTest/feature_test.jl b/test/NLPTest/feature_test.jl index da533294..cb536424 100644 --- a/test/NLPTest/feature_test.jl +++ b/test/NLPTest/feature_test.jl @@ -198,6 +198,49 @@ function test_nonunit_expr(backend) end end +function test_colon_index(backend) + @testset "Colon expands to the tuple of entries along an axis" begin + N, K = 4, 3 + itr = [(i, k, exp10(i / N)) for i in 1:N, k in 1:K] + f(x...) = -x[1] * 2 + x[2] * x[3] - x[4] + x[5] + g(x, θ) = x[1] * θ[1] - x[5] * θ[2] + function build(con) + c = ExaCore(; backend, concrete = Val(true)) + c, z = add_var(c, 1:5, 1:N, 0:K; start = 1.0) + c, θ = add_par(c, [2.0, 3.0]) + c, _ = con(c, z, θ) + ExaModel(c) + end + m = build( + (c, z, θ) -> + add_con(c, -hi * (f(z[:, i, k]...) + g(z[:, i, k], θ[:])) for (i, k, hi) in itr), + ) + mref = build( + (c, z, θ) -> add_con( + c, + -hi * ( + f(z[1, i, k], z[2, i, k], z[3, i, k], z[4, i, k], z[5, i, k]) + + g((z[1, i, k], z[2, i, k], z[3, i, k], z[4, i, k], z[5, i, k]), (θ[1], θ[2])) + ) for (i, k, hi) in itr + ), + ) + x0 = ExaModels.convert_array([sin(i) for i in 1:m.meta.nvar], backend) + @test Array(NLPModels.cons(m, x0)) ≈ Array(NLPModels.cons(mref, x0)) + + # column-major order, non-unit starts, `x[:]` as in Base, Expression, wrong rank + c = ExaCore(; backend, concrete = Val(true)) + c, x = add_var(c, 2, 3) + c, y = add_var(c, 2:4, 0:1) + c, s = add_expr(c, x[i, j]^2 for (i, j) in Iterators.product(1:2, 1:3)) + @test [v.i for v in x[:, 2]] == [3, 4] + @test [v.i for v in x[2, :]] == [2, 4, 6] + @test [v.i for v in y[:, 1]] == [10, 11, 12] + @test [v.i for v in x[:]] == [v.i for v in x[:, :]] == 1:6 + @test length(s[:, 1]) == 2 + @test_throws Exception x[:, 1, 1] + end +end + function test_generator_free_constr(backend) @testset "add_con(core, expr, itr)" begin c = ExaCore(; backend, concrete = Val(true)) @@ -225,6 +268,9 @@ function test_features(backend) @testset "Non-unit expression indexing" begin test_nonunit_expr(backend) end + @testset "Colon indexing" begin + test_colon_index(backend) + end @testset "Generator-free constraint" begin test_generator_free_constr(backend) end