Search code examples
arraysmultidimensional-arrayjuliabroadcastarray-broadcasting

Julia: mapslices() where the function being mapped has multiple input arguments?


I am trying to use mapslices (perhaps this isn't even the right approach) in a way that allows me to broadcast a function across the 'slices' of a multidimensional array, where the function involves another array that is always the same for each iteration of the broadcast. Can this be done with mapslices or map?

Here's the particular problem:

I have a function F(X,Y) that takes as input 2-Dimensional arrays X and Y. X is size (D, N) and has type = Array{Float32,2}), whereas Y is also of size (D, 1), with type = Array{Float32,2}.

F(X,Y) performs an operation on the two matrices X and Y to yield an array output with size (N, 1).

The broadcasting/mapping problem: I have a larger 3-D array X_all whose 'slices' (e.g. matrices stacked long the third dimension, like X[:,:,1]) are examples of input arrays X with dimensions (D,N), that are supposed to be inputs to F(X,Y). And Y is always constant.

So an example X_all would be of shape (D, N, 5), with Y still being (D,1), and I want the output of the broadcast to be a matrix whose columsn store the results of applying F to each slice of X_all and Y. So the final output shape in this example should be (N,5).

This isn't working with broadcast (the error I get suggests that F(X,Y) is being fed scalar floating point numbers, not arrays), so I'm assuming it needs to be done with mapslices. The only issue is that I can't seem to find examples of mapslices being used with multiple input arguments (e.g. X and Y) -- usually I just see it with one input argument.

Any thoughts?


Solution

  • Your written description of the problem is a little hard for me to follow, but if I understand the problem correctly, you just need to use mapslices with an anonymous function X -> F(X, Y)

    Here's an example:

    F(X, Y) = X'*Y 
    
    F(rand(3, 4), rand(3, 1))
    
    #+RESULTS:
    : 4×1 Array{Float64,2}:
    :  0.2038830017261133
    :  0.16952327765859904
    :  0.25986519462136687
    :  0.18604863369789784
    
    #+BEGIN_SRC jupyter-julia
    let D = 3, N = 4
        x_all = randn(D, N, 5)
        Y = randn(D, 1)
    
        mapslices(X -> F(X, Y), x_all, dims=(1,2))
    end
    
    #+RESULTS:
    4×1×5 Array{Float64,3}:
    [:, :, 1] =
      1.949979259990055
      0.3939816492611507
     -0.8464390217323449
      0.4555254821129745
    
    [:, :, 2] =
      1.6893553574839384
      1.9573396471517739
      1.5358035472026144
     -1.2706780393752082
    
    [:, :, 3] =
      0.999887609119746
      0.32959013947652727
      0.8756560406754215
     -1.0257977117576802
    
    [:, :, 4] =
     -0.2252063989697573
     -0.5968450960471975
     -0.7319338928241739
      0.236534278792568
    
    [:, :, 5] =
     -0.9523913128029378
      0.13537059021324083
      0.9514285829484852
      0.5464547535310538