diff --git a/Project.toml b/Project.toml index 0ba745a7f..9b39e662f 100644 --- a/Project.toml +++ b/Project.toml @@ -23,8 +23,15 @@ 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" Combinatorics = "1" @@ -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..c6e16fd76 --- /dev/null +++ b/ext/MPSKitAdaptExt.jl @@ -0,0 +1,57 @@ +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 + +Adapt.adapt_structure(to, mpo::MPO) = MPO(map(adapt(to), mpo.O)) + +function Adapt.adapt_structure(::Type{TorA}, W::MPSKit.JordanMPOTensor) where {TorA <: Union{Number, DenseVector{<:Number}}} + TT = MPSKit.jordanmpotensortype(spacetype(W), TorA) + W′ = TT(undef, space(W)) + ad = adapt(TorA) + + for (k, v) in nonzero_pairs(W.A) + W′.A[k] = ad(v) + end + for (k, v) in nonzero_pairs(W.B) + W′.B[k] = ad(v) + end + for (k, v) in nonzero_pairs(W.C) + W′.C[k] = ad(v) + end + for (k, v) in nonzero_pairs(W.D) + W′.D[k] = ad(v) + end + + return W′ +end +Adapt.adapt_structure(to, mpo::MPOHamiltonian) = MPOHamiltonian(map(adapt(to), mpo.W)) + +end diff --git a/src/operators/jordanmpotensor.jl b/src/operators/jordanmpotensor.jl index 4a3dbe798..02f36c987 100644 --- a/src/operators/jordanmpotensor.jl +++ b/src/operators/jordanmpotensor.jl @@ -121,12 +121,12 @@ 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} + TA = Union{tensormaptype(S, 2, 2, TorA), BraidingTensor{scalartype(TorA), S}} + TB = tensormaptype(S, 2, 1, TorA) + TC = tensormaptype(S, 1, 2, TorA) + TD = tensormaptype(S, 1, 1, TorA) + return JordanMPOTensor{scalartype(TorA), S, TA, TB, TC, TD} end function jordanmpotensortype(::Type{O}) where {O <: MPOTensor} return jordanmpotensortype(spacetype(O), scalartype(O)) diff --git a/src/states/abstractmps.jl b/src/states/abstractmps.jl index 3d277aff4..b80cdcc1c 100644 --- a/src/states/abstractmps.jl +++ b/src/states/abstractmps.jl @@ -199,6 +199,7 @@ TensorKit.spacetype(ψ::AbstractMPS) = spacetype(typeof(ψ)) 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 755825c74..1c294ae90 100644 --- a/src/states/finitemps.jl +++ b/src/states/finitemps.jl @@ -376,9 +376,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..34f2553db 100644 --- a/test/operators/mpo.jl +++ b/test/operators/mpo.jl @@ -4,8 +4,8 @@ println(" -------------------- ") -using .TestSetup using Test, TestExtras +using Adapt using MPSKit using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle, MPOStyle using TensorKit @@ -83,6 +83,7 @@ using TensorKit: ℙ end end + @testset "InfiniteMPO" begin P = ℂ^2 T = Float64 @@ -98,3 +99,26 @@ 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) ≈ 1 atol = 1.0e-4 + end + +end diff --git a/test/operators/mpohamiltonian.jl b/test/operators/mpohamiltonian.jl index 666ab2bee..7eae64ea9 100644 --- a/test/operators/mpohamiltonian.jl +++ b/test/operators/mpohamiltonian.jl @@ -4,8 +4,8 @@ println(" ---------------------------- ") -using .TestSetup using Test, TestExtras +using Adapt using MPSKit using MPSKit: GeometryStyle, FiniteChainStyle, InfiniteChainStyle, OperatorStyle, HamiltonianStyle using TensorKit @@ -239,3 +239,36 @@ 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 = @testinferred adapt(Vector{T}, H1) + @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 = @testinferred adapt(Vector{T}, H3) + @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 93e9e2b3a..d23e26343 100644 --- a/test/states/finitemps.jl +++ b/test/states/finitemps.jl @@ -4,8 +4,8 @@ println(" ---------------------- ") -using .TestSetup using Test, TestExtras +using Adapt using MPSKit using MPSKit: _transpose_front, _transpose_tail using MPSKit: GeometryStyle, FiniteChainStyle @@ -76,6 +76,20 @@ end end end +@testset "Adapt" begin + for (d, D) in [(ℂ^2, ℂ^4), (ℙ^2, ℙ^4)] + mps1 = FiniteMPS(rand, Float32, 3, d, D) + t1 = convert(TensorMap, mps1) + for T in (Float64, ComplexF64) + mps2 = @testinferred adapt(Vector{T}, mps1) + @test mps2 isa FiniteMPS + @test scalartype(mps2) == T + @test storagetype(mps2) == Vector{T} + @test convert(TensorMap, mps2) ≈ t1 + end + end +end + @testset "FiniteMPS center + (slice) indexing" begin L = 11 ψ = FiniteMPS(L, ℂ^2, ℂ^16) diff --git a/test/states/infinitemps.jl b/test/states/infinitemps.jl index 7398c6bb8..0ac4e269c 100644 --- a/test/states/infinitemps.jl +++ b/test/states/infinitemps.jl @@ -4,8 +4,8 @@ println(" ------------------------ ") -using .TestSetup using Test, TestExtras +using Adapt using MPSKit using MPSKit: GeometryStyle, InfiniteChainStyle, TransferMatrix using TensorKit @@ -45,3 +45,16 @@ using TensorKit: ℙ @test TransferMatrix(ψ.AR[i], ψ.AR[i]) * r_RR(ψ, i) ≈ r_RR(ψ, i + 1) end 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