Julia `mapslices` when output dimension exceeds input dimension

Viewed 189

I'm trying to apply a function f(x::Array{Float64, 1}) -> Array{Float64, 2} to each row in an m x n array in Julia (I'm using v1.1, if it's relevant).

Now, I was hoping that I could simply apply mapslices as follows:

# toy example of f
f = (x -> randn(length(x), length(x)))
A = randn(100, 50)

# intent: apply f to every row in A and collect the result into a 100 x 50 x 50 matrix.
result = mapslices(f, A, dims=2)

Unfortunately, mapslices fails with the following message:

ERROR: DimensionMismatch("tried to assign 2 elements to 1 destinations")

Is there a mapslices analogue for this situation? I know I can do this with a for loop but I was hoping for something easier.

Update: apparently one way to do this is to embed A into a 3-D array:

result = mapslices(f, reshape(A, (size(A)..., 1)), dims=[2, 3])
2 Answers

I tried only with Julia v1.6, but I think you can very literally "apply f to each row" with eachrow (which appeared in Julia v1.1):

julia> f(x) = randn(length(x), length(x))
f (generic function with 1 method)

julia> A = rand(3,2)
3×2 Matrix{Float64}:
 0.47239    0.179252
 0.542389   0.25828
 0.0513623  0.630193

julia> f.(eachrow(A))
3-element Vector{Matrix{Float64}}:
 [-2.183245554875081 -0.16762649791435957; 0.9124553173227186 0.9148972946316921]
 [-0.7322553194397725 -0.5844536492551982; 0.37738478201981623 -0.7056092457600269]
 [1.0091890849396576 1.6451194487283958; 0.674221636656597 1.0509408618443663]

mapslices isn't very flexible. Here's a variant of what you ask, done without it:

julia> f(x::AbstractVector) = x .* x';  # vector -> matrix

julia> A = randn(5, 7);

julia> f(A[1,:]) |> size
(7, 7)

julia> reduce(hcat, map(f, eachrow(A))) |> size
(7, 35)

julia> B = reshape(reduce(hcat, map(f, eachrow(A))), (7,7,5));

julia> B[:,:,3] ≈ f(A[3,:])
true

Here reduce(hcat, ...) is combining matrices from f along their 2nd direction; we need reshape to separate again the first index of A which has become the last index of B -- which of the 5 rows made this slice.

There are also many package solutions:

julia> using JuliennedArrays

julia> Slices(A,2) |> size  # like eachrow
(5,)

julia> C = Align(map(f, Slices(A,2)), 1,2);  # size (7, 7, 5)

julia> C ≈ B
true

julia> D = Align(map(f, Slices(A,2)), 2,3);  # size (5, 7, 7)

julia> D[3,:,:] ≈ f(A[3,:])
true
julia> using TensorCast

julia> @cast C2[i,j,k] := f(A[k,:])[i,j];  # := makes a new array

julia> C2 ≈ B
true

julia> @cast D2[k,i,j] := f(A[k,:])[i,j];

julia> D2 ≈ D
true
Related