Slicing and broadcasting multidimensional arrays in Julia : meshgrid example

Question

I recently started learning Julia by coding a simple implementation of Self Organizing Maps. I want the size and dimensions of the map to be specified by the user, which means I can't really use for loops to work on the map arrays because I don't know in advance how many layers of loops I will need. So I absolutely need broadcasting and slicing functions that work on arrays of arbitrary dimensions.

Right now, I need to construct an array of indices of the map. Say my map is defined by an array of size mapsize = (5, 10, 15), I need to construct an array indices of size (3, 5, 10, 15) where indices[:, a, b, c] should return [a, b, c].

I come from a Python/NumPy background, in which the solution is already given by a specific "function", mgrid :

indices = numpy.mgrid[:5, :10, :15]
print indices.shape # gives (3, 5, 10, 15)
print indices[:, 1, 2, 3] gives [1, 2, 3]

I didn't expect Julia to have such a function on the get-go, so I turned to broadcasting. In NumPy, broadcasting is based on a set of rules that I find quite clear and logical. You can use arrays of different dimensions in the same expression as long as the sizes in each dimension match or one of it is 1 :

(5, 10, 15)   broadcasts to  (5, 10, 15) 
   (10,  1)

(5,  1, 15)   also broadcasts to (5, 10, 15) 
(1, 10,  1)

To help with this, you can also use numpy.newaxis or None to easily add new dimensions to your array :

array = numpy.zeros((5, 15))
array[:,None,:] has shape (5, 1, 15)

This helps broadcast arrays easily :

A = numpy.arange(5)
B = numpy.arange(10)
C = numpy.arange(15)
bA, bB, bC = numpy.broadcast_arrays(A[:,None,None], B[None,:,None], C[None,None,:])
bA.shape == bB.shape == bC.shape = (5, 10, 15)

Using this, creating the indices array is rather straightforward :

indices = numpy.array(numpy.broadcast_arrays(A[:,None,None], B[None,:,None], C[None,None,:]))
(indices == numpy.mgrid[:5,:10,:15]).all() returns True

The general case is of course a bit more complicated, but can be worked around using list comprehension and slices :

arrays = [ numpy.arange(i)[tuple([None if m!=n else slice(None) for m in range(len(mapsize))])] for n, i in enumerate(mapsize) ]
indices = numpy.array(numpy.broadcast_arrays(*arrays))

So back to Julia. I tried to apply the same kind of rationale and ended up achieving the equivalent of the arrays list of the code above. This ended up being rather simpler than the NumPy counterpart thanks to the compound expression syntax :

arrays = [ (idx = ones(Int, length(mapsize)); idx[n] = i;reshape([1:i], tuple(idx...))) for (n,i)=enumerate(mapsize) ]

Now I'm stuck here, as I don't really know how to apply the broadcasting to my list of generating arrays here... The broadcast[!] functions ask for a function f to apply, and I don't have any. I tried using a for loop to try forcing the broadcasting:

indices = Array(Int, tuple(unshift!([i for i=mapsize], length(mapsize))...))
for i=1:length(mapsize)
    A[i] = arrays[i]
end

But this gives me an error : ERROR: convert has no method matching convert(::Type{Int64}, ::Array{Int64,3})

Am I doing this the right way? Did I overlook something important? Any help is appreciated.

Answer 1

If you're running julia 0.4, you can do this:

julia> function mgrid(mapsize)
           T = typeof(CartesianIndex(mapsize))
           indices = Array(T, mapsize)
           for I in eachindex(indices)
               indices[I] = I
           end
           indices
       end

It would be even nicer if one could just say

indices = [I for I in CartesianRange(CartesianIndex(mapsize))]

I'll look into that :-).

Answer 2

Broadcasting in Julia has been modelled pretty much on broadcasting in NumPy, so you should hopefully find that it obeys more or less the same simple rules (not sure if the way to pad dimensions when not all inputs have the same number of dimensions is the same though, since Julia arrays are column-major).

A number of useful things like newaxis indexing and broadcast_arrays have not been implemented (yet) however. (I hope they will.) Also note that indexing works a bit differently in Julia compared to NumPy: when you leave off indices for trailing dimensions in NumPy, the remaining indices default to colons. In Julia they could be said to default to ones instead.

I'm not sure if you actually need a meshgrid function, most things that you would want to use it for could be done by using the original entries of your arrays array with broadcasting operations. The major reason that meshgrid is useful in matlab is because it is terrible at broadcasting.

But it is quite straightforward to accomplish what you want to do using the broadcast! function:

# assume mapsize is a vector with the desired shape, e.g. mapsize = [2,3,4]

N = length(mapsize)
# Your line to create arrays below, with an extra initial dimension on each array
arrays = [ (idx = ones(Int, N+1); idx[n+1] = i;reshape([1:i], tuple(idx...))) for (n,i) in enumerate(mapsize) ]

# Create indices and fill it one coordinate at a time
indices = zeros(Int, tuple(N, mapsize...))
for (i,arr) in enumerate(arrays)
    dest = sub(indices, i, [Colon() for j=1:N]...)
    broadcast!(identity, dest, arr)
end

I had to add an initial singleton dimension on the entries of arrays to line up with the axes of indices (newaxis had been useful here...). Then I go through each coordinate, create a subarray (a view) on the relevant part of indices, and fill it. (Indexing will default to returning subarrays in Julia 0.4, but for now we have to use sub explicitly).

The call to broadcast! just evaluates the identity function identity(x)=x on the input arr=arrays[i], broadcasts to the shape of the output. There's no efficiency lost in using the identity function for this; broadcast! generates a specialized function based on the given function, number of arguments, and number of dimensions of the result.