Operation across the matching dimension of an array with a vector

General Usage
1

Hi there,

I am trying to apply a bivariate function f across an arbitrary dimensional array A and a vector B. I want to apply this function across the dimension of A that matches length with B. For example if A is 3-dimensional and the matching dimension is the second one, then I want to be able to do

C = copy(A)
for i in eachindex(B)
        @. C[:, i, :] = f(view(A, :, i, :), B[i])
end

For testing purposes, f can be *.

I have the following code:

function dimwise(f, A::AbstractArray, B::AbstractVector)
    m = findfirst(isequal(length(B)), size(A))
    isnothing(m) && error("A and B have no matching dimensions.")
    return dimwise(f, A, B, m)
end

dimwise(f, A::AbstractArray, B::AbstractVector, m::Int) = 
dimwise!(copy(A), f, A, B, m)

function dimwise!(C, f, A::AbstractArray, B::AbstractVector, m::Int)
    for i in eachindex(B)
        @. C[...] = f(view(A, [...]), B[i])
    end
    return C
end

This finds the number of the matching dimension, m. What I need to be able to do is apply the line @. C[...] = f(view(A, [...]), B[i]) by adding the correct sequence of colons : and index i within the brackets [...].

Anyone know of a simple, yet performant way to do this?

EDIT: importantly, I need to ensure that C has identical structure with A as far as size is concerned.

2

You probably should use broadcasting. First you need to shift the dimension of B so that it’s for example [1, n, 1], if it’s the second dimension that matches. Then just write:

C .= f.(A, B)

You just need to figure out a way to reshape(B, ...) correctly.

1 Like
3

Thank you!

defining:

function dimwise!(C, f, A::AbstractArray, B::AbstractVector, m::Int)
    r = ones(Int, length(size(A)))
    r[m] = length(B)
    b = reshape(B, r...)
    C .= f.(A, b)
    return C
end

actually did the trick!!!

Would be nice if I found a way to make r not be a vector but a tuple instead but that is minor!

4

Maybe

function shiftbroadcast!(C, f, A::AbstractArray{T, N}, b::AbstractVector) where {T,N}
   n = length(b)
   ind = findfirst(==(n), size(A))
   shape = ntuple(i -> (i==ind) ? n : 1, N)
   b_ = reshape(b, shape)
   C .= f.(A, b_)
   return C
end

?

1 Like