Adapt extension (#384)

* add adapt extension

* add tests

* centralize storagetype definition

* accept storagetype for `jordanmpotensortype`

* small fixes

Author: lkdvos

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"]

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

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))

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(ψ)])

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)

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

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

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)

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