Merge branch 'master' into feature/compression-v2

Author: josepamigoo

.github/workflows/CI.yml

@@ -48,7 +48,7 @@ jobs:
           include-all-prereleases: true
       - uses: julia-actions/cache@v1
         with:
-          cache-name: Unit Tests CI - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+          cache-name: CI / ${{ matrix.test_group }} / Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
       - name: Add Julia registries
         run: |
           using Pkg
@@ -63,6 +63,52 @@ jobs:
       - uses: codecov/codecov-action@v3
         with:
           files: lcov.info
+  python:
+    name: ${{ matrix.test_group }} / Julia ${{ matrix.version }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1'
+        test_group:
+          - python
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+          include-all-prereleases: true
+      - uses: julia-actions/cache@v1
+        with:
+          cache-name: CI / ${{ matrix.test_group }} / Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+      - name: Set and export registry flavor preference
+        run: echo "JULIA_PKG_SERVER_REGISTRY_PREFERENCE=${JULIA_PKG_SERVER_REGISTRY_PREFERENCE:-eager}" >> ${GITHUB_ENV}
+        shell: bash
+      - name: Add Julia registries
+        run: |
+          using Pkg
+          pkg"registry add https://github.com/bsc-quantic/Registry.git"
+          pkg"registry add General"
+        shell: julia --color=yes {0}
+      - name: Run tests
+        run: |
+          julia --color=yes --code-coverage=user --depwarn=yes --project=test/python/ -e '
+            using Pkg
+            Pkg.instantiate()
+            Pkg.resolve()
+            include("test/python/runtests.jl")'
+        shell: bash
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v3
+        with:
+          files: lcov.info
   docs:
     name: Documentation
     runs-on: ubuntu-latest

.gitignore

@@ -420,6 +420,6 @@ dist
 .julia
 *.excalidraw
 archive/
-test/.CondaPkg/
+**/*/.CondaPkg/
 .CondaPkg/
 CondaPkg.toml

Project.toml

@@ -1,7 +1,7 @@
 name = "Tenet"
 uuid = "85d41934-b9cd-44e1-8730-56d86f15f3ec"
 authors = ["Sergio Sánchez Ramírez <[email protected]>"]
-version = "0.8.0"
+version = "0.8.3"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -71,15 +71,15 @@ GraphMakie = "0.4,0.5"
 Graphs = "1.7"
 ITensorMPS = "0.2, 0.3"
 ITensorNetworks = "0.11"
-ITensors = "0.6, 0.7"
+ITensors = "0.6, 0.7, 0.8"
 KrylovKit = "0.7, 0.8, 0.9"
 LinearAlgebra = "1.10"
 Makie = "0.18,0.19,0.20, 0.21, 0.22"
 OMEinsum = "0.7, 0.8"
 PythonCall = "0.9"
 Quac = "0.3"
 Random = "1.10"
-Reactant = "0.2.18"
+Reactant = "0.2.22"
 ScopedValues = "1"
 Serialization = "1.10"
 SparseArrays = "1.10"

ext/TenetITensorMPSExt.jl

@@ -1,10 +1,9 @@
 module TenetITensorMPSExt
 
 using Tenet
-using Tenet: Tenet, MPS, tensors, form, inds, lanes, id, Site, Lane
-using ITensors
-using ITensorMPS
-using ITensors: ITensor, Index, dim
+using Tenet: Tenet, MPS, tensors, form, inds, lanes, Site, Lane
+using ITensors: ITensors, ITensor, Index, dim, siteinds
+using ITensorMPS: ITensorMPS, linkinds
 
 # Convert an AbstractMPS to an ITensor MPS
 function Base.convert(::Type{ITensorMPS.MPS}, mps::Tenet.AbstractMPS)
@@ -76,15 +75,15 @@ function Base.convert(::Type{MPS}, itensors_mps::ITensorMPS.MPS)
     links = linkinds(itensors_mps)
 
     tensors_vec = []
-    first_ten = array(itensors_mps[1], sites[1], links[1])
+    first_ten = ITensors.array(itensors_mps[1], sites[1], links[1])
     push!(tensors_vec, first_ten)
 
     # Extract the bulk tensors
     for j in 2:(length(itensors_mps) - 1)
-        ten = array(itensors_mps[j], sites[j], links[j - 1], links[j]) # Indices are ordered as (site index, left link, right link)
+        ten = ITensors.array(itensors_mps[j], sites[j], links[j - 1], links[j]) # Indices are ordered as (site index, left link, right link)
         push!(tensors_vec, ten)
     end
-    last_ten = array(itensors_mps[end], sites[end], links[end])
+    last_ten = ITensors.array(itensors_mps[end], sites[end], links[end])
     push!(tensors_vec, last_ten)
 
     mps = Tenet.MPS(tensors_vec)

ext/TenetReactantExt.jl

@@ -9,26 +9,41 @@ const stablehlo = MLIR.Dialects.stablehlo
 
 const Enzyme = Reactant.Enzyme
 
-@static if isdefined(Reactant, :traced_type_inner)
-    # we specify `mode` and `track_numbers` types due to ambiguity
-    Base.@nospecializeinfer function Reactant.traced_type_inner(
-        @nospecialize(TT::Type{<:Tensor}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
-    )
-        A_traced = Reactant.traced_type_inner(Tenet.parenttype(TT), seen, mode, track_numbers)
-        T = eltype(A_traced)
-        N = ndims(TT)
-        return Tensor{T,N,A_traced}
-    end
+# we specify `mode` and `track_numbers` types due to ambiguity
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(_::Type{Tensor}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+)
+    return Tensor
+end
 
-    # we specify `mode` and `track_numbers` types due to ambiguity
-    Base.@nospecializeinfer function Reactant.traced_type_inner(
-        @nospecialize(T::Type{<:Tenet.AbstractTensorNetwork}),
-        seen,
-        mode::Reactant.TraceMode,
-        @nospecialize(track_numbers::Type)
-    )
-        return T
-    end
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(_::Type{Tensor{T}}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+) where {T}
+    return Tensor{TracedRNumber{T}}
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(_::Type{Tensor{T,N}}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+) where {T,N}
+    return Tensor{TracedRNumber{T,N}}
+end
+
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(_::Type{Tensor{T,N,A}}), seen, mode::Reactant.TraceMode, @nospecialize(track_numbers::Type)
+) where {T,N,A}
+    A_traced = Reactant.traced_type_inner(A, seen, mode, track_numbers)
+    T_traced = eltype(A_traced)
+    return Tensor{T_traced,N,A_traced}
+end
+
+# we specify `mode` and `track_numbers` types due to ambiguity
+Base.@nospecializeinfer function Reactant.traced_type_inner(
+    @nospecialize(T::Type{<:Tenet.AbstractTensorNetwork}),
+    seen,
+    mode::Reactant.TraceMode,
+    @nospecialize(track_numbers::Type)
+)
+    return T
 end
 
 function Reactant.make_tracer(seen, @nospecialize(prev::RT), @nospecialize(path), mode; kwargs...) where {RT<:Tensor}
@@ -184,8 +199,8 @@ Base.@nospecializeinfer @noinline function Tenet.contract(
 
     # TODO replace for `Ops.convert`/`adapt` when it's available (there can be problems with nested array structures)
     T = Base.promote_eltype(a, b)
-    da = eltype(a) != T ? TracedRArray{T,ndims(a)}(parent(a)) : parent(a)
-    db = eltype(b) != T ? TracedRArray{T,ndims(b)}(parent(b)) : parent(b)
+    da = eltype(a) != T ? TracedRArray{Reactant.unwrapped_eltype(T),ndims(a)}(parent(a)) : parent(a)
+    db = eltype(b) != T ? TracedRArray{Reactant.unwrapped_eltype(T),ndims(b)}(parent(b)) : parent(b)
 
     data = Reactant.Ops.dot_general(da, db; contracting_dimensions, batching_dimensions)
 
@@ -221,11 +236,4 @@ function Base.conj(@nospecialize(x::Tensor{TracedRNumber{T},N,<:TracedRArray}))
     Tensor(conj(parent(x)), inds(x))
 end
 
-# fix infinite recursion on Reactant rewrite of invoke/call step
-@reactant_overlay @noinline function Base.replace!(
-    tn::Tenet.AbstractQuantum, old_new::Base.AbstractVecOrTuple{Pair{Symbol,Symbol}}
-)
-    Base.inferencebarrier(Base.replace!)(tn, old_new)
-end
-
 end

src/Ansatz.jl

@@ -165,8 +165,12 @@ function inds(kwargs::NamedTuple{(:bond,)}, tn::AbstractAnsatz)
     return only(inds(tensor1) ∩ inds(tensor2))
 end
 
-# TODO fix this properly when we do the mapping 
-tensors(kwargs::NamedTuple{(:at,),Tuple{L}}, tn::AbstractAnsatz) where {L<:Lane} = tensors(tn; at=Site(kwargs.at))
+# TODO fix this properly when we do the mapping
+function tensors(kwargs::NamedTuple{(:at,),Tuple{L}}, tn::AbstractAnsatz) where {L<:Lane}
+    hassite(tn, Site(kwargs.at)) && return tensors(tn; at=Site(kwargs.at))
+    hassite(tn, Site(kwargs.at; dual=true)) && return tensors(tn; at=Site(kwargs.at; dual=true))
+    throw(ArgumentError("Lane $kwargs.at not found"))
+end
 
 """
     tensors(tn::AbstractAnsatz; bond)
@@ -375,7 +379,8 @@ Compute the expectation value of an observable on a [`AbstractAnsatz`](@ref) Ten
 function expect(ψ::AbstractAnsatz, observable; bra=adjoint(ψ))
     @assert socket(ψ) == State() "ψ must be a state"
     @assert socket(bra) == State(; dual=true) "bra must be a dual state"
-    contract(merge(ψ, observable, bra))
+
+    return expect(form(ψ), ψ, observable; bra)
 end
 
 function expect(ψ::AbstractAnsatz, observables::AbstractVecOrTuple; bra=adjoint(ψ))
@@ -384,6 +389,50 @@ function expect(ψ::AbstractAnsatz, observables::AbstractVecOrTuple; bra=adjoint
     end
 end
 
+function expect(::NonCanonical, ψ::AbstractAnsatz, observable; bra=adjoint(ψ))
+    return contract(merge(ψ, observable, bra))
+end
+
+# TODO: Try to find a better way to do this
+function expect(::MixedCanonical, ψ::AbstractAnsatz, observable; bra=adjoint(ψ))
+    return contract(merge(ψ, observable, bra))
+end
+
+function expect(::Canonical, ψ::Tenet.AbstractAnsatz, observable; bra=adjoint(ψ))
+    obs_sites = unique(id.(sites(observable)))
+
+    ket_Λ = []
+    bra_Λ = []
+    ket_tensors = []
+    bra_tensors = []
+    for i in obs_sites
+        replace!(observable, inds(observable; at=Site(i)) => Symbol(:input, i))
+        replace!(observable, inds(observable; at=Site(i; dual=true)) => Symbol(:output, i))
+        replace!(ψ, inds(ψ; at=Site(i)) => Symbol(:input, i))
+        replace!(bra, inds(bra; at=Site(i; dual=true)) => Symbol(:output, i))
+
+        replace!(bra, inds(bra; bond=(Lane(i), Lane(i + 1))) => inds(ψ; bond=(Lane(i), Lane(i + 1))))
+        replace!(bra, inds(bra; bond=(Lane(i - 1), Lane(i))) => inds(ψ; bond=(Lane(i - 1), Lane(i))))
+
+        push!(ket_Λ, tensors(ψ; bond=(Lane(i - 1), Lane(i))))
+        push!(bra_Λ, tensors(bra; bond=(Lane(i - 1), Lane(i))))
+
+        push!(ket_tensors, tensors(ψ; at=Site(i)))
+        push!(bra_tensors, tensors(bra; at=Site(i; dual=true)))
+    end
+
+    push!(ket_Λ, tensors(ψ; bond=(Lane(obs_sites[end]), Lane(obs_sites[end] + 1))))
+    push!(bra_Λ, tensors(bra; bond=(Lane(obs_sites[end]), Lane(obs_sites[end] + 1))))
+
+    t = contract(
+        contract(ket_Λ..., ket_tensors...; dims=[]),
+        contract(bra_Λ..., bra_tensors...; dims=[]),
+        tensors(Quantum(observable))[1],
+    )
+
+    return t
+end
+
 """
     evolve!(ψ::AbstractAnsatz, gate; threshold = nothing, maxdim = nothing, normalize = false)
 
@@ -485,10 +534,10 @@ function simple_update_2site!(::NonCanonical, ψ::AbstractAnsatz, gate; kwargs..
     gate = copy(gate)
 
     # contract involved sites
-    bond = (sitel, siter) = extrema(lanes(gate))
+    bond = (lanel, laner) = extrema(lanes(gate))
     vind = inds(ψ; bond)
-    linds = filter(!=(vind), inds(tensors(ψ; at=sitel)))
-    rinds = filter(!=(vind), inds(tensors(ψ; at=siter)))
+    linds = filter(!=(vind), inds(tensors(ψ; at=lanel)))
+    rinds = filter(!=(vind), inds(tensors(ψ; at=laner)))
     contract!(ψ; bond)
 
     # TODO replace for `merge!` when #243 is fixed
@@ -519,20 +568,20 @@ end
 # TODO remove `normalize` argument?
 function simple_update_2site!(::Canonical, ψ::AbstractAnsatz, gate; threshold, maxdim, normalize=false, canonize=true)
     # Contract the exterior Λ tensors
-    sitel, siter = extrema(lanes(gate))
-    (0 < id(sitel) < nlanes(ψ) || 0 < id(siter) < nlanes(ψ)) ||
+    lanel, laner = extrema(lanes(gate))
+    (0 < id(lanel) < nlanes(ψ) || 0 < id(laner) < nlanes(ψ)) ||
         throw(ArgumentError("The sites in the bond must be between 1 and $(nlanes(ψ))"))
 
-    Λᵢ₋₁ = id(sitel) == 1 ? nothing : tensors(ψ; between=(Lane(id(sitel) - 1), sitel))
-    Λᵢ₊₁ = id(sitel) == nsites(ψ) - 1 ? nothing : tensors(ψ; between=(siter, Lane(id(siter) + 1)))
+    Λᵢ₋₁ = id(lanel) == 1 ? nothing : tensors(ψ; bond=(Lane(id(lanel) - 1), lanel))
+    Λᵢ₊₁ = id(lanel) == nsites(ψ) - 1 ? nothing : tensors(ψ; bond=(laner, Lane(id(laner) + 1)))
 
-    !isnothing(Λᵢ₋₁) && contract!(ψ; between=(Lane(id(sitel) - 1), sitel), direction=:right, delete_Λ=false)
-    !isnothing(Λᵢ₊₁) && contract!(ψ; between=(siter, Lane(id(siter) + 1)), direction=:left, delete_Λ=false)
+    !isnothing(Λᵢ₋₁) && absorb!(ψ; bond=(Lane(id(lanel) - 1), lanel), dir=:right, delete_Λ=false)
+    !isnothing(Λᵢ₊₁) && absorb!(ψ; bond=(laner, Lane(id(laner) + 1)), dir=:left, delete_Λ=false)
 
     simple_update_2site!(NonCanonical(), ψ, gate; threshold, maxdim, normalize=false, canonize=false)
 
     # contract the updated tensors with the inverse of Λᵢ and Λᵢ₊₂, to get the new Γ tensors
-    U, Vt = tensors(ψ; at=sitel), tensors(ψ; at=siter)
+    U, Vt = tensors(ψ; at=lanel), tensors(ψ; at=laner)
     Γᵢ₋₁ = if isnothing(Λᵢ₋₁)
         U
     else
@@ -545,13 +594,13 @@ function simple_update_2site!(::Canonical, ψ::AbstractAnsatz, gate; threshold,
     end
 
     # Update the tensors in the tensor network
-    replace!(ψ, tensors(ψ; at=sitel) => Γᵢ₋₁)
-    replace!(ψ, tensors(ψ; at=siter) => Γᵢ)
+    replace!(ψ, tensors(ψ; at=lanel) => Γᵢ₋₁)
+    replace!(ψ, tensors(ψ; at=laner) => Γᵢ)
 
     if canonize
         canonize!(ψ; normalize)
     else
-        normalize && normalize!(ψ, collect((sitel, siter)))
+        normalize && normalize!(ψ, collect((lanel, laner)))
     end
 
     return ψ

src/Lattice.jl

@@ -14,8 +14,24 @@ Graphs.dst(edge::Bond) = edge.dst
 Graphs.reverse(edge::Bond) = Bond(Graphs.dst(edge), Graphs.src(edge))
 Base.show(io::IO, edge::Bond) = write(io, "Bond: $(Graphs.src(edge)) - $(Graphs.dst(edge))")
 
-Pair(e::Bond) = src(e) => dst(e)
-Tuple(e::Bond) = (src(e), dst(e))
+Pair(e::Bond) = Graphs.src(e) => Graphs.dst(e)
+Tuple(e::Bond) = (Graphs.src(e), Graphs.dst(e))
+
+function Base.iterate(bond::Bond, state=0)
+    if state == 0
+        (Graphs.src(bond), 1)
+    elseif state == 1
+        (Graphs.dst(bond), 2)
+    else
+        nothing
+    end
+end
+
+Base.IteratorSize(::Type{Bond}) = Base.HasLength()
+Base.length(::Bond) = 2
+Base.IteratorEltype(::Type{Bond{L}}) where {L} = Base.HasEltype()
+Base.eltype(::Bond{L}) where {L} = L
+Base.isdone(::Bond, state) = state == 2
 
 """
     Lattice