a bit more functionality and tests

Author: Jutho

src/tensors/diagtensor.jl

@@ -198,6 +198,33 @@ function Base.zero(d::DiagonalTensorMap)
     return DiagonalTensorMap(zero.(d.data), d.domain)
 end
 
+function compose(d1::DiagonalTensorMap, d2::DiagonalTensorMap)
+    d1.domain == d2.domain || throw(SpaceMismatch())
+    return DiagonalTensorMap(d1.data .* d2.data, d1.domain)
+end
+Base.inv(d::DiagonalTensorMap) = DiagonalTensorMap(inv.(d.data), d.domain)
+function Base.:\(d1::DiagonalTensorMap, d2::DiagonalTensorMap)
+    d1.domain == d2.domain || throw(SpaceMismatch())
+    return DiagonalTensorMap(d1.data .\ d2.data, d1.domain)
+end
+function Base.:/(d1::DiagonalTensorMap, d2::DiagonalTensorMap)
+    d1.domain == d2.domain || throw(SpaceMismatch())
+    return DiagonalTensorMap(d1.data ./ d2.data, d1.domain)
+end
+function LinearAlgebra.pinv(d::DiagonalTensorMap; kwargs...)
+    T = eltype(d.data)
+    atol = get(kwargs, :atol, zero(real(T)))
+    if iszero(atol)
+        rtol = get(kwargs, :rtol, zero(real(T)))
+    else
+        rtol = sqrt(eps(real(float(oneunit(T))))) * length(d.data)
+    end
+    pdata = let tol = max(atol, rtol * maximum(abs, d.data))
+        map(x -> abs(x) < tol ? zero(x) : pinv(x), d.data)
+    end
+    return DiagonalTensorMap(pdata, d.domain)
+end
+
 function eig!(d::DiagonalTensorMap)
     return d, one(d)
 end

src/tensors/linalg.jl

@@ -345,6 +345,7 @@ end
 function LinearAlgebra.pinv(t::AbstractTensorMap; kwargs...)
     T = float(scalartype(t))
     tpinv = similar(t, T, domain(t) ← codomain(t))
+    # TODO: fix so that `rtol` used total tensor norm instead of per block
     for (c, b) in blocks(t)
         copy!(block(tpinv, c), pinv(b; kwargs...))
     end

test/diagtensors.jl

@@ -103,4 +103,23 @@ diagspacelist = ((ℂ^4)', ℂ[Z2Irrep](0 => 2, 1 => 3),
             @test convert(TensorMap, t5) == permute(convert(TensorMap, t), (((), (2, 1))))
         end
     end
+    @timedtestset "Trace, Multiplication and inverse" begin
+        t1 = DiagonalTensorMap(rand(Float64, reduceddim(V)), V)
+        t2 = DiagonalTensorMap(rand(ComplexF64, reduceddim(V)), V)
+        @test tr(TensorMap(t1)) == @constinferred tr(t1)
+        @test tr(TensorMap(t2)) == @constinferred tr(t2)
+        @test TensorMap(@constinferred t1 * t2) ≈ TensorMap(t1) * TensorMap(t2)
+        @test TensorMap(@constinferred t1 \ t2) ≈ TensorMap(t1) \ TensorMap(t2)
+        @test TensorMap(@constinferred t1 / t2) ≈ TensorMap(t1) / TensorMap(t2)
+        @test TensorMap(@constinferred inv(t1)) ≈ inv(TensorMap(t1))
+        @test TensorMap(@constinferred pinv(t1)) ≈ pinv(TensorMap(t1))
+        @test all(Base.Fix2(isa, DiagonalTensorMap),
+                  (t1 * t2, t1 \ t2, t1 / t2, inv(t1), pinv(t1)))
+
+        u = randn(Float64, V * V' * V, V)
+        @test u * t1 ≈ u * TensorMap(t1)
+        @test u / t1 ≈ u / TensorMap(t1)
+        @test t1 * u' ≈ TensorMap(t1) * u'
+        @test t1 \ u' ≈ TensorMap(t1) \ u'
+    end
 end