Add TensorOperations specializations

Author: lkdvos

src/tensors/diagonal.jl

@@ -185,6 +185,34 @@ function VectorInterface.add(ty::DiagonalTensorMap, tx::DiagonalTensorMap,
     return add!(scale!(zerovector(ty, T), ty, β), tx, α) # zerovector instead of similar preserves diagonal structure
 end
 
+# TensorOperations
+# ----------------
+function TO.tensoradd_type(TC, A::DiagonalTensorMap, ::Index2Tuple{1,1}, ::Bool)
+    M = similarstoragetype(A, TC)
+    return DiagonalTensorMap{TC,spacetype(A),M}
+end
+
+function TO.tensorcontract_type(TC, A::DiagonalTensorMap, ::Index2Tuple{1,1}, ::Bool,
+                                B::DiagonalTensorMap, ::Index2Tuple{1,1}, ::Bool,
+                                ::Index2Tuple{1,1})
+    M = similarstoragetype(A, TC)
+    M == similarstoragetype(B, TC) ||
+        throw(ArgumentError("incompatible storage types:\n$(M) ≠ $(similarstoragetype(B, TC))"))
+    spacetype(A) == spacetype(B) || throw(SpaceMismatch("incompatible space types"))
+    return DiagonalTensorMap{TC,spacetype(A),M}
+end
+
+function TO.tensoralloc(::Type{DiagonalTensorMap{T,S,M}},
+                        structure::TensorMapSpace{S,1,1},
+                        istemp::Val,
+                        allocator=TO.DefaultAllocator()) where {T,S,M}
+    domain(structure) == codomain(structure) || throw(ArgumentError("domain ≠ codomain"))
+    V = only(domain(structure))
+    dim = reduceddim(V)
+    data = TO.tensoralloc(M, dim, istemp, allocator)
+    return DiagonalTensorMap{T,S,M}(data, V)
+end
+
 # Linear Algebra and factorizations
 # ---------------------------------
 function one!(d::DiagonalTensorMap)
@@ -198,13 +226,6 @@ function Base.zero(d::DiagonalTensorMap)
     return DiagonalTensorMap(zero.(d.data), d.domain)
 end
 
-function compose_dest(d1::DiagonalTensorMap, d2::DiagonalTensorMap)
-    A = promote_type(storagetype(d1), storagetype(d2))
-    S = spacetype(d1)
-    T = scalartype(A)
-    return DiagonalTensorMap{T,S,A}(undef, d1.domain)
-end
-
 function LinearAlgebra.mul!(dC::DiagonalTensorMap,
                             dA::DiagonalTensorMap,
                             dB::DiagonalTensorMap,

test/diagonal.jl

@@ -130,6 +130,9 @@ diagspacelist = ((ℂ^4)', ℂ[Z2Irrep](0 => 2, 1 => 3),
         if BraidingStyle(I) isa SymmetricBraiding
             @tensor C[a b c; d] := A[a b c; e] * d[e, d]
             @test C ≈ A * d
+            @tensor D[a; b] := d[a, c] * d[c, b]
+            @test D ≈ d * d
+            @test D isa DiagonalTensorMap
         end
         @planar E1[-1 -2 -3; -4 -5] := B[-1 -2 -3; 1 -5] * d[1; -4]
         @planar E2[-1 -2 -3; -4 -5] := B[-1 -2 -3; 1 -5] * t[1; -4]