Implement VectorInterface scale and add functions for InfinitePEPS

Author: pbrehmer

src/PEPSKit.jl

@@ -4,6 +4,7 @@ using LinearAlgebra, Statistics, Base.Threads, Base.Iterators, Printf
 using Compat
 using Accessors: @set, @reset
 using VectorInterface
+import VectorInterface as VI
 using TensorKit, KrylovKit, OptimKit, TensorOperations
 using ChainRulesCore, Zygote
 using LoggingExtras

src/environments/ctmrg_environments.jl

@@ -488,8 +488,6 @@ end
 # big vector. In other words, the associated vector space is not the natural one associated
 # to the original (physical) system, and addition, scaling, etc. are performed element-wise.
 
-import VectorInterface as VI
-
 function VI.scalartype(::Type{CTMRGEnv{C,T}}) where {C,T}
     S₁ = scalartype(C)
     S₂ = scalartype(T)

src/operators/infinitepepo.jl

@@ -177,12 +177,8 @@ end
 
 ## Vector interface
 
-function VectorInterface.scalartype(::Type{NT}) where {NT<:InfinitePEPO}
-    return scalartype(eltype(NT))
-end
-function VectorInterface.zerovector(A::InfinitePEPO)
-    return InfinitePEPO(zerovector(unitcell(A)))
-end
+VI.scalartype(::Type{NT}) where {NT<:InfinitePEPO} = scalartype(eltype(NT))
+VI.zerovector(A::InfinitePEPO) = InfinitePEPO(zerovector(unitcell(A)))
 
 ## (Approximate) equality
 function Base.:(==)(A₁::InfinitePEPO, A₂::InfinitePEPO)

src/states/infinitepartitionfunction.jl

@@ -157,12 +157,8 @@ end
 
 ## Vector interface
 
-function VectorInterface.scalartype(::Type{NT}) where {NT<:InfinitePartitionFunction}
-    return scalartype(eltype(NT))
-end
-function VectorInterface.zerovector(A::InfinitePartitionFunction)
-    return InfinitePartitionFunction(zerovector(unitcell(A)))
-end
+VI.scalartype(::Type{NT}) where {NT<:InfinitePartitionFunction} = scalartype(eltype(NT))
+VI.zerovector(A::InfinitePartitionFunction) = InfinitePartitionFunction(zerovector(unitcell(A)))
 
 ## (Approximate) equality
 function Base.:(==)(A₁::InfinitePartitionFunction, A₂::InfinitePartitionFunction)

src/states/infinitepeps.jl

@@ -157,10 +157,29 @@ end
 
 ## Vector interface
 
-function VectorInterface.scalartype(::Type{NT}) where {NT<:InfinitePEPS}
-    return scalartype(eltype(NT))
+VI.scalartype(::Type{NT}) where {NT<:InfinitePEPS} = scalartype(eltype(NT))
+VI.zerovector(A::InfinitePEPS) = InfinitePEPS(zerovector(unitcell(A)))
+
+function VI.scale(ψ::InfinitePEPS, α::Number)
+    _scale = Base.Fix2(scale, α)
+    return InfinitePEPS(map(_scale, unitcell(ψ)))
+end
+function VI.scale!(ψ::InfinitePEPS, α::Number)
+    _scale! = Base.Fix2(scale!, α)
+    return foreach(_scale!, unitcell(ψ))
+end
+VI.scale!!(ψ::InfinitePEPS, α::Number) = scale!(ψ, α)
+
+function VI.add(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS, α::Number, β::Number)
+    _add(x, y) = add(x, y, α, β)
+    return InfinitePEPS(map(_add, unitcell(ψ₁), unitcell(ψ₂)))
+end
+function VI.add!(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS, α::Number, β::Number)
+    _add!(x, y) = add!(x, y, α, β)
+    foreach(_add!, unitcell(ψ₁), unitcell(ψ₂))
+    return ψ₁
 end
-VectorInterface.zerovector(A::InfinitePEPS) = InfinitePEPS(zerovector(unitcell(A)))
+VI.add!!(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS, α::Number, β::Number) = add!(ψ₁, ψ₂, α, β)
 
 ## Math
 

src/states/infiniteweightpeps.jl

@@ -58,7 +58,7 @@ Base.size(W::SUWeight, i) = size(W.data, i)
 Base.length(W::SUWeight) = length(W.data)
 Base.eltype(W::SUWeight) = eltype(typeof(W))
 Base.eltype(::Type{SUWeight{E}}) where {E} = E
-VectorInterface.scalartype(::Type{T}) where {T<:SUWeight} = scalartype(eltype(T))
+VI.scalartype(::Type{T}) where {T<:SUWeight} = scalartype(eltype(T))
 
 Base.getindex(W::SUWeight, args...) = Base.getindex(W.data, args...)
 Base.setindex!(W::SUWeight, args...) = (Base.setindex!(W.data, args...); W)