Particles as indices

[cscherrer]

Hi @baggepinnen ,

For some Soss work, I have something like (simplified example)

julia> a = [Particles(), Particles()+1, Particles()+2]
3-element Array{Particles{Float64,2000},1}:
 -5.33e-18 ± 1.0
  1.0 ± 1.0
  2.0 ± 1.0

julia> i = Particles(DiscreteUniform(1,3))
Particles{Int64,2000}
 1.984 ± 0.821

and I'd like to be able to do

julia> a[i]
Particles{Float64,2000}
 0.995411 ± 1.29

My current approach is

function Base.getindex(a::AbstractArray{P}, i::Particles) where P <: AbstractParticles
    return Particles([a[i.particles[n]][n] for n in eachindex(i.particles)])
end

This works fine, but maybe I have 2, 3, or 4 indices instead of just one. And for n indices we might have some particles and some Ints, with 2^n possibilities.

Do you have a better way to set up this kind of dispatch? Ideally this would be in MCM, but I'm not sure how to even approach the general case

[cscherrer]

For now I just have

function Base.getindex(a::AbstractArray{P}, i::Particles) where P <: AbstractParticles
    return Particles([a[i.particles[n]][n] for n in eachindex(i.particles)])
end

function Base.getindex(a::AbstractArray{P}, i, j::Particles) where P <: AbstractParticles
    return Particles([a[i,j.particles[n]][n] for n in eachindex(j.particles)])
end

function Base.getindex(a::AbstractArray{P}, i::Particles, j) where P <: AbstractParticles
    return Particles([a[i.particles[n], j][n] for n in eachindex(i.particles)])
end

function Base.getindex(a::AbstractArray{P}, i::Particles, j::Particles) where P <: AbstractParticles
    return Particles([a[i.particles[n], j.particles[n]][n] for n in eachindex(i.particles)])
end

Vectors and matrices cover the most common cases, but it would be nice to have something more general. If generalizing would be too complicated I can just use this and extend as I need to.

[baggepinnen]

Hi Chad! Interesting use case :) I think the tricky part is to avoid method ambiguities. Maybe one can get away with a linear number of methods like this

function Base.getindex(a::AbstractArray{P}, i::Particles{<:Integer}, args::Union{AbstractParticles{<:Integer}, Integer}...) where P <: AbstractParticles
    args = promote(i, args...) # all args are now particles
    # logic
end

function Base.getindex(a::AbstractArray{P}, i::Int, j::Particles{<:Integer}, args::Union{AbstractParticles{<:Integer}, Integer}...) where P <: AbstractParticles
    args = promote(i, j, args...) # all args are now particles
    # logic
end

unction Base.getindex(a::AbstractArray{P}, i::Int,, j::Int, k::Particles{<:Integer}, args::Union{AbstractParticles{<:Integer}, Integer}...)  
...

I.e., a triangle of methods that specify 0,1,2,3,... leading integer indices, then one particle index and then a varying number of union arguments, and immediately promotes all of them to particles so that they can be treated like they were all particles.

There are multiple packages that implement "indexing vectors", like StaticArrays
https://juliaarrays.github.io/StaticArrays.jl/latest/pages/api/#Indexing-1
OffsetArrays
https://juliaarrays.github.io/OffsetArrays.jl/dev/internals/#The-axes-of-OffsetArrays
and probably similar packages like NamedArrays etc. Maybe one of these packages has a nice solution to this problem?

[cscherrer]

Thanks @baggepinnen , the "triangle of methods" idea is pretty clever, and seems like it might work