`map` for `Diagonal` matrix is inefficient

[eschnett]

trafficstars

Applying a function to a Diagonal matrix can convert it to a dense matrix without need:

julia> using LinearAlgebra
julia> using StaticArrays
julia> map(SMatrix{1,1}, Diagonal([1, 2]))
2×2 Matrix{SMatrix{1, 1, Int64, 1}}:
 [1;;]  [0;;]
 [0;;]  [2;;]

The result of map is a dense matrix (Matrix) instead of a diagonal matrix (Diagonal).

Although iszero(SMatrix{1,1}(1)), map converts the matrix to a dense matrix.

[mikmoore]

In this particular example you can get away with this, but in general you cannot. For example, map(cos, Diagonal([1, 2])) does not preserve diagonal structure. It would be difficult to make this work in general so I really don't think you're likely to find somebody (else) who will put much effort into improving it.

If you have a specific case where you know this pattern is legal, I'll suggest you use something like Diagonal(map(SMatrix{1,1}, diag(Diagonal([1, 2])))).

[gbaraldi]

For things it knows are fine it does do an efficient mapping like

julia> map(sin, Diagonal([1, 2]))
2×2 Diagonal{Float64, Vector{Float64}}:
 0.841471   ⋅
  ⋅        0.909297

So I'm not sure there is much to be done here

[eschnett]

I thought map performed a run-time check on the return value for a zero input to the mapped function. See e.g. https://github.com/JuliaSparse/SparseArrays.jl/blob/090474bec9de69d2573e816a6187568822f1d85a/src/higherorderfns.jl#L160 for sparse arrays:

function _noshapecheck_map!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N}
    fofzeros = f(_zeros_eltypes(A, Bs...)...)
    fpreszeros = _iszero(fofzeros)
    return fpreszeros ? _map_zeropres!(f, C, A, Bs...) :
                        _map_notzeropres!(f, fofzeros, C, A, Bs...)
end

Diagonal arrays and their friends should do the same.

[mikmoore]

Neat! I didn't notice that it already had that infrastructure in place. So what appears to be the issue here is that it isn't recognizing that f(zero(T)) == zero(f(_::T)) (or do we need ===?) when f=SMatrix{1,1} and T=Int (although the T=Int doesn't seem to be the main issue here).

[eschnett]

Calling iszero on the result of f is more efficient that comparing to zero.

[mikmoore]

iszero cannot catch every case. Consider f(x) = fill(x,1,1,1). If we only check that iszero(f(zero(T))) then we will determine that the Diagonal-result map is legal. But It isn't.

julia> map(x->fill(x,1,1,1),Diagonal(1:2)) # this is the current behavior
2×2 Matrix{Array{Int64, 3}}:
 [1;;;]  [0;;;]
 [0;;;]  [2;;;]

julia> Diagonal(map(x->fill(x,1,1,1),1:2)) # if we only check for `iszero` we might try this
2×2 Diagonal{Array{Int64, 3}, Vector{Array{Int64, 3}}}:
Error showing value of type Diagonal{Array{Int64, 3}, Vector{Array{Int64, 3}}}:
ERROR: MethodError: no method matching zero(::Type{Array{Int64, 3}})

I had to go to Array{T,3} elements because apparently we've already specialized Matrix{T} off-diagonals to infer their dimensions from the corresponding diagonals. Although even those are fraught because I could make a function that does something much different with the size of off-diagonals. For example,

julia> map(x->fill(x,x+1,2x+1),Diagonal(1:2))[1,2] # intended behavior
1×1 Matrix{Int64}:
 0

julia> Diagonal(map(x->fill(x,x+1,2x+1),1:2))[1,2] # if we over-specialize for `iszero` off-diagonals
2×5 Matrix{Int64}:
 0  0  0  0  0
 0  0  0  0  0

The current implementation is correct but a version that specializes when iszero(f(zero(T))) would get this wrong.

[eschnett]

This problem isn't specific to map. One can define types where the result of sin(zero(T)) has the same problem:

struct IntFillMatrix
    val::Int
end

and its values are equivalent to

Base.convert(::Type{Matrix}, ifm::IntFillMatrix) = ifm.val * [1 1; 1 1]

Applying sin to it, we need to convert to floating point. Our type doesn't support that, so we implement

Base.sin(ifm::IntFillMatrix) = sin(Matrix(ifm))

I agree that my example is contrived, and I haven't encountered this problem in practice. However, given that sparse arrays also currently don't handle this, I would treat this as a separate problem, and implement the sparse array optimization to diagonal (and related) matrices as well

[eschnett]

Also, the code e.g. for inv or pinv for diagonal matrices calls zero(T) on many occasions. I think we can deduce that T = Matrix isn't currently handled well.