# Operation across the matching dimension of an array with a vector

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.

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

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!

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