Interaction between broadcasting and custom indexing

Hi,

I have some weird interaction between broadcasting and custom indices for a type. This is a minimal example.

Say I have this type, which holds data in a matrix format

struct Foo <: AbstractMatrix{Float64}
    data::Matrix{Float64}
    nm::Tuple{Int, Int}
    Foo(data) = new(data, size(data))
end

and I then implement the custom indexing behaviour I need using Cartesian indices


# reflect indices
@inline mapindices(k::Int, j::Int) = 
    (ifelse(k≥0, k, -k)+1, ifelse(j≥0, j, -j)+1)

# ~~ cartesian indexing ~~
@noinline function Base.getindex(f::Foo, k::Int, j::Int)
    f.data[mapindices(k, j)...]
end

@noinline function Base.setindex!(f::Foo, val, k::Int, j::Int)
    f.data[mapindices(k, j)...] = val
    val
end

In some cases, however, I prefer to use linear indices for performance, so I define:

# ~~ linear indexing ~~
@noinline function Base.getindex(f::Foo, i::Int)
    f.data[i]
end

@noinline function Base.setindex!(f::Foo, val, i::Int)
    f.data[i] = val
    val
end

I then implement linearindices and IndexStyle, so that by default the IndexLinear behaviour is used, over the indices of the underlying data, while indices allows bounds checking when the special cartesian indexing I have defined above is used.

function Base.indices(f::Foo)
    n, m = f.nm 
    return -n+1:n-1, -m+1:m-1
end
Base.linearindices(f::Foo) = linearindices(f.data)

# tell broadcast to use fast linear ndexing
Base.IndexStyle(::Type{Foo}) = IndexLinear()

Here is where my issue arises. I might not be understanding completely how indexing and broadcast interact, so here is the problem.

# ~~~ tests ~~~
# this represents the data
# 4 3 4
# 2 1 2
# 4 3 4
u = Foo([1 2; 
         3 4])
# this is what the `mapindices` function above does, it just mirrors the data.

# here is some other data
b = Foo([1 2; 
         3 4])

# this is what it looks like
display(u); println()
@show indices(u)
@show linearindices(u)
@show IndexStyle(u)

Now, when I use the new dot syntax in v0.6, I want broadcast to use the fast linear indexing, to loop over the underlying data efficiently. However

# use broadcast
u .= u .+ b

# results in
# ERROR: LoadError: BoundsError: attempt to access 2×2 Array{Float64,2} at index [5]

where it seems to me that it is trying to access the underlying data at index 5, ignoring my definition of linearindices.

Any help is welcome.

Davide

1 Like

Bump.

This is a show stopper at the moment for me. Has anyone successfully integrated broadcasting behaviour and custom indices?