Slicing Custom Indexes

I am trying to write a custom indexed array (with reflective boundary conditions). I would really like to be able to slice the array in the same way that the amazing developers of ShiftedArrays did.

x = reshape(1:100,10,10)
x1 = ShiftedArrays.circshift(x,(1,1))
x1[-10:10,-10:10]

which just gives an array using the index look-up scheme.

Doing what I thought was copying their approach, I modified code from other custom indexing posts here and came to the following.

module tmp
    struct ReflectArray{T, N} <: AbstractArray{T, N}
        parent::AbstractArray{T,N}
        function ReflectArray(x::AbstractArray{T,N}) where {T,N} 
            return new{T,N}(x)
        end
    end
    
    function reflectindex(i::Int,N::Int)
        if (0 < i <= N)
            return i
        elseif (-N < i <= 0)
            return -i+1
        elseif (N < i <= 2*N)
            return 2*N - i + 1
        else
            error("Outside of first order reflective bounds!")
        end
    end

    Base.size(A::ReflectArray) = size(A.parent)

    function Base.getindex(A::ReflectArray{T, N}, ind::Vararg{Int, N}) where {T, N}
        v = A.parent
        @boundscheck checkbounds(v, ind...)
        i = map(reflectindex, ind, size(A))
        @inbounds ret = v[i...]
        ret
    end
end

This works for grabbing single elements out of an array, but the same slicing call as above is failing. Any help/thoughts as to why.

c = tmp.ReflectArray(collect(x));
c[-10:10,-10:10]
BoundsError: attempt to access 10Ă—10 Main.tmp.ReflectArray{Int64,2} at index [-10:10, -10:10]

Edit: I was working off of ShiftedArrays.jl/circshiftedarray.jl at master · JuliaArrays/ShiftedArrays.jl · GitHub assuming all of the relevant code was there.

1 Like

I haven’t dug into the packages you mentioned, but based on the code you have it seems you are doing a boundscheck which should be the source of the error you see. I don’t think you want a boundscheck here since the whole idea is that you should be able to index outside of the bounds, so I would remove it.

You also only dispatch on Int as the index argument to getindex, you might want a UnitRange here to allow for slices.

Got curious and got something to work as I understood your instructions. Also updated the reflectindex to be more general.

struct ReflectArray{T, N} <: AbstractArray{T, N}
    parent::AbstractArray{T,N}
    function ReflectArray(x::AbstractArray{T,N}) where {T,N} 
        return new{T,N}(x)
    end
end

function reflectindex(i::Int,N::Int)
    i -= 1 # Starting at zero makes mod easier
    i = mod(i, 2N)
    if i < N
        # Here we are in normal mode, and can just return i with fixed index
        return i + 1
    else
        # Here we are in reflected mode, and need to flip and fix index
        return 2N - i
    end
end

function reflectindex(i::UnitRange, N::Int) 
    # This captures the ranges, and will then broadcast over each index in them
    reflectindex.(i, N)
end

Base.size(A::ReflectArray) = size(A.parent)

# Now it takes a Union{Int, UnitRange}, not sure if splitting them in two functions is better for some reason but this seems to work fine
function Base.getindex(A::ReflectArray{T, N}, ind::Vararg{Union{Int, UnitRange}, N}) where {T, N}
    v = A.parent
    #@boundscheck checkbounds(v, ind...)
    i = map(reflectindex, ind, size(A))
    @inbounds ret = v[i...]
    ret
end
3 Likes

This is very helpful thank you. I have one follow up question. I noticed that this implementation does not seem to support views. Can someone point out what needs to be done so that the following also does not throw a bounds error? (Note that again I am taking ShiftedArrays as my point of comparison where this does work).

@views c[-10:10,-10:10]

To get a view to work, I just specialized view to not check bounds in the spirit of the solution’s suggestion. I suspect there is a better/safer way to do this, but this should be fine for now. If someone knows the safer way, I would love to learn however! Thanks so much for all the help!

function Base.view(A::ReflectArray, I::Vararg{Any,N}) where {N}
     J = map(i->Base.unalias(A,i), to_indices(A, I))
     Base.unsafe_view(Base._maybe_reshape_parent(A, Base.index_ndims(J...)), J...)
end

I think you can avoid defining methods for view etc, by just disabling bounds checks completely – there are no indices which are out of bounds here:

Base.checkbounds(::Type{Bool}, ::ReflectArray, i...) = true

Then you should only need the getindex(A::ReflectArray{T, N}, ind::Vararg{Int, N}) method.

2 Likes

Ah, that works for me and is much cleaner!