Fix Adapt extension (#389)

* Revert "Adapt extension (#384)"

This reverts commit 2db4cb1.

* Deep copying MPS with copy (#387)

* broadcast copy over arrays

* add basic tests for copying mps

* fix variable naming mismatch

* small fixes

* format

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* Refactor entropy function to also use spectrum directly (#377)

* Refactor entropy function to also use spectrum directly

* specialize to SectorVector and add test

* add infinite test

* update docstring

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* Deep copying MPS with copy (#387)

* broadcast copy over arrays

* add basic tests for copying mps

* fix variable naming mismatch

* small fixes

* format

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* fix wrong testcase

* alter JordanMPOTensor Adapt implementation

* type stability chenanigans

* bump TensorKit minimal version

* more bypassing of type stability issues...

* fix some tests

* lost of git struggles

* more git misery!

* handle storagetype better

* disable lts type stability tests

---------

Co-authored-by: Akshay Shankar <sakshays.2000@gmail.com>
Co-authored-by: Andreas Feuerpfeil <andreas.feuerpfeil@gmail.com>

Author: 3 people

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

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

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

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)

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

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

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

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