How to specify the i-th axis of array in julia

Viewed 107

How can I simplify the code using for loop?

eta and T are 6-order tensor arrays and size(eta) == size(T) .

using  InvertedIndices
one_to_N = [1:6;]
eta[:,1,1,1,1,1] = vec(sum(T, dims=one_to_N[Not(1)]))
eta[1,:,1,1,1,1] = vec(sum(T, dims=one_to_N[Not(2)]))
eta[1,1,:,1,1,1] = vec(sum(T, dims=one_to_N[Not(3)]))
eta[1,1,1,:,1,1] = vec(sum(T, dims=one_to_N[Not(4)]))
eta[1,1,1,1,:,1] = vec(sum(T, dims=one_to_N[Not(5)]))
eta[1,1,1,1,1,:] = vec(sum(T, dims=one_to_N[Not(6)]))
2 Answers

Instead of creating the sums and then copying them into eta, you can also just write directly using sum!. Making views with index 1:1 rather than 1 means that these dimensions are not dropped, and hence sum! can deduce which dimensions to sum over:

data = rand(Int8, 4,7,2);  # originally T, with ndims(T)==6, reduced for clarity
eta = ones(4,7,2);

for dim in 1:ndims(data)
    ind = ntuple(d -> d==dim ? (:) : (1:1), ndims(data))
    sum!(view(eta, ind...), data)
end

This gives:

julia> eta
4×7×2 Array{Float64, 3}:
[:, :, 1] =
   -1.0  -29.0  -581.0  118.0  -106.0  189.0  198.0
 -100.0    1.0     1.0    1.0     1.0    1.0    1.0
  468.0    1.0     1.0    1.0     1.0    1.0    1.0
 -214.0    1.0     1.0    1.0     1.0    1.0    1.0

[:, :, 2] =
 -130.0  1.0  1.0  1.0  1.0  1.0  1.0
    1.0  1.0  1.0  1.0  1.0  1.0  1.0
    1.0  1.0  1.0  1.0  1.0  1.0  1.0
    1.0  1.0  1.0  1.0  1.0  1.0  1.0

julia> dim = 2;

julia> ind = ntuple(d -> d==dim ? (:) : (1:1), ndims(data))
(1:1, Colon(), 1:1)

julia> view(eta, ind...)
1×7×1 view(::Array{Float64, 3}, 1:1, :, 1:1) with eltype Float64:
[:, :, 1] =
 -1.0  -29.0  -581.0  118.0  -106.0  189.0  198.0

julia> sum(data, dims=(1,3))
1×7×1 Array{Int64, 3}:
[:, :, 1] =
 80  -29  -581  118  -106  189  198

Notice that eta[1] is overwritten several times here, so the result depends on the order of the operations, not just on the data array.

There are many ways to do it. Here are some that quickly come to mind:

You can do it like this:

for i in 1:6
    eta[ifelse.(1:6 .== i, :, 1)...] = vec(sum(T, dims=one_to_N[Not(i)]))
end

or like this:

for i in 1:6
    eta[ntuple(j -> j == i ? (:) : 1, 6)...] = vec(sum(T, dims=one_to_N[Not(i)]))
end

or like this:

for i in 1:6
    eta[(j == i ? (:) : 1 for j in 1:6)...] = vec(sum(T, dims=one_to_N[Not(i)]))
end
Related