diff --git a/Project.toml b/Project.toml index 0ba745a7f..e36c838f5 100644 --- a/Project.toml +++ b/Project.toml @@ -23,10 +23,17 @@ TensorKitManifolds = "11fa318c-39cb-4a83-b1ed-cdc7ba1e3684" TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2" VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8" +[weakdeps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" + +[extensions] +MPSKitAdaptExt = "Adapt" + [compat] Accessors = "0.1" +Adapt = "4" Aqua = "0.8.9" -BlockTensorKit = "0.3.4" +BlockTensorKit = "0.3.8" Combinatorics = "1" Compat = "3.47, 4.10" DocStringExtensions = "0.9.3" @@ -43,7 +50,7 @@ Plots = "1.40" Printf = "1" Random = "1" RecipesBase = "1.1" -TensorKit = "0.16" +TensorKit = "0.16.3" TensorKitManifolds = "0.7" TensorKitTensors = "0.2" TensorOperations = "5" @@ -53,6 +60,7 @@ VectorInterface = "0.2, 0.3, 0.4, 0.5" julia = "1.10" [extras] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" @@ -63,4 +71,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [targets] -test = ["Aqua", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"] +test = ["Aqua", "Adapt", "Pkg", "Test", "TestExtras", "Plots", "Combinatorics", "ParallelTestRunner", "TensorKitTensors"] diff --git a/ext/MPSKitAdaptExt.jl b/ext/MPSKitAdaptExt.jl new file mode 100644 index 000000000..6f0b0c7f2 --- /dev/null +++ b/ext/MPSKitAdaptExt.jl @@ -0,0 +1,39 @@ +module MPSKitAdaptExt + +using TensorKit: space, spacetype +using MPSKit +using BlockTensorKit: nonzero_pairs +using Adapt + +function Adapt.adapt_structure(to, mps::FiniteMPS) + ad = adapt(to) + adapt_not_missing(x) = ismissing(x) ? x : ad(x) + + TA = Base.promote_op(ad, MPSKit.site_type(mps)) + TB = Base.promote_op(ad, MPSKit.bond_type(mps)) + + ALs = map!(adapt_not_missing, similar(mps.ALs, Union{Missing, TA}), mps.ALs) + ARs = map!(adapt_not_missing, similar(mps.ARs, Union{Missing, TA}), mps.ARs) + ACs = map!(adapt_not_missing, similar(mps.ACs, Union{Missing, TA}), mps.ACs) + Cs = map!(adapt_not_missing, similar(mps.Cs, Union{Missing, TB}), mps.Cs) + + return FiniteMPS{TA, TB}(ALs, ARs, ACs, Cs) +end + +function Adapt.adapt_structure(to, mps::InfiniteMPS) + ad = adapt(to) + AL = map(ad, mps.AL) + AR = map(ad, mps.AR) + C = map(ad, mps.C) + AC = map(ad, mps.AC) + return InfiniteMPS{eltype(AL), eltype(C)}(AL, AR, C, AC) +end + +# inline to improve type stability with closures +@inline Adapt.adapt_structure(to, mpo::MPO) = MPO(map(adapt(to), mpo.O)) +@inline Adapt.adapt_structure(to, W::MPSKit.JordanMPOTensor) = + 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)) + +end diff --git a/src/states/abstractmps.jl b/src/states/abstractmps.jl index 3d277aff4..e442100db 100644 --- a/src/states/abstractmps.jl +++ b/src/states/abstractmps.jl @@ -200,6 +200,8 @@ TensorKit.spacetype(ψtype::Type{<:AbstractMPS}) = spacetype(site_type(ψtype)) TensorKit.sectortype(ψ::AbstractMPS) = sectortype(typeof(ψ)) TensorKit.sectortype(ψtype::Type{<:AbstractMPS}) = sectortype(site_type(ψtype)) +TensorKit.storagetype(ψtype::Type{<:AbstractMPS}) = storagetype(site_type(ψtype)) + """ left_virtualspace(ψ::AbstractMPS, [pos=1:length(ψ)]) diff --git a/src/states/finitemps.jl b/src/states/finitemps.jl index c956b9730..5bfdcba92 100644 --- a/src/states/finitemps.jl +++ b/src/states/finitemps.jl @@ -387,9 +387,6 @@ end site_type(::Type{<:FiniteMPS{A}}) where {A} = A bond_type(::Type{<:FiniteMPS{<:Any, B}}) where {B} = B -function TensorKit.storagetype(::Union{MPS, Type{MPS}}) where {A, MPS <: FiniteMPS{A}} - return storagetype(A) -end function left_virtualspace(ψ::FiniteMPS, n::Integer) checkbounds(ψ, n) diff --git a/test/operators/mpo.jl b/test/operators/mpo.jl index 4f6d2cad9..f93d4a9c2 100644 --- a/test/operators/mpo.jl +++ b/test/operators/mpo.jl @@ -10,6 +10,7 @@ using MPSKit using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle, MPOStyle using TensorKit using TensorKit: ℙ +using Adapt @testset "FiniteMPO" begin # start from random operators @@ -98,3 +99,25 @@ end @test OperatorStyle(typeof(H)) == MPOStyle() @test OperatorStyle(H) == MPOStyle() end + +@testset "Adapt" for V in (ℂ^2, U1Space(-1 => 1, 0 => 1, 1 => 1)) + L = 3 + o = rand(Float32, V^L ← V^L) + mpo1 = FiniteMPO(o) + for T in (Float64, ComplexF64) + mpo2 = @testinferred adapt(Vector{T}, mpo1) + @test mpo2 isa FiniteMPO + @test scalartype(mpo2) == T + @test storagetype(mpo2) == Vector{T} + @test convert(TensorMap, mpo2) ≈ o + end + + mpo3 = InfiniteMPO(mpo1[2:2]) + for T in (Float64, ComplexF64) + mpo4 = @testinferred adapt(Vector{T}, mpo3) + @test mpo4 isa InfiniteMPO + @test scalartype(mpo4) == T + @test storagetype(mpo4) == Vector{T} + @test dot(mpo3, mpo4) ≈ norm(mpo3)^2 atol = 1.0e-4 + end +end diff --git a/test/operators/mpohamiltonian.jl b/test/operators/mpohamiltonian.jl index 666ab2bee..e199df30f 100644 --- a/test/operators/mpohamiltonian.jl +++ b/test/operators/mpohamiltonian.jl @@ -10,6 +10,7 @@ using MPSKit using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle, HamiltonianStyle using TensorKit using TensorKit: ℙ +using Adapt pspaces = (ℙ^4, Rep[U₁](0 => 2), Rep[SU₂](1 => 1)) vspaces = (ℙ^10, Rep[U₁]((0 => 20)), Rep[SU₂](1 // 2 => 10, 3 // 2 => 5, 5 // 2 => 1)) @@ -239,3 +240,44 @@ end h4 = H4 * H4 @test real(expectation_value(ψ2, H4)) >= 0 end + +@testset "Adapt" for V in (ℂ^2, U1Space(-1 => 1, 0 => 1, 1 => 1)) + h = rand(Float32, V^2 ← V^2) + h += h' + + L = 4 + H1 = FiniteMPOHamiltonian( + fill(V, L), + ((i, i + 1) => h for i in 1:(L - 1))..., + ((i, i + 2) => h for i in 1:(L - 2))..., + ((i, i + 3) => h for i in 1:(L - 3))..., + ) + mps1 = FiniteMPS(physicalspace(H1), oneunit(V)) + + for T in (Float64, ComplexF64) + H2 = if VERSION <= v"1.12" + adapt(Vector{T}, H1) + else + @testinferred adapt(Vector{T}, H1) + end + @test H2 isa FiniteMPOHamiltonian + @test scalartype(H2) == T + @test storagetype(H2) == Vector{T} + @test expectation_value(mps1, H1) ≈ expectation_value(mps1, H2) + end + + H3 = InfiniteMPOHamiltonian(fill(V, L), (1, 2) => h, (1, 3) => h, (1, 4) => h) + mps2 = InfiniteMPS(physicalspace(H3), [oneunit(V)]) + for T in (Float64, ComplexF64) + H4 = if VERSION <= v"1.12" + # this is type unstable for LTS for some reason + adapt(Vector{T}, H3) + else + @testinferred adapt(Vector{T}, H3) + end + @test H4 isa InfiniteMPOHamiltonian + @test scalartype(H4) == T + @test storagetype(H4) == Vector{T} + @test expectation_value(mps2, H3) ≈ expectation_value(mps2, H4) + end +end diff --git a/test/states/finitemps.jl b/test/states/finitemps.jl index c99240782..e098c9c59 100644 --- a/test/states/finitemps.jl +++ b/test/states/finitemps.jl @@ -11,6 +11,7 @@ using MPSKit: _transpose_front, _transpose_tail using MPSKit: GeometryStyle, FiniteChainStyle using TensorKit using TensorKit: ℙ +using Adapt @testset "FiniteMPS ($(sectortype(D)), $elt)" for (D, d, elt) in [ (ℙ^10, ℙ^2, ComplexF64), diff --git a/test/states/infinitemps.jl b/test/states/infinitemps.jl index ec4f46f55..fb27a84d0 100644 --- a/test/states/infinitemps.jl +++ b/test/states/infinitemps.jl @@ -10,6 +10,7 @@ using MPSKit using MPSKit: GeometryStyle, InfiniteChainStyle, TransferMatrix using TensorKit using TensorKit: ℙ +using Adapt @testset "InfiniteMPS ($(sectortype(D)), $elt)" for (D, d, elt) in [(ℙ^10, ℙ^2, ComplexF64), (Rep[U₁](1 => 3), Rep[U₁](0 => 1), ComplexF64)] @@ -71,6 +72,19 @@ end @test mps1.C[1] !== mps2.C[1] end +@testset "Adapt" begin + for (d, D) in [(ℂ^2, ℂ^4), (ℙ^2, ℙ^4)] + mps1 = InfiniteMPS(rand, Float32, d, D) + for T in (Float64, ComplexF64) + mps2 = @testinferred adapt(Vector{T}, mps1) + @test mps2 isa InfiniteMPS + @test scalartype(mps2) == T + @test storagetype(mps2) == Vector{T} + @test dot(mps1, mps2) ≈ 1 atol = 1.0e-4 + end + end +end + @testset "InfiniteMPS entropy ($(sectortype(D)), $elt)" for (D, d, elt) in [(ℙ^10, ℙ^2, ComplexF64), (Rep[U₁](1 => 3), Rep[U₁](0 => 1), ComplexF64)] ψ = InfiniteMPS([d, d], [D, D]) @@ -92,3 +106,28 @@ end Ss_product = entropy(ψ_product) @test all(S -> isapprox(S, 0; atol = 1.0e-10), Ss_product) end + +@testset "InfiniteMPS copying" begin + mps1 = InfiniteMPS(rand, ComplexF64, ℂ^2, ℂ^5) + mps2 = copy(mps1) + + @test mps1 !== mps2 + + # elements are equal + @test mps1.AL[1] == mps2.AL[1] + @test mps1.AR[1] == mps2.AR[1] + @test mps1.AC[1] == mps2.AC[1] + @test mps1.C[1] == mps2.C[1] + + # arrays are distinct + @test mps1.AL !== mps2.AL + @test mps1.AR !== mps2.AR + @test mps1.AC !== mps2.AC + @test mps1.C !== mps2.C + + # tensors are distinct + @test mps1.AL[1] !== mps2.AL[1] + @test mps1.AR[1] !== mps2.AR[1] + @test mps1.AC[1] !== mps2.AC[1] + @test mps1.C[1] !== mps2.C[1] +end