Support indexing A[I1, i, I2] for CartesianIndexes I1, I2

[timholy]

This is useful for writing algorithms that perform a calculation along a particular dimension. Not only are the resulting algorithms much prettier, but it's easy to make them cache-friendly even when the calculation is not along dimension 1. This is in contrast with mapslices, for example, which slices out 1-dimensional vectors before performing the algorithm, and hence is not cache-friendly.

It's worth acknowledging that this is another step along the slow march towards support for general getindex(A, indexes::Union(Integer, CartesianIndex)...). But currently we can't add that without causing massive ambiguity warnings in packages, so until #10525 gets merged I think we'd better stick with special cases.

Demo script:

## Cache-unfriendly version (like mapslices, but optimized for 1d)
function myfilt_slow!(b::Number, A, dim)
    indexes = Any[Colon() for i = 1:ndims(A)]
    sz = ntuple(ndims(A), d->d==dim ? 1 : size(A,d))::NTuple{ndims(A),Int}
    for I in CartesianRange(sz)
        for d = 1:ndims(A)
            if d != dim
                indexes[d] = I[d]
            end
        end
        v = slice(A, indexes...)
        myfilt1d!(b, v)
    end
    A
end

function myfilt1d!(b, v)
    for i = 2:length(v)
        v[i] -= b*v[i-1]
    end
    v
end

## Cache-friendly version
function myfilt!(b::Number, A, dim)
    R1 = CartesianRange(size(A)[1:dim-1])
    R2 = CartesianRange(size(A)[dim+1:end])
    _myfilt!(b, A, R1, R2, size(A,dim))  # for type-stability
end

function _myfilt!(b, A, R1, R2, n)
    for I2 in R2
        for i = 2:n
            for I1 in R1
                A[I1, i, I2] -= b*A[I1, i-1, I2]  # in addition to being faster, isn't this just plain nicer?
            end
        end
    end
    A
end

A = rand(100, 100, 100)
B1 = mapslices(x->myfilt1d!(0.4, x), copy(A), (2,))
B2 = myfilt_slow!(0.4, copy(A), 2)
B3 = myfilt!(0.4, copy(A), 2)
using Base.Test
@test_approx_eq B1 B2
@test_approx_eq B1 B3

Ac = copy(A)
@time 1
@time mapslices(x->myfilt1d!(0.4, x), Ac, (2,))
Ac = copy(A)
@time myfilt_slow!(0.4, Ac, 2)
Ac = copy(A)
@time myfilt!(0.4, Ac, 2)
nothing

Results:

elapsed time: 0.002351136 seconds (35 kB allocated)   # @time 1   (i.e., warmup)
elapsed time: 0.045129864 seconds (20 MB allocated, 7.28% gc time in 1 pauses with 0 full sweep)  # mapslices
elapsed time: 0.020325132 seconds (3 MB allocated)   # myfilt_slow!
elapsed time: 0.003462904 seconds (664 bytes allocated)    # myfilt!

In addition to the usual "array & indexing ninjas" (@Jutho, @mbauman, @carlobaldassi, @JeffBezanson), CCing @tlycken (this will be useful for interp_invert! style operations).

[mbauman]

Very clever. I wasn't sure there was a use case for having multiple cartesian indices… but this is a great example. This seems reasonable until we can define dispatch for the fully generic version without worry.

[mbauman]

One option would be to define the general version on a different function, somewhat in the spirit of #10525. It'd be a whole lot less code that way, and the general version is very straightforward.

[timholy]

You mean, a function named something different from getindex? That would definitely work, but I'm hopeful we'll get #10525 in 0.4, so I'm somewhat inclined to just wait. But open to persuasion.

Also, I should confess that the particular examples I showed could have been written even more simply as A[I-Ioffset], where the iteration on I starts at 2 along dimension dim and Ioffset = CartesianIndex{N}(0,..0,1,0,...,0) is 1 only along dimension dim. Where you really need this functionality (and the example that convinced me to write it) is if you want to do something fancy like traverse dimension dim in reverse order. You need that for LU-backsubstitution or writing filtfilt!.

@andreasnoack, ready to write all those A_ldiv_B!(A, B, dim) methods? (Before you have a heart attack: I am joking. However, I am putting together a FilterMD package that will implement at least Tridiagonal.)

[kmsquire]

Is it possible to write the general version as a different function, and have special cases call the other-named generic version? (That might be what @mbauman was referring to...?)

[mbauman]

Yup, exactly. It's not a big deal, though, and hopefully these definitions don't last long.

[timholy]

Sure, I'll implement it that way.

[mbauman]

I already have some of this written in the context of #10525. I've not pushed it yet (or really tested it thoroughly).

function _expand_cartesian_indices(idx_types, I)
    idxs = Expr[]
    for i=1:length(idx_types)
        if idx_types[i] <: CartesianIndex
            append!(idxs, Expr[:($I[$i][$d]) for d in 1:length(idx_types[i])])
        else
            push!(idxs, :($I[$i]))
        end
    end
    idxs
end
stagedfunction _getindex{T,N}(l::LinearIndexing, A::AbstractArray{T,N}, I::Union(Real,AbstractArray,Colon,CartesianIndex))
    idxs = _expand_cartesian_indices(I, :I)
    if (l <: LinearSlow && length(idxs) == N) || (l <: LinearFast && length(idxs) == 1)
        :($(Expr(:meta, :inline)); getindex(A, $(idxs...))
    else
        :($(Expr(:meta, :inline)); _getindex(l, A, $(idxs...))
    end
end

[timholy]

Ah, I should have looked here again before writing it. But thanks for the head start! If you prefer your longer name, I can change it.

[mbauman]

Nah, your spelling is just fine. 👍

[Jutho]

Looks great to me! I certainly have use cases for multiplying a higher rank tensor with a matrix along a specific dimension, so this will certainly be very useful for implementing this. But it sounds like your FilterMD might be offering this functionality.

[tomasaschan]

Wow. Just wow.

This, along with #10525, will go a very long way to making my life a lot easier when working on Interpolations.jl. I don't have much else to chip in (this is a little above my competence level, and a lot above my time allowance...) but I'm cheering loudly from the sidelines! + 💯 !