Support ProjMPSs

tensor4all · Nov 17, 2024 · 046fcf1 · 046fcf1
1 parent 550c441
commit 046fcf1
Show file tree

Hide file tree

Showing 7 changed files with 398 additions and 27 deletions.
diff --git a/Project.toml b/Project.toml
@@ -9,6 +9,7 @@ FastMPOContractions = "f6e391d2-8ffa-4d7a-98cd-7e70024481ca"
 ITensorMPS = "0d1a4710-d33b-49a5-8f18-73bdf49b47e2"
 ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+ProjMPSs = "3cac4759-b6c4-45be-a543-65afe6a1360a"
 SparseIR = "4fe2279e-80f0-4adb-8463-ee114ff56b7d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
@@ -20,6 +21,7 @@ ITensors = "0.7"
 SparseIR = "^0.96, 0.97, 1"
 StaticArrays = "1"
 julia = "1"
+ProjMPSs = "0.4.1"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

diff --git a/src/Quantics.jl b/src/Quantics.jl
@@ -1,19 +1,12 @@
-#__precompile__(false)
-
 module Quantics
 
-#@everywhere begin
-#using Pkg
-#Pkg.activate(".")
-#Pkg.instantiate()
-#end
-
 using ITensors
 import ITensors
 using ITensors.SiteTypes: siteinds
 import ITensors.NDTensors: Tensor, BlockSparseTensor, blockview
 using ITensorMPS: MPS, MPO, AbstractMPS
-using ITensorMPS: findsite, linkinds, linkind
+using ITensorMPS: findsite, linkinds, linkind, findsites
+import ProjMPSs: ProjMPSs, ProjMPS, BlockedMPS, isprojectedat, project
 
 import SparseIR: Fermionic, Bosonic, Statistics
 import LinearAlgebra: I

diff --git a/src/mul.jl b/src/mul.jl
@@ -70,15 +70,17 @@ end
 """
 Convert an MPS tensor to an MPO tensor with a diagonal structure
 """
-function _asdiagonal(t, site::Index{T})::ITensor where {T<:Number}
-    hasinds(t, site') && error("Found $(site')")
-    links = uniqueinds(inds(t), site)
+function _asdiagonal(t, site::Index{T}; baseplev=0)::ITensor where {T<:Number}
+    ITensors.hasinds(t, site') && error("Found $(site')")
+    links = ITensors.uniqueinds(ITensors.inds(t), site)
     rawdata = Array(t, links..., site)
-    tensor = zeros(eltype(t), size(rawdata)..., dim(site))
-    for i in 1:dim(site)
+    tensor = zeros(eltype(t), size(rawdata)..., ITensors.dim(site))
+    for i in 1:ITensors.dim(site)
         tensor[.., i, i] = rawdata[.., i]
     end
-    return ITensor(tensor, links..., site', site)
+    return ITensor(
+        tensor, links..., ITensors.prime(site, baseplev + 1), ITensors.prime(site, baseplev)
+    )
 end
 
 function _todense(t, site::Index{T}) where {T<:Number}
@@ -161,3 +163,67 @@ function automul(M1::MPS, M2::MPS; tag_row::String="", tag_shared::String="",
 
     return asMPS(M)
 end
+
+
+"""
+By default, elementwise multiplication will be performed.
+"""
+function automul(
+    M1::BlockedMPS,
+    M2::BlockedMPS;
+    tag_row::String="",
+    tag_shared::String="",
+    tag_col::String="",
+    alg="naive",
+    maxdim=typemax(Int),
+    cutoff=1e-25,
+    kwargs...,
+)
+    all(length.(siteinds(M1)) .== 1) || error("M1 should have only 1 site index per site")
+    all(length.(siteinds(M2)) .== 1) || error("M2 should have only 1 site index per site")
+
+    sites_row = _findallsiteinds_by_tag(M1; tag=tag_row)
+    sites_shared = _findallsiteinds_by_tag(M1; tag=tag_shared)
+    sites_col = _findallsiteinds_by_tag(M2; tag=tag_col)
+    sites_matmul = Set(Iterators.flatten([sites_row, sites_shared, sites_col]))
+
+    sites1 = only.(siteinds(M1))
+    sites1_ewmul = setdiff(only.(siteinds(M1)), sites_matmul)
+    sites2_ewmul = setdiff(only.(siteinds(M2)), sites_matmul)
+    sites2_ewmul == sites1_ewmul || error("Invalid sites for elementwise multiplication")
+
+    M1 = _makesitediagonal(M1, sites1_ewmul; baseplev=1)
+    M2 = _makesitediagonal(M2, sites2_ewmul; baseplev=0)
+
+    sites_M1_diag = [collect(x) for x in siteinds(M1)]
+    sites_M2_diag = [collect(x) for x in siteinds(M2)]
+
+    M1 = rearrange_siteinds(M1, combinesites(sites_M1_diag, sites_row, sites_shared))
+
+    M2 = rearrange_siteinds(M2, combinesites(sites_M2_diag, sites_shared, sites_col))
+
+    M = ProjMPSs.contract(M1, M2; alg=alg, kwargs...)
+
+    M = extractdiagonal(M, sites1_ewmul)
+
+    ressites = Vector{eltype(siteinds(M1)[1])}[]
+    for s in siteinds(M)
+        s_ = unique(ITensors.noprime.(s))
+        if length(s_) == 1
+            push!(ressites, s_)
+        else
+            if s_[1] ∈ sites1
+                push!(ressites, [s_[1]])
+                push!(ressites, [s_[2]])
+            else
+                push!(ressites, [s_[2]])
+                push!(ressites, [s_[1]])
+            end
+        end
+    end
+    return ProjMPSs.truncate(rearrange_siteinds(M, ressites); cutoff=cutoff, maxdim=maxdim)
+end
+
+function _findallsiteinds_by_tag(M::BlockedMPS; tag=tag)
+    return Quantics.findallsiteinds_by_tag(only.(siteinds(M)); tag=tag)
+end
diff --git a/src/util.jl b/src/util.jl
@@ -364,6 +364,39 @@ function combinesites(M::MPO, site1::Index, site2::Index)
     return MPO(tensors)
 end
 
+
+function combinesites(
+    sites::Vector{Vector{Index{IndsT}}},
+    site1::AbstractVector{Index{IndsT}},
+    site2::AbstractVector{Index{IndsT}},
+) where {IndsT}
+    length(site1) == length(site2) || error("Length mismatch")
+    for (s1, s2) in zip(site1, site2)
+        sites = combinesites(sites, s1, s2)
+    end
+    return sites
+end
+
+function combinesites(
+    sites::Vector{Vector{Index{IndsT}}}, site1::Index, site2::Index
+) where {IndsT}
+    sites = deepcopy(sites)
+    p1 = findfirst(x -> x[1] == site1, sites)
+    p2 = findfirst(x -> x[1] == site2, sites)
+    if p1 === nothing || p2 === nothing
+        error("Site not found")
+    end
+    if abs(p1 - p2) != 1
+        error("Sites are not adjacent")
+    end
+    deleteat!(sites, min(p1, p2))
+    deleteat!(sites, min(p1, p2))
+    insert!(sites, min(p1, p2), [site1, site2])
+    return sites
+end
+
+
+
 function directprod(::Type{T}, sites, indices) where {T}
     length(sites) == length(indices) || error("Length mismatch between sites and indices")
     any(0 .== indices) && error("indices must be 1-based")
@@ -500,6 +533,100 @@ function makesitediagonal(M::AbstractMPS, tag::String)::MPS
     return MPS(collect(M_))
 end
 
+function _makesitediagonal(
+    projmps::ProjMPS, sites::AbstractVector{Index{IndsT}}; baseplev=0
+) where {IndsT}
+    M_ = deepcopy(MPO(collect(MPS(projmps))))
+    for site in sites
+        target_site::Int = only(findsites(M_, site))
+        M_[target_site] = _asdiagonal(M_[target_site], site; baseplev=baseplev)
+    end
+    return project(M_, projmps.projector)
+end
+
+function _makesitediagonal(projmps::ProjMPS, site::Index; baseplev=0)
+    return _makesitediagonal(projmps, [site]; baseplev=baseplev)
+end
+
+function makesitediagonal(projmps::ProjMPS, site::Index)
+    return _makesitediagonal(projmps, site; baseplev=0)
+end
+
+function makesitediagonal(projmps::ProjMPS, sites::AbstractVector{Index})
+    return _makesitediagonal(projmps, sites; baseplev=0)
+end
+
+function makesitediagonal(projmps::ProjMPS, tag::String)
+    mps_diagonal = Quantics.makesitediagonal(MPS(projmps), tag)
+    projmps_diagonal = ProjMPS(mps_diagonal)
+
+    target_sites = Quantics.findallsiteinds_by_tag(
+        unique(ITensors.noprime.(Iterators.flatten(siteinds(projmps)))); tag=tag
+    )
+
+    newproj = deepcopy(projmps.projector)
+    for s in target_sites
+        if isprojectedat(projmps.projector, s)
+            newproj[ITensors.prime(s)] = newproj[s]
+        end
+    end
+
+    return project(projmps_diagonal, newproj)
+end
+
+# FIXME: may be type unstable
+function _find_site_allplevs(tensor::ITensor, site::Index; maxplev=10)
+    ITensors.plev(site) == 0 || error("Site index must be unprimed.")
+    return [
+        ITensors.prime(site, plev) for
+        plev in 0:maxplev if ITensors.prime(site, plev) ∈ ITensors.inds(tensor)
+    ]
+end
+
+function extractdiagonal(
+    projmps::ProjMPS, sites::AbstractVector{Index{IndsT}}
+) where {IndsT}
+    tensors = collect(projmps.data)
+    for i in eachindex(tensors)
+        for site in intersect(sites, ITensors.inds(tensors[i]))
+            sitewithallplevs = _find_site_allplevs(tensors[i], site)
+            tensors[i] = if length(sitewithallplevs) > 1
+                tensors[i] = Quantics._extract_diagonal(tensors[i], sitewithallplevs...)
+            else
+                tensors[i]
+            end
+        end
+    end
+
+    projector = deepcopy(projmps.projector)
+    for site in sites
+        if site' in keys(projector.data)
+            delete!(projector.data, site')
+        end
+    end
+    return ProjMPS(MPS(tensors), projector)
+end
+
+function extractdiagonal(projmps::ProjMPS, site::Index{IndsT}) where {IndsT}
+    return Quantics.extractdiagonal(projmps, [site])
+end
+
+function extractdiagonal(projmps::ProjMPS, tag::String)::ProjMPS
+    targetsites = Quantics.findallsiteinds_by_tag(
+        unique(ITensors.noprime.(ProjMPSs._allsites(projmps))); tag=tag
+    )
+    return extractdiagonal(projmps, targetsites)
+end
+
+function rearrange_siteinds(projmps::ProjMPS, sites)
+    return ProjMPSs.rearrange_siteinds(projmps, sites)
+end
+
+function rearrange_siteinds(bmps::BlockedMPS, sites)
+    return ProjMPSs.rearrange_siteinds(bmps, sites)
+end
+
+
 """
 Extract diagonal components
 """
@@ -550,3 +677,28 @@ function _apply(A::MPO, Ψ::MPS; alg::String="fit", cutoff::Real=1e-25, kwargs..
     AΨ = noprime.(FastMPOContractions.contract_mpo_mpo(A, asMPO(Ψ); alg, kwargs...))
     MPS(collect(AΨ))
 end
+
+
+
+"""
+Make the BlockedMPS diagonal for a given site index `s` by introducing a dummy index `s'`.
+"""
+function makesitediagonal(obj::BlockedMPS, site)
+    return BlockedMPS([
+        _makesitediagonal(prjmps, site; baseplev=baseplev) for prjmps in values(obj)
+    ])
+end
+
+function _makesitediagonal(obj::BlockedMPS, site; baseplev=0)
+    return BlockedMPS([
+        _makesitediagonal(prjmps, site; baseplev=baseplev) for prjmps in values(obj)
+    ])
+end
+
+"""
+Extract diagonal of the BlockedMPS for `s`, `s'`, ... for a given site index `s`,
+where `s` must have a prime level of 0.
+"""
+function extractdiagonal(obj::BlockedMPS, site)
+    return BlockedMPS([extractdiagonal(prjmps, site) for prjmps in values(obj)])
+end
diff --git a/test/_util.jl b/test/_util.jl
@@ -0,0 +1,23 @@
+using ITensors
+import ITensorMPS: random_mps
+using Random
+
+function _random_mpo(
+    rng::AbstractRNG, sites::AbstractVector{<:AbstractVector{Index{T}}}; linkdims::Int=1
+) where {T}
+    sites_ = collect(Iterators.flatten(sites))
+    Ψ = random_mps(rng, sites_; linkdims)
+    tensors = ITensor[]
+    pos = 1
+    for i in 1:length(sites)
+        push!(tensors, prod(Ψ[pos:(pos + length(sites[i]) - 1)]))
+        pos += length(sites[i])
+    end
+    return MPO(tensors)
+end
+
+function _random_mpo(
+    sites::AbstractVector{<:AbstractVector{Index{T}}}; linkdims::Int=1
+) where {T}
+    return _random_mpo(Random.default_rng(), sites; linkdims)
+end