More tests for CUDA + MPS

Project.toml

@@ -12,6 +12,7 @@ HalfIntegers = "f0d1745a-41c9-11e9-1dd9-e5d34d218721"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
 OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
@@ -20,6 +21,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorKitManifolds = "11fa318c-39cb-4a83-b1ed-cdc7ba1e3684"
+TensorKitTensors = "41b62e7d-e9d1-4e23-942c-79a97adf954b"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 
@@ -35,21 +37,29 @@ projects = ["test", "docs"]
 [compat]
 Accessors = "0.1"
 Adapt = "4"
-BlockTensorKit = "0.3.11"
+BlockTensorKit = "0.3"
 Compat = "3.47, 4.10"
 DocStringExtensions = "0.9.3"
 HalfIntegers = "1.6.0"
 KrylovKit = "0.8.3, 0.9.2, 0.10"
 LinearAlgebra = "1.6"
 LoggingExtras = "~1.0"
 MatrixAlgebraKit = "0.6.5"
+MPSKitModels = "0.4"
 OhMyThreads = "0.7, 0.8"
 OptimKit = "0.3.1, 0.4"
 Printf = "1"
 Random = "1"
 RecipesBase = "1.1"
-TensorKit = "0.16.5"
+TensorKit = "0.16, 0.17"
 TensorKitManifolds = "0.7, 0.8"
 TensorOperations = "5.5.1"
 VectorInterface = "0.2, 0.3, 0.4, 0.5"
 julia = "1.10"
+
+[sources]
+MPSKitModels = {url = "https://github.com/QuantumKitHub/MPSKitModels.jl", rev = "ksh/bump"}
+TensorKit = {url = "https://github.com/QuantumKitHub/TensorKit.jl", rev = "main"}
+TensorKitManifolds = {url = "https://github.com/QuantumKitHub/TensorKitManifolds.jl", rev = "ksh/bump"}
+TensorKitTensors = {url = "https://github.com/QuantumKitHub/TensorKitTensors.jl", rev = "ksh/compat"}
+BlockTensorKit = {url = "https://github.com/QuantumKitHub/BlockTensorKit.jl", rev = "ksh/mps"}

ext/MPSKitAdaptExt.jl

@@ -35,5 +35,7 @@ end
     MPSKit.JordanMPOTensor(space(W), adapt(to, W.A), adapt(to, W.B), adapt(to, W.C), adapt(to, W.D))
 @inline Adapt.adapt_structure(to, mpo::MPOHamiltonian) =
     MPOHamiltonian(map(x -> adapt(to, x), mpo.W))
+@inline Adapt.adapt_structure(to, ml::MPSKit.Multiline) = MPSKit.Multiline(map(adapt(to), ml.data))
+@inline Adapt.adapt_structure(to, pa::MPSKit.PeriodicArray) = MPSKit.PeriodicArray(map(adapt(to), pa.data))
 
 end

src/operators/abstractmpo.jl

@@ -99,7 +99,7 @@ Base.:\(α::Number, mpo::AbstractMPO) = scale(mpo, inv(α))
 
 function VectorInterface.scale(mpo::AbstractMPO, α::Number)
     T = VectorInterface.promote_scale(scalartype(mpo), scalartype(α))
-    dst = similar(mpo, T)
+    dst = similar(mpo, TensorKit.similarstoragetype(storagetype(mpo), T))
     return scale!(dst, mpo, α)
 end
 
@@ -132,7 +132,8 @@ function _fuse_mpo_mpo(O1::MPOTensor, O2::MPOTensor, Fₗ, Fᵣ)
         T = promote_type(scalartype(O1), scalartype(O2))
         V = fuse(left_virtualspace(O2) ⊗ left_virtualspace(O1)) ⊗ physicalspace(O1) ←
             physicalspace(O2) ⊗ fuse(right_virtualspace(O2) ⊗ right_virtualspace(O1))
-        return BraidingTensor{T}(V)
+        A = promote_storagetype(storagetype(O1), storagetype(O2), T)
+        return BraidingTensor{T, typeof(V), A}(V)
     elseif O1 isa BraidingTensor
         @plansor O′[-1 -2; -3 -4] := Fₗ[-1; 1 2] * O2[1 3; -3 5] *
             τ[2 -2; 3 4] * conj(Fᵣ[-4; 5 4])
@@ -152,16 +153,21 @@ Compute the mpo tensor that arises from multiplying MPOs.
 """
 function fuse_mul_mpo(O1, O2)
     TT = promote_type(scalartype(O1), scalartype(O2))
-    T = TensorKit.similarstoragetype(storagetype(O1), TT)
+    T = if O1 isa BraidingTensor
+        TensorKit.similarstoragetype(storagetype(O2), TT)
+    else
+        TensorKit.similarstoragetype(storagetype(O1), TT)
+    end
     F_left = fuser(T, left_virtualspace(O2), left_virtualspace(O1))
     F_right = fuser(T, right_virtualspace(O2), right_virtualspace(O1))
     return _fuse_mpo_mpo(O1, O2, F_left, F_right)
 end
-function fuse_mul_mpo(O1::BraidingTensor, O2::BraidingTensor)
-    T = promote_type(scalartype(O1), scalartype(O2))
+function fuse_mul_mpo(O1::BraidingTensor{T1, S1, A1}, O2::BraidingTensor{T2, S2, A2}) where {T1, T2, S1, S2, A1, A2}
+    T = promote_type(T1, T2)
+    A = TensorKit.promote_storagetype(A1, A2)
     V = fuse(left_virtualspace(O2) ⊗ left_virtualspace(O1)) ⊗ physicalspace(O1) ←
         physicalspace(O2) ⊗ fuse(right_virtualspace(O2) ⊗ right_virtualspace(O1))
-    return BraidingTensor{T}(V)
+    return BraidingTensor{T, spacetype(V), A}(V)
 end
 function fuse_mul_mpo(
         O1::AbstractBlockTensorMap{T₁, S, 2, 2}, O2::AbstractBlockTensorMap{T₂, S, 2, 2}
@@ -198,7 +204,7 @@ function add_physical_charge(O::BraidingTensor, charge::Sector)
     auxspace = Vect[typeof(charge)](charge => 1)'
     V = left_virtualspace(O) ⊗ fuse(physicalspace(O), auxspace) ←
         fuse(physicalspace(O), auxspace) ⊗ right_virtualspace(O)
-    return BraidingTensor{scalartype(O)}(V)
+    return BraidingTensor{scalartype(O), typeof(V), storagetype(O)}(V)
 end
 function add_physical_charge(O::AbstractBlockTensorMap{<:Any, <:Any, 2, 2}, charge::Sector)
     sectortype(O) == typeof(charge) || throw(SectorMismatch())

src/operators/jordanmpotensor.jl

@@ -60,7 +60,7 @@ end
 
 const JordanMPOTensorMap{T, S, A <: DenseVector{T}} = JordanMPOTensor{
     T, S,
-    Union{TensorMap{T, S, 2, 2, A}, BraidingTensor{T, S}},
+    Union{TensorMap{T, S, 2, 2, A}, BraidingTensor{T, S, A}},
     TensorMap{T, S, 2, 1, A},
     TensorMap{T, S, 1, 2, A},
     TensorMap{T, S, 1, 1, A},
@@ -124,7 +124,8 @@ end
 function jordanmpotensortype(::Type{S}, ::Type{E}) where {S <: VectorSpace, E}
     TA = tensormaptype(S, 2, 2, E)
     T = scalartype(TA)
-    Tτ = BraidingTensor{T, S}
+    TT = storagetype(TA)
+    Tτ = BraidingTensor{T, S, TT}
     TB = tensormaptype(S, 2, 1, E)
     TC = tensormaptype(S, 1, 2, E)
     TD = tensormaptype(S, 1, 1, E)
@@ -135,6 +136,10 @@ function jordanmpotensortype(::Type{O}) where {O <: AbstractTensorMap}
 end
 Base.similar(W::JordanMPOTensor, ::Type{TorA}) where {TorA} =
     jordanmpotensortype(spacetype(W), TorA)(undef, space(W))
+Base.similar(W::JordanMPOTensor, V::VectorSpace) =
+    jordanmpotensortype(spacetype(V), storagetype(W))(undef, V)
+Base.similar(W::JordanMPOTensor, V::VectorSpace, ::Type{TorA}) where {TorA} =
+    jordanmpotensortype(spacetype(V), TorA)(undef, V)
 
 # Properties
 # ----------
@@ -172,7 +177,7 @@ BlockTensorKit.issparse(W::JordanMPOTensor) = true
 # Converters
 # ----------
 function BlockTensorKit.SparseBlockTensorMap(W::JordanMPOTensor)
-    τ = BraidingTensor{scalartype(W)}(eachspace(W)[1])
+    τ = BraidingTensor{scalartype(W), typeof(eachspace(W)[1]), storagetype(W)}(eachspace(W)[1])
     W′ = SparseBlockTensorMap{AbstractTensorMap{scalartype(W), spacetype(W), 2, 2}}(
         undef_blocks, space(W)
     )
@@ -205,7 +210,8 @@ end
 for f in (:real, :complex)
     @eval function Base.$f(W::JordanMPOTensor)
         E = $f(scalartype(W))
-        W′ = JordanMPOTensor{E}(undef, space(W))
+        TE = TensorKit.similarstoragetype(TensorKit.storagetype(W), E)
+        W′ = similar(W, TE)
         for (I, v) in nonzero_pairs(W.A)
             W′.A[I] = $f(v)
         end
@@ -232,7 +238,10 @@ end
     j = I[4]
     if (size(W, 4) > 1 && i == 1 && j == 1) ||
             (size(W, 1) > 1 && i == size(W, 1) && j == size(W, 4))
-        return BraidingTensor{scalartype(W)}(eachspace(W)[1])
+        T = scalartype(W)
+        TA = storagetype(W)
+        S = spacetype(eachspace(W)[1])
+        return BraidingTensor{T, S, TA}(eachspace(W)[1])
     elseif i == 1 && j == size(W, 4)
         return insertrightunit(insertleftunit(only(W.D), 1), 3)
     elseif i == 1
@@ -242,7 +251,7 @@ end
     elseif 1 < i < size(W, 1) && 1 < j < size(W, 4)
         return W.A[i - 1, 1, 1, j - 1]
     else
-        return zeros(scalartype(W), eachspace(W)[i, 1, 1, j])
+        return zeros(storagetype(W), eachspace(W)[i, 1, 1, j])
     end
 end
 
@@ -357,7 +366,7 @@ function add_physical_charge(O::JordanMPOTensor, charge::Sector)
     Vdst = left_virtualspace(O) ⊗
         fuse(physicalspace(O), auxspace) ←
         fuse(physicalspace(O), auxspace) ⊗ right_virtualspace(O)
-    Odst = JordanMPOTensor{scalartype(O)}(undef, Vdst)
+    Odst = similar(O, Vdst)
     for (I, v) in nonzero_pairs(O)
         Odst[I] = add_physical_charge(v, charge)
     end

src/operators/mpo.jl

@@ -62,7 +62,7 @@ eachsite(mpo::InfiniteMPO) = PeriodicArray(eachindex(mpo))
 function Base.similar(mpo::MPO{<:MPOTensor}, ::Type{O}, L::Int) where {O <: MPOTensor}
     return MPO(similar(parent(mpo), O, L))
 end
-function Base.similar(mpo::MPO, ::Type{T}) where {T <: Number}
+function Base.similar(mpo::MPO, ::Type{T}) where {T <: Union{Number, DenseVector}}
     return MPO(similar.(parent(mpo), T))
 end
 

src/operators/mpohamiltonian.jl

@@ -59,19 +59,19 @@ function InfiniteMPOHamiltonian(Ws::AbstractVector{O}) where {O <: MPOTensor}
 end
 
 """
-    FiniteMPOHamiltonian(Ws::Vector{<:Matrix})
+    FiniteMPOHamiltonian(Ws::Vector{<:AbstractMatrix})
 
 Create a `FiniteMPOHamiltonian` from a vector of matrices, such that `Ws[i][j, k]` represents
 the operator at site `i`, left level `j` and right level `k`. Here, the entries can be
 either `MPOTensor`, `Missing` or `Number`.
 """
-function FiniteMPOHamiltonian(Ws::Vector{<:Matrix})
+function FiniteMPOHamiltonian(Ws::Vector{<:AbstractMatrix})
     T = promote_type(_split_mpoham_types.(Ws)...)
     W = jordanmpotensortype(T)
     return FiniteMPOHamiltonian{W}(Ws)
 end
-function FiniteMPOHamiltonian{O}(W_mats::Vector{<:Matrix}) where {O <: JordanMPOTensor}
-    T = scalartype(O)
+function FiniteMPOHamiltonian{O}(W_mats::Vector{<:AbstractMatrix}) where {O <: JordanMPOTensor}
+    T = storagetype(O)
     L = length(W_mats)
     # initialize sumspaces
     S = spacetype(O)
@@ -140,7 +140,8 @@ function FiniteMPOHamiltonian{O}(W_mats::Vector{<:Matrix}) where {O <: JordanMPO
             if v isa MPOTensor
                 W[I] = v
             elseif !iszero(v)
-                τ = BraidingTensor{T}(eachspace(W)[I])
+                A = similarstoragetype(T, eltype(T))
+                τ = BraidingTensor{eltype(T), typeof(eachspace(W)[I]), A}(eachspace(W)[I])
                 W[I] = isone(v) ? τ : τ * v
             end
         end
@@ -157,12 +158,12 @@ Create a `InfiniteMPOHamiltonian` from a vector of matrices, such that `Ws[i][j,
 the the operator at site `i`, left level `j` and right level `k`. Here, the entries can be
 either `MPOTensor`, `Missing` or `Number`.
 """
-function InfiniteMPOHamiltonian(Ws::Vector{<:Matrix})
+function InfiniteMPOHamiltonian(Ws::Vector{<:AbstractMatrix})
     T = promote_type(_split_mpoham_types.(Ws)...)
     TW = jordanmpotensortype(T)
     return InfiniteMPOHamiltonian{TW}(Ws)
 end
-function InfiniteMPOHamiltonian{O}(W_mats::Vector{<:Matrix}) where {O <: MPOTensor}
+function InfiniteMPOHamiltonian{O}(W_mats::Vector{<:AbstractMatrix}) where {O <: MPOTensor}
     # InfiniteMPOHamiltonian only works for square matrices:
     for W_mat in W_mats
         size(W_mat, 1) == size(W_mat, 2) ||
@@ -172,7 +173,7 @@ function InfiniteMPOHamiltonian{O}(W_mats::Vector{<:Matrix}) where {O <: MPOTens
         throw(ArgumentError("matrices should have the same size"))
     nlvls = size(W_mats[1], 1)
 
-    T = scalartype(O)
+    T = storagetype(O)
     L = length(W_mats)
     # initialize sumspaces
     S = spacetype(O)
@@ -261,7 +262,8 @@ function InfiniteMPOHamiltonian{O}(W_mats::Vector{<:Matrix}) where {O <: MPOTens
             if v isa MPOTensor
                 W[I] = v
             elseif !iszero(v)
-                τ = BraidingTensor{T}(eachspace(W)[I])
+                A = similarstoragetype(T, eltype(T))
+                τ = BraidingTensor{eltype(T), typeof(eachspace(W)[I]), A}(eachspace(W)[I])
                 W[I] = isone(v) ? τ : τ * v
             end
         end
@@ -477,7 +479,8 @@ function FiniteMPOHamiltonian(lattice::AbstractArray{<:VectorSpace}, local_opera
             key_R = key_R′ == 0 ? length(virtualsumspaces[site + 1]) : key_R′
             O[key_L, 1, 1, key_R] += if o isa Number
                 iszero(o) && continue
-                τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
+                S = spacetype(eachspace(O)[key_L, 1, 1, key_R])
+                τ = BraidingTensor{scalartype(TW), S, storagetype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
                 isone(o) ? τ : τ * o
             else
                 o
@@ -504,10 +507,8 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
     end
 
     # partial sort by interaction range
-    local_mpos = sort!(
-        map(Base.Fix1(instantiate_operator, lattice), collect(local_operators));
-        by = x -> length(x[1])
-    )
+    unsorted_mpos = map(Base.Fix1(instantiate_operator, lattice), [local_operators...])
+    local_mpos = sort!(unsorted_mpos; by = x -> length(x[1]))
 
     for (sites, local_mpo) in local_mpos
         local key_R # trick to define key_R before the first iteration
@@ -600,7 +601,8 @@ function InfiniteMPOHamiltonian(lattice′::AbstractArray{<:VectorSpace}, local_
             key_R = key_R′ == 0 ? length(virtualspaces[site]) : key_R′
             O[key_L, 1, 1, key_R] += if o isa Number
                 iszero(o) && continue
-                τ = BraidingTensor{scalartype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
+                S = typeof(eachspace(O)[key_L, 1, 1, key_R])
+                τ = BraidingTensor{scalartype(TW), S, storagetype(TW)}(eachspace(O)[key_L, 1, 1, key_R])
                 isone(o) ? τ : τ * o
             else
                 o
@@ -703,8 +705,8 @@ end
 function Base.similar(H::MPOHamiltonian, ::Type{O}, L::Int) where {O <: MPOTensor}
     return MPOHamiltonian(similar(parent(H), O, L))
 end
-function Base.similar(H::MPOHamiltonian, ::Type{T}) where {T <: Number}
-    return MPOHamiltonian(similar.(parent(H), T))
+function Base.similar(H::MPOHamiltonian, ::Type{TorA}) where {TorA <: Union{Number, DenseVector}}
+    return MPOHamiltonian(similar.(parent(H), TorA))
 end
 
 # Linear Algebra
@@ -729,7 +731,7 @@ function Base.:+(
             ⊞(Vtriv, right_virtualspace(A), Vtriv)
         V = Vleft ⊗ physicalspace(A) ← physicalspace(A) ⊗ Vright
 
-        H[i] = eltype(H)(V, A, B, C, D)
+        H[i] = JordanMPOTensor(V, A, B, C, D)
     end
     return FiniteMPOHamiltonian(H)
 end
@@ -751,7 +753,7 @@ function Base.:+(
         Vright = ⊞(Vtriv, right_virtualspace(A), Vtriv)
         V = Vleft ⊗ physicalspace(A) ← physicalspace(A) ⊗ Vright
 
-        H[i] = eltype(H)(V, A, B, C, D)
+        H[i] = JordanMPOTensor(V, A, B, C, D)
     end
     return InfiniteMPOHamiltonian(H)
 end
@@ -821,7 +823,7 @@ function Base.:*(H::FiniteMPOHamiltonian, mps::FiniteMPS)
         A,
         tensormaptype(
             spacetype(mps), numout(eltype(mps)), numin(eltype(mps)),
-            promote_type(scalartype(H), scalartype(mps))
+            promote_type(storagetype(H), storagetype(mps))
         )
     )
     # left to middle

src/operators/ortho.jl

@@ -11,7 +11,14 @@ function left_canonicalize!(
     d = sqrt(dim(P))
 
     # orthogonalize second column against first
-    WI = removeunit(W[1, 1, 1, 1], 1)
+    Wi = W[1, 1, 1, 1]
+    if Wi isa BraidingTensor
+        Wi′ = removeunit(Wi, 1)
+        WI = similar(Wi′, TensorKit.storagetype(H))
+        copy!(WI, Wi′)
+    else
+        WI = removeunit(Wi, 1)
+    end
     @plansor t[l; r] := conj(WI[p; p' l]) * W.C[p; p' r]
     # TODO: the following is currently broken due to a TensorKit bug
     # @plansor C′[p; p' r] := W.C[p; p' r] - WI[p; p' l] * t[l; r]
@@ -99,7 +106,14 @@ function right_canonicalize!(
     d = sqrt(dim(P))
 
     # orthogonalize second row against last
-    WI = removeunit(W[end, 1, 1, end], 4)
+    Wi = W[end, 1, 1, end]
+    if Wi isa BraidingTensor
+        Wi′ = removeunit(Wi, 4)
+        WI = similar(Wi′, TensorKit.storagetype(H))
+        copy!(WI, Wi′)
+    else
+        WI = removeunit(Wi, 4)
+    end
     @plansor t[l; r] := conj(WI[r p; p']) * W.B[l p; p']
     # TODO: the following is currently broken due to a TensorKit bug
     # @plansor B′[l p; p'] := W.B[l p; p'] - WI[r p; p'] * t[l; r]
@@ -124,7 +138,8 @@ function right_canonicalize!(
         end
         H[i] = JordanMPOTensor(V ⊗ P ← domain(W), Q1, Q2, W.C, W.D)
     else
-        tmp = transpose(cat(insertleftunit(B′, 4), W.A; dims = 4), ((1,), (3, 4, 2)))
+        B′′ = insertleftunit(B′, 4)
+        tmp = transpose(cat(B′′, W.A; dims = 4), ((1,), (3, 4, 2)))
         R, Q = right_orth!(tmp; alg)
         if dim(R) == 0
             V = _rightunit ⊞ _rightunit

test/Project.toml

@@ -24,12 +24,12 @@ MPSKit = {path = ".."}
 
 [compat]
 Aqua = "0.8.9"
-CUDA = "5.9"
+CUDA = "5, 6"
 Combinatorics = "1"
 ParallelTestRunner = "2"
 Plots = "1.40"
 Pkg = "1"
 TensorKitTensors = "0.2"
 Test = "1"
 TestExtras = "0.3"
-cuTENSOR = "2.3"
+cuTENSOR = "1, 6"