Add TensorKit.storagetype for more PEPSKit types

src/environments/bp_environments.jl

@@ -27,6 +27,7 @@ struct BPEnv{T}
     "4 x rows x cols array of message tensors, where the first dimension specifies the spatial direction"
     messages::Array{T, 3}
 end
+TensorKit.storagetype(::Type{BPEnv{T}}) where {T} = storagetype(T)
 
 """
 Construct a message tensor on a certain bond of a network,

src/environments/ctmrg_environments.jl

@@ -235,6 +235,11 @@ CTMRGEnv(env::CTMRGEnv) = CTMRGEnv(env.corners, env.edges)
 
 @non_differentiable CTMRGEnv(state::Union{InfinitePartitionFunction, InfinitePEPS}, args...)
 
+function TensorKit.storagetype(::Type{CTMRGEnv{C, E}}) where {C, E}
+    T = promote_type(scalartype(C), scalartype(E))
+    return TensorKit.promote_storagetype(T, C, E)
+end
+
 # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation
 function ChainRulesCore.rrule(
         ::Type{CTMRGEnv}, corners::Array{C, 3}, edges::Array{T, 3}

src/environments/suweight.jl

@@ -139,6 +139,7 @@ TensorKit.spacetype(w::SUWeight) = spacetype(typeof(w))
 TensorKit.spacetype(::Type{T}) where {E, T <: SUWeight{E}} = spacetype(E)
 TensorKit.sectortype(w::SUWeight) = sectortype(typeof(w))
 TensorKit.sectortype(::Type{<:SUWeight{T}}) where {T} = sectortype(spacetype(T))
+TensorKit.storagetype(::Type{SUWeight{T}}) where {T} = storagetype(T)
 
 ## Bipartite check
 function _is_bipartite(wts::SUWeight)

src/networks/infinitesquarenetwork.jl

@@ -56,9 +56,7 @@ end
 ## Spaces
 
 TensorKit.spacetype(::Type{T}) where {T <: InfiniteSquareNetwork} = spacetype(eltype(T))
-function TensorKit.storagetype(::Type{T}) where {T <: InfiniteSquareNetwork}
-    return storagetype(eltype(T))
-end
+TensorKit.storagetype(::Type{InfiniteSquareNetwork{O}}) where {O} = storagetype(O)
 virtualspace(n::InfiniteSquareNetwork, r::Int, c::Int, dir) = virtualspace(n[r, c], dir)
 
 ## Vector interface

src/operators/localoperator.jl

@@ -87,6 +87,7 @@ function add_term!(
     return operator
 end
 
+TensorKit.storagetype(lo::LocalOperator{T, S}) where {T, S} = TensorKit.promote_storagetype(values(lo.terms)...)
 
 """
     checklattice(Bool, args...)

src/operators/transfermatrix.jl

@@ -115,6 +115,9 @@ function MultilineTransferPEPO(T::InfinitePEPS, O::InfinitePEPO, dir)
     return MPSKit.Multiline(map(cr -> InfiniteTransferPEPO(T, O, dir, cr), 1:rowsize))
 end
 
+TensorKit.storagetype(::Type{InfiniteTransferPEPS{T}}) where {T} = storagetype(T)
+TensorKit.storagetype(::Type{InfiniteTransferPEPO{H, T, O}}) where {H, T, O} = TensorKit.promote_storagetype(T, O)
+
 #
 # Common interface
 #

test/testsuite/boundarymps/vumps.jl

@@ -16,6 +16,7 @@ function boundary_mps_one_one_peps(AT)
         Vpeps = ComplexSpace(2)
         psi = adapt(AT, InfinitePEPS(Vpeps, Vpeps))
         T = PEPSKit.InfiniteTransferPEPS(psi, 1, 1)
+        @test TensorKit.storagetype(T) <: AT
         foreach(V -> (@test V == Vpeps ⊗ Vpeps'), physicalspace(T))
         mps = initialize_mps(T, [ComplexSpace(20)])
 
@@ -117,6 +118,7 @@ function boundary_mps_pepo_runthrough(AT)
         O = ising_pepo(1)
         psi = adapt(AT, PEPSKit.initializePEPS(O, Vpeps))
         T = InfiniteTransferPEPO(psi, O, 1, 1)
+        @test TensorKit.storagetype(T) <: AT
         foreach(V -> (@test V == Vpeps ⊗ Vpepo ⊗ Vpeps'), physicalspace(T))
 
         mps = initialize_mps(rand, scalartype(T), T, [ComplexSpace(10)])

test/types/localoperator.jl

@@ -17,6 +17,7 @@ if !is_buildkite
 
         terms = ((CartesianIndex(1, 1), CartesianIndex(1, 2)) => t,)
         op = LocalOperator(physical_spaces, terms...)
+        @test TensorKit.storagetype(op) == TensorKit.storagetype(t)
 
         @test op isa LocalOperator
         @test length(op.terms) == 1