Overload LinearAlgebra.tr (#210)

Author: dkarrasch

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearMaps"
 uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
-version = "3.10.2"
+version = "3.11.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

docs/src/history.md

@@ -1,5 +1,12 @@
 # Version history
 
+## What's new in v3.11
+
+* The `tr` function from `LinearAlgebra.jl` is now overloaded both for generic `LinearMap`
+  types and specialized for most provided `LinearMap` types. In the generic case, this is
+  computationally as expensive as computing the whole matrix representation, though the
+  latter is, of course, not stored.
+
 ## What's new in v3.10
 
 * A new `MulStyle` trait called `TwoArg` has been added. It should be used for `LinearMap`s

docs/src/types.md

@@ -179,6 +179,10 @@ as in the usual matrix case: `transpose(A) * x` and `mul!(y, A', x)`, for instan
   a linear map for which you only have a function definition (e.g. to be able
   to use its `transpose` or `adjoint`).
 
+!!! note
+    In Julia versions v1.9 and higher, conversion to sparse matrices requires loading
+    `SparseArrays.jl` by the user in advance.
+
 ### Slicing methods
 
 Complete slicing, i.e., `A[:,j]`, `A[:,J]`, `A[i,:]`, `A[I,:]` and `A[:,:]` for `i`, `j`
@@ -188,3 +192,12 @@ slicing) to standard unit vectors of appropriate length. By complete slicing we
 two-dimensional Cartesian indexing where at least one of the "indices" is a colon. This is
 facilitated by overloads of `Base.getindex`. Partial slicing à la `A[I,J]` and scalar or
 linear indexing are _not_ supported.
+
+### Sum, product, mean and trace
+
+Natural function overloads for `Base.sum`, `Base.prod`, `Statistics.mean` and `LinearAlgebra.tr`
+exist.
+
+!!! note
+    In Julia versions v1.9 and higher, creating the mean linear operator requires loading
+    `Statistics.jl` by the user in advance.

src/LinearMaps.jl

@@ -5,7 +5,7 @@ export ⊗, squarekron, kronsum, ⊕, sumkronsum, khatrirao, facesplitting
 
 using LinearAlgebra
 using LinearAlgebra: AbstractQ
-import LinearAlgebra: mul!
+import LinearAlgebra: mul!, tr
 
 using Base: require_one_based_indexing
 
@@ -348,6 +348,7 @@ include("conversion.jl") # conversion of linear maps to matrices
 include("show.jl") # show methods for LinearMap objects
 include("getindex.jl") # getindex functionality
 include("inversemap.jl")
+include("trace.jl")
 
 """
     LinearMap(A::LinearMap; kwargs...)::WrappedMap

src/kronecker.jl

@@ -108,7 +108,7 @@ compared to [`kron`](@ref), but benchmarking intended use cases is highly recomm
 function squarekron(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...)
     maps = (A, B, C, Ds...)
     T = promote_type(map(eltype, maps)...)
-    all(_issquare, maps) || throw(ArgumentError("operators need to be square in Kronecker sums"))
+    all(_issquare, maps) || throw(ArgumentError("operators need to be square in squarekron"))
     ns = map(a -> size(a, 1), maps)
     firstmap = first(maps) ⊗ UniformScalingMap(true, prod(ns[2:end]))
     lastmap  = UniformScalingMap(true, prod(ns[1:end-1])) ⊗ last(maps)
@@ -376,7 +376,7 @@ true
 [^1]: Fernandes, P. and Plateau, B. and Stewart, W. J. ["Efficient Descriptor-Vector Multiplications in Stochastic Automata Networks"](https://doi.org/10.1145/278298.278303), _Journal of the ACM_, 45(3), 381–414, 1998.
 """
 function sumkronsum(A::MapOrMatrix, B::MapOrMatrix)
-    LinearAlgebra.checksquare(A, B)
+    (_issquare(A) && _issquare(B)) || throw(ArgumentError("operators need to be square in Kronecker sums"))
     A ⊗ UniformScalingMap(true, size(B,1)) + UniformScalingMap(true, size(A,1)) ⊗ B
 end
 function sumkronsum(A::MapOrMatrix, B::MapOrMatrix, C::MapOrMatrix, Ds::MapOrMatrix...)

src/trace.jl

@@ -0,0 +1,62 @@
+function tr(A::LinearMap)
+    _issquare(A) || throw(ArgumentError("operator needs to be square in tr"))
+    _tr(A)
+end
+
+function _tr(A::LinearMap{T}) where {T}
+    S = typeof(oneunit(eltype(A)) + oneunit(eltype(A)))
+    ax1, ax2 = axes(A)
+    xi = zeros(eltype(A), ax2)
+    y = similar(xi, T, ax1)
+    o = one(T)
+    z = zero(T)
+    s = zero(S)
+    @inbounds for (i, j) in zip(ax1, ax2)
+        xi[j] = o
+        mul!(y, A, xi)
+        xi[j] = z
+        s += y[i]
+    end
+    return s
+end
+function _tr(A::OOPFunctionMap{T}) where {T}
+    S = typeof(oneunit(eltype(A)) + oneunit(eltype(A)))
+    ax1, ax2 = axes(A)
+    xi = zeros(eltype(A), ax2)
+    o = one(T)
+    z = zero(T)
+    s = zero(S)
+    @inbounds for (i, j) in zip(ax1, ax2)
+        xi[j] = o
+        s += (A * xi)[i]
+        xi[j] = z
+    end
+    return s
+end
+# specialiations
+_tr(A::AbstractVecOrMat) = tr(A)
+_tr(A::WrappedMap) = _tr(A.lmap)
+_tr(A::TransposeMap) = _tr(A.lmap)
+_tr(A::AdjointMap) = conj(_tr(A.lmap))
+_tr(A::UniformScalingMap) = A.M * A.λ
+_tr(A::ScaledMap) = A.λ * _tr(A.lmap)
+function _tr(L::KroneckerMap)
+    if all(_issquare, L.maps)
+        return prod(_tr, L.maps)
+    else
+        return invoke(_tr, Tuple{LinearMap}, L)
+    end
+end
+function _tr(L::OuterProductMap{<:RealOrComplex})
+    a, bt = L.maps
+    return bt.lmap*a.lmap
+end
+function _tr(L::OuterProductMap)
+    a, bt = L.maps
+    mapreduce(*, +, a.lmap, bt.lmap)
+end
+function _tr(L::KroneckerSumMap)
+    A, B = L.maps # A and B are square by construction
+    return _tr(A) * size(B, 1) + _tr(B) * size(A, 1)
+end
+_tr(A::FillMap) = A.size[1] * A.λ

test/runtests.jl

@@ -43,3 +43,5 @@ include("inversemap.jl")
 include("rrules.jl")
 
 include("khatrirao.jl")
+
+include("trace.jl")

test/trace.jl

@@ -0,0 +1,28 @@
+using LinearMaps, LinearAlgebra, Test
+
+@testset "trace" begin
+    for A in (randn(5, 5), randn(ComplexF64, 5, 5))
+        @test tr(LinearMap(A)) == tr(A)
+        @test tr(transpose(LinearMap(A))) == tr(A)
+        @test tr(adjoint(LinearMap(A))) == tr(A')
+    end
+    @test tr(LinearMap(3I, 10)) == 30
+    @test tr(LinearMap{Int}(cumsum, 10)) == 10
+    @test tr(LinearMap{Int}(cumsum, reverse∘cumsum∘reverse, 10)') == 10
+    @test tr(LinearMap{Complex{Int}}(cumsum, reverse∘cumsum∘reverse, 10)') == 10
+    @test tr(LinearMap{Int}(cumsum!, 10)) == 10
+    @test tr(2LinearMap{Int}(cumsum!, 10)) == 20
+    A = randn(3, 5); B = copy(transpose(A))
+    @test tr(A ⊗ B) == tr(kron(A, B))
+    @test tr(A ⊗ B ⊗ A ⊗ B) ≈ tr(kron(A, B, A, B))
+    A = randn(5, 5); B = copy(transpose(A))
+    @test tr(A ⊗ B) ≈ tr(kron(A, B))
+    @test tr(A ⊗ B ⊗ A) ≈ tr(kron(A, B, A))
+    @test tr(A ⊗ B ⊗ A ⊗ B) ≈ tr(kron(A, B, A, B))
+    v = A[:,1]
+    @test tr(v ⊗ v') ≈ norm(v)^2
+    v = [randn(2,2) for _ in 1:3]
+    @test tr(v ⊗ v') ≈ mapreduce(*, +, v, v')
+    @test tr(LinearMap{Int}(cumsum!, 10) ⊕ LinearMap{Int}(cumsum!, 10)) == 200
+    @test tr(FillMap(true, 5, 5)) == 5
+end