Expanding array along a dimension

Question

The following Python function expands an array along a dimension

def expand(x, dim,copies):
        trans_cmd = list(range(0,len(x.shape)))
        trans_cmd.insert(dim,len(x.shape))
        new_data = x.repeat(copies).reshape(list(x.shape) + [copies]).transpose(trans_cmd)
        return new_data

>>x = np.array([[1,2,3],[4,5,6]])
>>x.shape
(2, 3)
>>x_new = expand(x,2,4)
>>x_new.shape
(2,3,4)
>>x_new
array([[[1, 1, 1, 1],
        [2, 2, 2, 2],
        [3, 3, 3, 3]],

       [[4, 4, 4, 4],
        [5, 5, 5, 5],
        [6, 6, 6, 6]]])

How can this function be replicated in Julia?

user_176071 · Answer

Instead of doing a repeat -> reshape -> permutedims (like you have in numpy), I'd just do a reshape -> repeat. Accounting for the translation of row-major to column-major, it looks like this:

julia> x = [[1,2,3] [4,5,6]]
3×2 Array{Int64,2}:
 1  4
 2  5
 3  6

julia> repeat(reshape(x, (1, 3, 2)), outer=(4,1,1))
4×3×2 Array{Int64,3}:
[:, :, 1] =
 1  2  3
 1  2  3
 1  2  3
 1  2  3

[:, :, 2] =
 4  5  6
 4  5  6
 4  5  6
 4  5  6

The tricky part about doing this efficiently in Julia is the construction of the tuples (1, 3, 2) and (4, 1, 1). While you could convert them to arrays and use array mutation (like you convert it to a python list), it's far more efficient to keep them as tuples. ntuple is your friend here:

julia> function expand(x, dim, copies)
           sz = size(x)
           rep = ntuple(d->d==dim ? copies : 1, length(sz)+1)
           new_size = ntuple(d->d expand(x, 1, 4)
4×3×2 Array{Int64,3}:
[:, :, 1] =
 1  2  3
 1  2  3
 1  2  3
 1  2  3

[:, :, 2] =
 4  5  6
 4  5  6
 4  5  6
 4  5  6

julia> expand(x, 2, 4)
3×4×2 Array{Int64,3}:
[:, :, 1] =
 1  1  1  1
 2  2  2  2
 3  3  3  3

[:, :, 2] =
 4  4  4  4
 5  5  5  5
 6  6  6  6

julia> expand(x, 3, 4)
3×2×4 Array{Int64,3}:
[:, :, 1] =
 1  4
 2  5
 3  6

[:, :, 2] =
 1  4
 2  5
 3  6

[:, :, 3] =
 1  4
 2  5
 3  6

[:, :, 4] =
 1  4
 2  5
 3  6

Expanding array along a dimension

1 Answers