Couple more fixes and a CUDA+MPS example

kshyatt · kshyatt · commit 1798f870c708 · 2026-02-10T05:32:06.000-05:00
diff --git a/examples/cu_classic2d/1.hard-hexagon/main.jl b/examples/cu_classic2d/1.hard-hexagon/main.jl
@@ -0,0 +1,96 @@
+"""
+# [The Hard Hexagon model](@id demo_hardhexagon)
+
+![logo](hexagon.svg)
+
+Tensor networks are a natural way to do statistical mechanics on a lattice.
+As an example of this we will extract the central charge of the hard hexagon model.
+This model is known to have central charge 0.8, and has very peculiar non-local (anyonic) symmetries.
+Because TensorKit supports anyonic symmetries, so does MPSKit.
+To follow the tutorial you need the following packages.
+"""
+
+using MatrixAlgebraKit, MPSKit, MPSKitModels, TensorKit, Plots, Polynomials, CUDA, cuTENSOR
+
+#src # for reproducibility:
+#src using Random
+#src Random.seed!(123)
+
+"""
+The [hard hexagon model](https://en.wikipedia.org/wiki/Hard_hexagon_model) is a 2-dimensional lattice model of a gas, where particles are allowed to be on the vertices of a triangular lattice, but no two particles may be adjacent.
+This can be encoded in a transfer matrix with a local MPO tensor using anyonic symmetries, and the resulting MPO has been implemented in MPSKitModels.
+
+In order to use these anyonic symmetries, we need to generalise the notion of the bond dimension and define how it interacts with the symmetry. Thus, we implement away of converting integers to symmetric spaces of the given dimension, which provides a crude guess for how the final MPS would distribute its Schmidt spectrum.
+"""
+mpo = hard_hexagon(CuMatrix{ComplexF64, CUDA.DeviceMemory})
+P = physicalspace(mpo, 1)
+function virtual_space(D::Integer)
+    _D = round(Int, D / sum(dim, values(FibonacciAnyon)))
+    return Vect[FibonacciAnyon](sector => _D for sector in (:I, :τ))
+end
+
+@assert isapprox(dim(virtual_space(100)), 100; atol = 3)
+
+"""
+## The leading boundary
+
+One way to study statistical mechanics in infinite systems with tensor networks is by approximating the dominant eigenvector of the transfer matrix by an MPS. 
+This dominant eigenvector contains a lot of hidden information.
+For example, the free energy can be extracted by computing the expectation value of the mpo.
+Additionally, we can compute the entanglement entropy as well as the correlation length of the state:
+"""
+
+D = 10
+V = virtual_space(D)
+ψ₀ = InfiniteMPS(CuVector{ComplexF64, CUDA.DeviceMemory}, [P], [V]; alg_orth = CUSOLVER_HouseholderQR(; positive=true))
+ψ, envs, = leading_boundary(
+    ψ₀, mpo,
+    VUMPS(; verbosity = 0, alg_eigsolve = MPSKit.Defaults.alg_eigsolve(; ishermitian = false))
+) # use non-hermitian eigensolver
+F = real(expectation_value(ψ, mpo))
+S = real(first(entropy(ψ)))
+ξ = correlation_length(ψ)
+println("F = $F\tS = $S\tξ = $ξ")
+
+"""
+## The scaling hypothesis
+
+The dominant eigenvector is of course only an approximation. The finite bond dimension enforces a finite correlation length, which effectively introduces a length scale in the system. This can be exploited to formulate a scaling hypothesis [pollmann2009](@cite), which in turn allows to extract the central charge.
+
+First we need to know the entropy and correlation length at a bunch of different bond dimensions. Our approach will be to re-use the previous approximated dominant eigenvector, and then expanding its bond dimension and re-running VUMPS.
+According to the scaling hypothesis we should have ``S \propto \frac{c}{6} log(ξ)``. Therefore we should find ``c`` using
+"""
+
+function scaling_simulations(
+        ψ₀, mpo, Ds; verbosity = 0, tol = 1.0e-6,
+        alg_eigsolve = MPSKit.Defaults.alg_eigsolve(; ishermitian = false)
+    )
+    entropies = similar(Ds, Float64)
+    correlations = similar(Ds, Float64)
+    alg = VUMPS(; verbosity, tol, alg_eigsolve)
+
+    ψ, envs, = leading_boundary(ψ₀, mpo, alg)
+    entropies[1] = real(entropy(ψ)[1])
+    correlations[1] = correlation_length(ψ)
+
+    for (i, d) in enumerate(diff(Ds))
+        ψ, envs = changebonds(ψ, mpo, OptimalExpand(; trscheme = truncrank(d)), envs)
+        ψ, envs, = leading_boundary(ψ, mpo, alg, envs)
+        entropies[i + 1] = real(entropy(ψ)[1])
+        correlations[i + 1] = correlation_length(ψ)
+    end
+    return entropies, correlations
+end
+
+bond_dimensions = 10:5:25
+ψ₀ = InfiniteMPS(CuVector{ComplexF64, CUDA.DeviceMemory}, [P], [virtual_space(bond_dimensions[1])])
+Ss, ξs = scaling_simulations(ψ₀, mpo, bond_dimensions)
+
+f = fit(log.(ξs), 6 * Ss, 1)
+c = f.coeffs[2]
+
+#+
+
+p = plot(; xlabel = "logarithmic correlation length", ylabel = "entanglement entropy")
+p = plot(log.(ξs), Ss; seriestype = :scatter, label = nothing)
+plot!(p, ξ -> f(ξ) / 6; label = "fit")
diff --git a/src/environments/abstract_envs.jl b/src/environments/abstract_envs.jl
@@ -15,52 +15,64 @@ Base.unlock(envs::AbstractMPSEnvironments) = unlock(envs.lock);
 # ------------------
 function allocate_GL(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int)
     T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
-    TA = similarstoragetype(storagetype(mpo), T)
+    TO = TensorKit.similarstoragetype(storagetype(mpo), T)
+    TB = TensorKit.similarstoragetype(storagetype(bra), T)
+    TK = TensorKit.similarstoragetype(storagetype(ket), T)
+    TA = reduce(TensorKit.promote_storagetype, (TB, TO, TK))
     V = left_virtualspace(bra, i) ⊗ left_virtualspace(mpo, i)' ←
         left_virtualspace(ket, i)
     if V isa BlockTensorKit.TensorMapSumSpace
         TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), TA)
     else
-        TT = TensorKit.TensorMapWithStorage{T, TA}
+        TT = TensorKit.tensormaptype(spacetype(bra), numout(V), numin(V), TA)
     end
     return TT(undef, V)
 end
 
 function allocate_GR(bra::AbstractMPS, mpo::AbstractMPO, ket::AbstractMPS, i::Int)
     T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
-    TA = similarstoragetype(storagetype(mpo), T)
+    TO = TensorKit.similarstoragetype(storagetype(mpo), T)
+    TB = TensorKit.similarstoragetype(storagetype(bra), T)
+    TK = TensorKit.similarstoragetype(storagetype(ket), T)
+    TA = reduce(TensorKit.promote_storagetype, (TB, TO, TK))
     V = right_virtualspace(ket, i) ⊗ right_virtualspace(mpo, i) ←
         right_virtualspace(bra, i)
     if V isa BlockTensorKit.TensorMapSumSpace
         TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), TA)
     else
-        TT = TensorKit.TensorMapWithStorage{T, TA}
+        TT = TensorKit.tensormaptype(spacetype(bra), numout(V), numin(V), TA)
     end
     return TT(undef, V)
 end
 
 function allocate_GBL(bra::QP, mpo::AbstractMPO, ket::QP, i::Int)
     T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
-    TA = similarstoragetype(storagetype(mpo), T)
+    TO = TensorKit.similarstoragetype(storagetype(mpo), T)
+    TB = TensorKit.similarstoragetype(storagetype(bra), T)
+    TK = TensorKit.similarstoragetype(storagetype(ket), T)
+    TA = reduce(TensorKit.promote_storagetype, (TB, TO, TK))
     V = left_virtualspace(bra.left_gs, i) ⊗ left_virtualspace(mpo, i)' ←
         auxiliaryspace(ket)' ⊗ left_virtualspace(ket.right_gs, i)
     if V isa BlockTensorKit.TensorMapSumSpace
         TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), TA)
     else
-        TT = TensorKit.TensorMapWithStorage{T, TA}
+        TT = TensorKit.tensormaptype(spacetype(bra), numout(V), numin(V), TA)
     end
     return TT(undef, V)
 end
 
 function allocate_GBR(bra::QP, mpo::AbstractMPO, ket::QP, i::Int)
     T = Base.promote_type(scalartype(bra), scalartype(mpo), scalartype(ket))
-    TA = similarstoragetype(storagetype(mpo), T)
+    TO = TensorKit.similarstoragetype(storagetype(mpo), T)
+    TB = TensorKit.similarstoragetype(storagetype(bra), T)
+    TK = TensorKit.similarstoragetype(storagetype(ket), T)
+    TA = reduce(TensorKit.promote_storagetype, (TB, TO, TK))
     V = right_virtualspace(ket.left_gs, i) ⊗ right_virtualspace(mpo, i) ←
         auxiliaryspace(ket)' ⊗ right_virtualspace(bra.right_gs, i)
     if V isa BlockTensorKit.TensorMapSumSpace
         TT = blocktensormaptype(spacetype(bra), numout(V), numin(V), TA)
     else
-        TT = TensorKit.TensorMapWithStorage{T, TA}
+        TT = TensorKit.tensormaptype(spacetype(bra), numout(V), numin(V), TA)
     end
     return TT(undef, V)
 end
diff --git a/src/operators/jordanmpotensor.jl b/src/operators/jordanmpotensor.jl
@@ -65,6 +65,9 @@ const JordanMPOTensorMap{T, S, A <: DenseVector{T}} = JordanMPOTensor{
     TensorMap{T, S, 1, 2, A},
     TensorMap{T, S, 1, 1, A},
 }
+# TODO FIX ME
+TensorKit.storagetype(jmpo::JordanMPOTensorMap{E, S, TA}) where {E, S, TA} = TA
+TensorKit.storagetype(::Type{JordanMPOTensorMap{E, S, TA}}) where {E, S, TA} = TA
 
 function JordanMPOTensor{E, S}(::UndefInitializer, V::TensorMapSumSpace{S}) where {E, S}
     return jordanmpotensortype(S, E)(undef, V)
@@ -121,12 +124,13 @@ function JordanMPOTensor(W::SparseBlockTensorMap{TT, E, S, 2, 2}) where {TT, E,
     )
 end
 
-function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E <: Number}
-    TA = Union{tensormaptype(S, 2, 2, E), BraidingTensor{E, S}}
-    TB = tensormaptype(S, 2, 1, E)
-    TC = tensormaptype(S, 1, 2, E)
-    TD = tensormaptype(S, 1, 1, E)
-    return JordanMPOTensor{E, S, TA, TB, TC, TD}
+function jordanmpotensortype(::Type{S}, ::Type{TorA}) where {S <: VectorSpace, TorA}
+    BT = BraidingTensor{eltype(TorA), S}
+    TA = Union{tensormaptype(S, 2, 2, TorA), BT}
+    TB = tensormaptype(S, 2, 1, TorA)
+    TC = tensormaptype(S, 1, 2, TorA)
+    TD = tensormaptype(S, 1, 1, TorA)
+    return JordanMPOTensor{eltype(TorA), S, TA, TB, TC, TD}
 end
 function jordanmpotensortype(::Type{O}) where {O <: MPOTensor}
     return jordanmpotensortype(spacetype(O), scalartype(O))
diff --git a/src/operators/mpohamiltonian.jl b/src/operators/mpohamiltonian.jl
@@ -48,6 +48,8 @@ end
 const InfiniteMPOHamiltonian{O <: MPOTensor} = MPOHamiltonian{O, PeriodicVector{O}}
 Base.isfinite(::Type{<:InfiniteMPOHamiltonian}) = false
 GeometryStyle(::Type{<:InfiniteMPOHamiltonian}) = InfiniteChainStyle()
+TensorKit.storagetype(impo::InfiniteMPOHamiltonian{O}) where {O} = storagetype(O)
+TensorKit.storagetype(::Type{InfiniteMPOHamiltonian{O}}) where {O} = storagetype(O)
 
 function InfiniteMPOHamiltonian(Ws::AbstractVector{O}) where {O <: MPOTensor}
     for i in eachindex(Ws)
@@ -522,6 +524,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
     T = _find_tensortype(nonzero_opps)
     E = scalartype(T)
     S = spacetype(T)
+    TA = storagetype(T)
 
     # construct the virtual spaces
     MissingS = Union{Missing, S}
@@ -588,7 +591,7 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
     end
 
     # construct the tensor
-    TW = jordanmpotensortype(S, E)
+    TW = jordanmpotensortype(S, TA)
     Os = map(1:length(lattice)) do site
         V = virtualsumspaces[site - 1] * lattice[site] ←
             lattice[site] * virtualsumspaces[site]
diff --git a/src/states/abstractmps.jl b/src/states/abstractmps.jl
@@ -34,8 +34,13 @@ Construct an `MPSTensor` with given physical and virtual spaces.
 """
 function MPSTensor(
         ::UndefInitializer, ::Type{TorA}, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
-    ) where {S <: ElementarySpace, TorA}
-    return TensorKit.TensorMapWithStorage{TorA}(undef, Vₗ ⊗ P ← Vᵣ)
+    ) where {S <: ElementarySpace, TorA <: Number}
+    return TensorKit.TensorMap{TorA}(undef, Vₗ ⊗ P ← Vᵣ)
+end
+function MPSTensor(
+        ::UndefInitializer, ::Type{TorA}, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
+    ) where {S <: ElementarySpace, TorA <: DenseVector}
+    return TensorKit.TensorMapWithStorage{eltype(TorA), TorA}(undef, Vₗ ⊗ P ← Vᵣ)
 end
 function MPSTensor(
         f, ::Type{TorA}, P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
@@ -51,6 +56,12 @@ function MPSTensor(
         throw(ArgumentError("Unsupported initializer function: $f"))
     end
 end
+function MPSTensor(
+        ::Type{TorA},
+        P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
+    ) where {S <: ElementarySpace, TorA}
+    return MPSTensor(rand, TorA, P, Vₗ, Vᵣ)
+end
 # TODO: reinstate function initializers?
 function MPSTensor(
         P::Union{S, CompositeSpace{S}}, Vₗ::S, Vᵣ::S = Vₗ
diff --git a/src/states/infinitemps.jl b/src/states/infinitemps.jl
@@ -102,6 +102,8 @@ struct InfiniteMPS{A <: GenericMPSTensor, B <: MPSBondTensor} <: AbstractMPS
         return new{A, B}(AL, AR, C, AC)
     end
 end
+TensorKit.storagetype(imps::InfiniteMPS{A, B}) where {A, B} = storagetype(A)
+TensorKit.storagetype(::Type{InfiniteMPS{A, B}}) where {A, B} = storagetype(A)
 
 #===========================================================================================
 Constructors
@@ -116,35 +118,42 @@ function InfiniteMPS(
         convert(PeriodicVector{B}, C), convert(PeriodicVector{A}, AC)
     )
 end
-
 function InfiniteMPS(
+        ::Type{TorA},
         pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
         kwargs...
-    ) where {S <: IndexSpace}
-    return InfiniteMPS(MPSTensor.(pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
+    ) where {S <: IndexSpace, TorA}
+    return InfiniteMPS(MPSTensor.(TorA, pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
 end
 function InfiniteMPS(
-        f, elt::Type{<:Number}, pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
+        pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
         kwargs...
     ) where {S <: IndexSpace}
+    return InfiniteMPS(MPSTensor.(ComplexF64, pspaces, circshift(Dspaces, 1), Dspaces); kwargs...)
+end
+function InfiniteMPS(
+        f, ::Type{TorA}, pspaces::AbstractVector{S}, Dspaces::AbstractVector{S};
+        kwargs...
+    ) where {S <: IndexSpace, TorA}
     return InfiniteMPS(
-        MPSTensor.(f, elt, pspaces, circshift(Dspaces, 1), Dspaces);
+        MPSTensor.(f, TorA, pspaces, circshift(Dspaces, 1), Dspaces);
         kwargs...
     )
 end
-InfiniteMPS(d::S, D::S) where {S <: Union{Int, <:IndexSpace}} = InfiniteMPS([d], [D])
+InfiniteMPS(::Type{TorA}, d::S, D::S; kwargs...) where {S <: Union{Int, <:IndexSpace}, TorA} = InfiniteMPS(TorA, [d], [D]; kwargs...)
+InfiniteMPS(d::S, D::S; kwargs...) where {S <: Union{Int, <:IndexSpace}} = InfiniteMPS([d], [D]; kwargs...)
 function InfiniteMPS(
-        f, elt::Type{<:Number}, d::S, D::S
-    ) where {S <: Union{Int, <:IndexSpace}}
-    return InfiniteMPS(f, elt, [d], [D])
+        f, ::Type{TorA}, d::S, D::S; kwargs...
+    ) where {S <: Union{Int, <:IndexSpace}, TorA}
+    return InfiniteMPS(f, TorA, [d], [D]; kwargs...)
 end
-function InfiniteMPS(ds::AbstractVector{Int}, Ds::AbstractVector{Int})
-    return InfiniteMPS(ComplexSpace.(ds), ComplexSpace.(Ds))
+function InfiniteMPS(ds::AbstractVector{Int}, Ds::AbstractVector{Int}; kwargs...)
+    return InfiniteMPS(ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
 end
 function InfiniteMPS(
-        f, elt::Type{<:Number}, ds::AbstractVector{Int}, Ds::AbstractVector{Int}, kwargs...
-    )
-    return InfiniteMPS(f, elt, ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
+        f, ::Type{TorA}, ds::AbstractVector{Int}, Ds::AbstractVector{Int}, kwargs...
+    ) where {TorA}
+    return InfiniteMPS(f, TorA, ComplexSpace.(ds), ComplexSpace.(Ds); kwargs...)
 end
 
 function InfiniteMPS(A::AbstractVector{<:GenericMPSTensor}; kwargs...)
