Elegant broadcasting over Matrix{Any} with different functions

How do I apply function f over the diagonal and g over off-diagonal of a matrix, and more generally over indices I and J? Here I and J are subsets of product of the row and column index sets.

matrix = rand(3,3)
result = [i==j ? f(m[i,j]) : g(m[i,j]) for i in 1:size(matrix,1), j in 1:size(matrix,2)]

I believe that is the cleanest way to express it.

Maybe this one is more elegant (though probably slower)?

m = rand(3,3)
indices = CartesianIndices(m)
result = [I in diag(indices) ? f(m[I]) : g(m[I]) for I in indices]

Here’s another option using broadcasting:

m = rand(3,3)
I = diagind(m)
J = setdiff(LinearIndices(m), I)
result = similar(m)
result[I] .= f.(m[I])
result[J] .= g.(m[J])

and with the InvertedIndices.jl package you can replace setdiff(LinearIndices(m), I) with Not(I).

(Edit: @Seif_Shebl made me realize I had inadvertently removed the broadcasting assignment.)

2 Likes

Yes, you can use broadcasting too.

m = rand(3,3)
I = diagind(m)
r = g.(m)
r[I] .= f.(m[I])
2 Likes