small improvements

QuantumKitHub · Dec 11, 2024 · 993577d · 993577d
1 parent 187efb1
commit 993577d
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 14 deletions.
diff --git a/src/operators/infinitepepo.jl b/src/operators/infinitepepo.jl
@@ -109,12 +109,13 @@ end
 Base.size(T::InfinitePEPO) = size(T.A)
 Base.size(T::InfinitePEPO, i) = size(T.A, i)
 Base.length(T::InfinitePEPO) = length(T.A)
-Base.eltype(T::InfinitePEPO) = eltype(T.A[1])
-VectorInterface.scalartype(T::InfinitePEPO) = scalartype(T.A)
+Base.eltype(T::InfinitePEPO) = eltype(typeof(T))
+Base.eltype(::Type{<:InfinitePEPO{T}}) where {T} = T
+VectorInterface.scalartype(::Type{T}) where {T<:InfinitePEPO} = scalartype(eltype(T))
 
 ## Copy
 Base.copy(T::InfinitePEPO) = InfinitePEPO(copy(T.A))
-Base.similar(T::InfinitePEPO, args...) = InfinitePEPO(similar.(T.A, args...))
+Base.similar(T::InfinitePEPO, args...) = InfinitePEPO(similar(T.A, args...))
 Base.repeat(T::InfinitePEPO, counts...) = InfinitePEPO(repeat(T.A, counts...))
 
 Base.getindex(T::InfinitePEPO, args...) = Base.getindex(T.A, args...)

diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
@@ -113,12 +113,13 @@ end
 Base.size(T::InfinitePEPS) = size(T.A)
 Base.size(T::InfinitePEPS, i) = size(T.A, i)
 Base.length(T::InfinitePEPS) = length(T.A)
-Base.eltype(T::InfinitePEPS) = eltype(T.A[1])
-VectorInterface.scalartype(T::InfinitePEPS) = scalartype(T.A)
+Base.eltype(T::InfinitePEPS) = eltype(typeof(T))
+Base.eltype(::Type{<:InfinitePEPS{T}}) where {T} = T
+VectorInterface.scalartype(::Type{T}) where {T<:InfinitePEPS} = scalartype(eltype(T))
 
 ## Copy
 Base.copy(T::InfinitePEPS) = InfinitePEPS(copy(T.A))
-Base.similar(T::InfinitePEPS, args...) = InfinitePEPS(similar.(T.A, args...))
+Base.similar(T::InfinitePEPS, args...) = InfinitePEPS(similar(T.A, args...))
 Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...))
 
 Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...)

diff --git a/src/states/infiniteweightpeps.jl b/src/states/infiniteweightpeps.jl
@@ -36,18 +36,17 @@ end
 Base.size(W::SUWeight) = size(W.data)
 Base.size(W::SUWeight, i) = size(W.data, i)
 Base.length(W::SUWeight) = length(W.data)
-Base.eltype(W::SUWeight) = eltype(W.data[1])
+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))
 
 Base.getindex(W::SUWeight, args...) = Base.getindex(W.data, args...)
 Base.setindex!(W::SUWeight, args...) = (Base.setindex!(W.data, args...); W)
 Base.axes(W::SUWeight, args...) = axes(W.data, args...)
 
 function compare_weights(wts1::SUWeight, wts2::SUWeight)
     @assert size(wts1) == size(wts2)
-    wtdiff = sum(
-        _singular_value_distance((wt1, wt2)) for (wt1, wt2) in zip(wts1.data, wts2.data)
-    )
-    return wtdiff / length(wts1)
+    return sum(_singular_value_distance, zip(wts1.data, wts2.data)) / length(wts1)
 end
 
 """

diff --git a/src/utility/util.jl b/src/utility/util.jl
@@ -20,7 +20,7 @@ function _elementwise_mult(a::AbstractTensorMap, b::AbstractTensorMap)
     return dst
 end
 
-_safe_pow(a, pow, tol) = (pow < 0 && abs(a) < tol) ? zero(a) : a .^ pow
+_safe_pow(a, pow, tol) = (pow < 0 && abs(a) < tol) ? zero(a) : a^pow
 
 """
     sdiag_pow(S::AbstractTensorMap, pow::Real; tol::Real=eps(scalartype(S))^(3 / 4))
@@ -60,9 +60,9 @@ function ChainRulesCore.rrule(
 )
     tol *= norm(S, Inf)
     spow = sdiag_pow(S, pow; tol)
-    spow_minus1_conj = sdiag_pow(S', pow - 1; tol)
+    spow_minus1_conj = scale!(sdiag_pow(S', pow - 1; tol), pow)
     function sdiag_pow_pullback(c̄)
-        return (ChainRulesCore.NoTangent(), pow * _elementwise_mult(c̄, spow_minus1_conj))
+        return (ChainRulesCore.NoTangent(), _elementwise_mult(c̄, spow_minus1_conj))
     end
     return spow, sdiag_pow_pullback
 end