From 44baa13111171c2a99f2823d90a7c6e7c61c44d7 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Tue, 20 Jul 2021 12:03:35 -0500 Subject: [PATCH 1/5] Add rrule for split_links --- src/ITensorNetworks/itensor_network.jl | 8 +++++--- src/Optimizations/Optimizations.jl | 1 + src/Optimizations/itensor_network.jl | 26 ++++++++++++++++++++++++++ 3 files changed, 32 insertions(+), 3 deletions(-) create mode 100644 src/Optimizations/itensor_network.jl diff --git a/src/ITensorNetworks/itensor_network.jl b/src/ITensorNetworks/itensor_network.jl index 04af495..7ac1dba 100644 --- a/src/ITensorNetworks/itensor_network.jl +++ b/src/ITensorNetworks/itensor_network.jl @@ -190,7 +190,9 @@ end # Split the links of an ITensor network by changing the prime levels # or tags of pairs of links. -function split_links(H::Union{MPS,MPO}; split_tags=("" => ""), split_plevs=(0 => 1)) +function split_links( + H::Union{MPS,MPO}; split_tags=("" => ""), split_plevs=(0 => 1), tag_f=addtags +) left_tags, right_tags = split_tags left_plev, right_plev = split_plevs l = outlinkinds(H) @@ -198,8 +200,8 @@ function split_links(H::Union{MPS,MPO}; split_tags=("" => ""), split_plevs=(0 => for bond in keys(l) n1, n2 = bond lₙ = l[bond] - left_l_n = setprime(addtags(lₙ, left_tags), left_plev) - right_l_n = setprime(addtags(lₙ, right_tags), right_plev) + left_l_n = prime(tag_f(lₙ, left_tags), left_plev) + right_l_n = prime(tag_f(lₙ, right_tags), right_plev) Hsplit[n1] = replaceinds(Hsplit[n1], lₙ => left_l_n) Hsplit[n2] = replaceinds(Hsplit[n2], lₙ => right_l_n) end diff --git a/src/Optimizations/Optimizations.jl b/src/Optimizations/Optimizations.jl index eebcb15..2814dd7 100644 --- a/src/Optimizations/Optimizations.jl +++ b/src/Optimizations/Optimizations.jl @@ -5,6 +5,7 @@ using ITensors export gradient_descent, generate_inner_network include("peps.jl") +include("itensor_network.jl") include("run.jl") include("optimizers.jl") diff --git a/src/Optimizations/itensor_network.jl b/src/Optimizations/itensor_network.jl new file mode 100644 index 0000000..dad51fc --- /dev/null +++ b/src/Optimizations/itensor_network.jl @@ -0,0 +1,26 @@ +using ChainRulesCore +using ..ITensorNetworks +using ..ITensorNetworks: split_links + +inv_op(::typeof(addtags)) = :removetags +inv_op(::typeof(removetags)) = :addtags + +function ChainRulesCore.rrule( + ::typeof(split_links), + H::Union{MPS,MPO}; + split_tags=("" => ""), + split_plevs=(0 => 1), + tag_f=addtags, +) + function pullback(dHsplit) + dH = split_links( + dHsplit; + split_tags=split_tags, + split_plevs=(split_plevs[2] => split_plevs[1]), + tag_f=inv_op(tag_f), + ) + return (NoTangent(), dH, NoTangent(), NoTangent(), NoTangent()) + end + return split_links(H; split_tags=split_tags, split_plevs=split_plevs, tag_f=tag_f), + pullback +end From 9bb1b9f1eeecee8a5587223284b4eb5c4280b119 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 22 Jul 2021 11:30:37 -0500 Subject: [PATCH 2/5] Insert projectors to PEPS optimizations --- src/ITensorNetworks/boundary_mps.jl | 30 ++++++------ src/ITensorNetworks/itensor_network.jl | 4 ++ src/ITensorNetworks/peps.jl | 31 +++++++++++++ src/Optimizations/Optimizations.jl | 2 +- src/Optimizations/itensor_network.jl | 4 +- src/Optimizations/peps.jl | 64 +++++++++++++++++++++++++- test/Optimizations/runtests.jl | 35 +++++++++++++- 7 files changed, 148 insertions(+), 22 deletions(-) diff --git a/src/ITensorNetworks/boundary_mps.jl b/src/ITensorNetworks/boundary_mps.jl index 1694d91..0aceb72 100644 --- a/src/ITensorNetworks/boundary_mps.jl +++ b/src/ITensorNetworks/boundary_mps.jl @@ -253,6 +253,19 @@ function insert_projectors(tn, boundary_mps::BoundaryMPS; center, projector_cent return tn_split, projectors_left, projectors_right end +function insert_projectors(tn::Matrix{ITensor}; center, cutoff, maxdim) + #TODO + # Contract in every direction + combiner_gauge = combiners(linkinds, tn) + tnᶜ = insert_gauge(tn, combiner_gauge) + boundary_mpsᶜ = contract_approx(tnᶜ; maxdim=maxdim, cutoff=cutoff) + + tn_cacheᶜ = contraction_cache(tnᶜ, boundary_mpsᶜ) + tn_cache = insert_gauge(tn_cacheᶜ, combiner_gauge) + _boundary_mps = boundary_mps(tn_cache) + return insert_projectors(tn, _boundary_mps; center=center) +end + function contraction_cache_top(tn, boundary_mps::Vector{MPS}, n) tn_cache = fill(ITensor(1.0), size(tn)) for nrow in 1:size(tn, 1) @@ -359,27 +372,12 @@ function sqnorm_approx(ψ::Matrix{ITensor}; center, cutoff, maxdim) ψ′ = addtags(linkinds, ψ, "ket") # TODO: implement contract(commoninds, ψ′, ψᴴ) tn = ψ′ .* ψᴴ - - # Contract in every direction - combiner_gauge = combiners(linkinds, tn) - tnᶜ = insert_gauge(tn, combiner_gauge) - boundary_mpsᶜ = contract_approx(tnᶜ; maxdim=maxdim, cutoff=cutoff) - - tn_cacheᶜ = contraction_cache(tnᶜ, boundary_mpsᶜ) - tn_cache = insert_gauge(tn_cacheᶜ, combiner_gauge) - _boundary_mps = boundary_mps(tn_cache) - - # # Insert projectors horizontally (to measure e.g. properties # in a row of the network) - # - - tn_projected = insert_projectors(tn, _boundary_mps; center=center) + tn_projected = insert_projectors(tn; center=center, cutoff=cutoff, maxdim=maxdim) tn_split, Pl, Pr = tn_projected - ψᴴ_split = split_network(ψᴴ) ψ′_split = split_network(ψ′) - Pl_flat = reduce(vcat, Pl) Pr_flat = reduce(vcat, Pr) return mapreduce(vec, vcat, (ψᴴ_split, ψ′_split, Pl_flat, Pr_flat)) diff --git a/src/ITensorNetworks/itensor_network.jl b/src/ITensorNetworks/itensor_network.jl index 7ac1dba..06ac641 100644 --- a/src/ITensorNetworks/itensor_network.jl +++ b/src/ITensorNetworks/itensor_network.jl @@ -150,6 +150,10 @@ function ITensors.addtags(::typeof(linkinds), tn, args...) return mapinds(x -> addtags(x, args...), linkinds, tn) end +function ITensors.removetags(::typeof(linkinds), tn, args...) + return mapinds(x -> removetags(x, args...), linkinds, tn) +end + # Compute the sets of combiners that combine the link indices # of the tensor network so that neighboring tensors only # share a single larger index. diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index e42627a..62f15bc 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -81,10 +81,32 @@ broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> v[], +, A.data .* B.data) ITensors.prime(P::PEPS, n::Integer=1) = PEPS(map(x -> prime(x, n), P.data)) +function ITensors.prime(indices::Array{<:Index,1}, P::PEPS, n::Integer=1) + function primeinds(tensor) + prime_inds = [ind for ind in inds(tensor) if ind in indices] + return replaceinds(tensor, prime_inds => prime(prime_inds, n)) + end + return PEPS(map(x -> primeinds(x), P.data)) +end + function ITensors.prime(::typeof(linkinds), P::PEPS, n::Integer=1) return PEPS(mapinds(x -> prime(x, n), linkinds, P.data)) end +function ITensors.addtags(::typeof(linkinds), P::PEPS, args...) + return PEPS(addtags(linkinds, P.data, args...)) +end + +function ITensors.removetags(::typeof(linkinds), P::PEPS, args...) + return PEPS(removetags(linkinds, P.data, args...)) +end + +split_network(P::PEPS) = PEPS(split_network(P.data)) + +function ITensors.commoninds(p1::PEPS, p2::PEPS) + return mapreduce(a -> commoninds(a...), vcat, zip(p1.data, p2.data)) +end + # Get the tensor network of function inner_network(peps::PEPS, peps_prime::PEPS) return vcat(vcat(peps.data...), vcat(peps_prime.data...)) @@ -117,3 +139,12 @@ function flatten(v::Array{<:PEPS}) tensor_list = [vcat(peps.data...) for peps in v] return vcat(tensor_list...) end + +function insert_projectors(peps::PEPS, center, cutoff=1e-15, maxdim=100) + # Square the tensor network + psi_bra = addtags(linkinds, dag.(peps.data), "bra") + psi_ket = addtags(linkinds, peps.data, "ket") + tn = psi_bra .* psi_ket + tn_split, pl, pr = insert_projectors(tn; center=center, cutoff=cutoff, maxdim=maxdim) + return tn_split, vcat(reduce(vcat, pl), reduce(vcat, pr)) +end diff --git a/src/Optimizations/Optimizations.jl b/src/Optimizations/Optimizations.jl index 2814dd7..cdfa315 100644 --- a/src/Optimizations/Optimizations.jl +++ b/src/Optimizations/Optimizations.jl @@ -2,7 +2,7 @@ module Optimizations using ITensors -export gradient_descent, generate_inner_network +export gradient_descent, generate_inner_network, rayleigh_quotient include("peps.jl") include("itensor_network.jl") diff --git a/src/Optimizations/itensor_network.jl b/src/Optimizations/itensor_network.jl index dad51fc..6a4640a 100644 --- a/src/Optimizations/itensor_network.jl +++ b/src/Optimizations/itensor_network.jl @@ -2,8 +2,8 @@ using ChainRulesCore using ..ITensorNetworks using ..ITensorNetworks: split_links -inv_op(::typeof(addtags)) = :removetags -inv_op(::typeof(removetags)) = :addtags +inv_op(::typeof(addtags)) = removetags +inv_op(::typeof(removetags)) = addtags function ChainRulesCore.rrule( ::typeof(split_links), diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index 51cc621..96896f2 100644 --- a/src/Optimizations/peps.jl +++ b/src/Optimizations/peps.jl @@ -2,7 +2,7 @@ using AutoHOOT, ChainRulesCore, Zygote using ..ITensorAutoHOOT using ..ITensorNetworks using ITensors: setinds -using ..ITensorNetworks: PEPS, inner_network, flatten +using ..ITensorNetworks: PEPS, inner_network, flatten, insert_projectors, split_network using ..ITensorAutoHOOT: batch_tensor_contraction function ChainRulesCore.rrule(::typeof(PEPS), data::Matrix{ITensor}) @@ -13,6 +13,26 @@ function ChainRulesCore.rrule(::typeof(ITensors.prime), P::PEPS, n::Integer=1) return prime(P, n), dprime -> (NoTangent(), prime(dprime, -n), NoTangent()) end +function ChainRulesCore.rrule( + ::typeof(ITensors.addtags), ::typeof(linkinds), P::PEPS, args... +) + function pullback(dtag_peps) + dP = ITensors.removetags(linkinds, P, args...) + return (NoTangent(), NoTangent(), dP, broadcast_notangent(args)...) + end + return ITensors.addtags(linkinds, P, args...), pullback +end + +function ChainRulesCore.rrule( + ::typeof(ITensors.removetags), ::typeof(linkinds), P::PEPS, args... +) + function pullback(dtag_peps) + dP = ITensors.addtags(linkinds, P, args...) + return (NoTangent(), NoTangent(), dP, broadcast_notangent(args)...) + end + return ITensors.removetags(linkinds, P, args...), pullback +end + function ChainRulesCore.rrule( ::typeof(ITensors.prime), ::typeof(linkinds), P::PEPS, n::Integer=1 ) @@ -68,6 +88,21 @@ end peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, Hlocal::Array ) +function generate_inner_network( + peps::PEPS, + peps_prime::PEPS, + peps_prime_ham::PEPS, + projectors::Array{<:ITensor,1}, + Hlocal::Array, +) + network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, Hlocal) + return map(network -> vcat(network, projectors), network_list) +end + +@non_differentiable generate_inner_network( + peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, projectors::Array, Hlocal::Array +) + function rayleigh_quotient(inners::Array) self_inner = inners[length(inners)][] expectations = sum(inners[1:(length(inners) - 1)])[] @@ -86,3 +121,30 @@ function loss_grad_wrap(peps::PEPS, Hlocal::Array) loss_w_grad(peps::PEPS) = loss(peps), gradient(loss, peps)[1] return loss_w_grad end + +@non_differentiable insert_projectors(peps::PEPS, center) + +@non_differentiable ITensors.commoninds(p1::PEPS, p2::PEPS) + +function loss_grad_wrap(peps::PEPS, Hlocal::Array, ::typeof(insert_projectors)) + function loss(peps::PEPS) + # insert projectors here + center = (div(size(peps.data)[1], 2), :) + _, projectors = insert_projectors(peps, center) + peps_bra = addtags(linkinds, peps, "bra") + peps_ket = addtags(linkinds, peps, "ket") + sites = commoninds(peps_bra, peps_ket) + peps_bra_split = split_network(peps_bra) #TODO: implement rrule + peps_ket_split = split_network(peps_ket) + peps_ket_split_ham = prime(sites, peps_ket_split) #TODO: implement rrule + # generate network + network_list = generate_inner_network( + peps_bra_split, peps_ket_split, peps_ket_split_ham, projectors, Hlocal + ) + variables = flatten([peps_bra_split, peps_ket_split, peps_ket_split_ham]) + inners = batch_tensor_contraction(network_list, variables...) + return rayleigh_quotient(inners) + end + loss_w_grad(peps::PEPS) = loss(peps), gradient(loss, peps)[1] + return loss_w_grad +end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index d400e6f..1a19be9 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -1,6 +1,8 @@ using ITensors, ITensorNetworkAD, AutoHOOT, Zygote, OptimKit -using ITensorNetworkAD.ITensorNetworks: PEPS, inner_network, Models, flatten -using ITensorNetworkAD.Optimizations: gradient_descent, generate_inner_network +using ITensorNetworkAD.ITensorNetworks: + PEPS, inner_network, Models, flatten, insert_projectors, split_network +using ITensorNetworkAD.Optimizations: + gradient_descent, generate_inner_network, rayleigh_quotient using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction @testset "test monotonic loss decrease of optimization" begin @@ -42,3 +44,32 @@ end g_true_first_site = 2 * g_true_first_site @test isapprox(g[1].data[1, 1], g_true_first_site) end + +@testset "test approximate rayleigh quotient objective" begin + Nx = 3 + Ny = 3 + sites = siteinds("S=1/2", Ny, Nx) + Hlocal = Models.localham(Models.Model("tfim"), sites; h=1.0) + peps = PEPS(sites; linkdims=2) + randn!(peps) + function loss(peps::PEPS) + # insert projectors here + center = (div(size(peps.data)[1] - 1, 2) + 1, :) + _, projectors = insert_projectors(peps, center) + peps_bra = addtags(linkinds, peps, "bra") + peps_ket = addtags(linkinds, peps, "ket") + sites = commoninds(peps_bra, peps_ket) + peps_bra_split = split_network(peps_bra) + peps_ket_split = split_network(peps_ket) + peps_ket_split_ham = prime(sites, peps_ket_split) + # generate network + network_list = generate_inner_network( + peps_bra_split, peps_ket_split, peps_ket_split_ham, projectors, Hlocal + ) + variables = flatten([peps_bra_split, peps_ket_split, peps_ket_split_ham]) + inners = batch_tensor_contraction(network_list, variables...) + return rayleigh_quotient(inners) + end + out = loss(peps) + @assert size(out) == () +end From 2c8aa86a4dd300688f711c68f6a59fd75f31b207 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 25 Jul 2021 22:30:27 -0500 Subject: [PATCH 3/5] Get gradient optimization with approximation peps contraction working --- src/ITensorNetworks/peps.jl | 6 ++- src/Optimizations/itensor_network.jl | 38 ++++++++------ src/Optimizations/peps.jl | 39 ++++++++++----- src/Optimizations/run.jl | 44 ++++++++++++----- test/Optimizations/runtests.jl | 74 ++++++++++++++++++++++------ 5 files changed, 147 insertions(+), 54 deletions(-) diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index 62f15bc..0e6d087 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -81,6 +81,7 @@ broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> v[], +, A.data .* B.data) ITensors.prime(P::PEPS, n::Integer=1) = PEPS(map(x -> prime(x, n), P.data)) +# prime a PEPS with specified indices function ITensors.prime(indices::Array{<:Index,1}, P::PEPS, n::Integer=1) function primeinds(tensor) prime_inds = [ind for ind in inds(tensor) if ind in indices] @@ -89,6 +90,7 @@ function ITensors.prime(indices::Array{<:Index,1}, P::PEPS, n::Integer=1) return PEPS(map(x -> primeinds(x), P.data)) end +# prime linkinds of a PEPS function ITensors.prime(::typeof(linkinds), P::PEPS, n::Integer=1) return PEPS(mapinds(x -> prime(x, n), linkinds, P.data)) end @@ -101,7 +103,9 @@ function ITensors.removetags(::typeof(linkinds), P::PEPS, args...) return PEPS(removetags(linkinds, P.data, args...)) end -split_network(P::PEPS) = PEPS(split_network(P.data)) +ITensors.data(P::PEPS) = P.data + +split_network(P::PEPS) = PEPS(split_network(data(P))) function ITensors.commoninds(p1::PEPS, p2::PEPS) return mapreduce(a -> commoninds(a...), vcat, zip(p1.data, p2.data)) diff --git a/src/Optimizations/itensor_network.jl b/src/Optimizations/itensor_network.jl index 6a4640a..627d6b4 100644 --- a/src/Optimizations/itensor_network.jl +++ b/src/Optimizations/itensor_network.jl @@ -1,26 +1,32 @@ using ChainRulesCore using ..ITensorNetworks -using ..ITensorNetworks: split_links +using ..ITensorNetworks: split_network, default_projector_center inv_op(::typeof(addtags)) = removetags inv_op(::typeof(removetags)) = addtags function ChainRulesCore.rrule( - ::typeof(split_links), - H::Union{MPS,MPO}; - split_tags=("" => ""), - split_plevs=(0 => 1), - tag_f=addtags, + ::typeof(split_network), + tn::Matrix{ITensor}; + projector_center=default_projector_center(tn), ) - function pullback(dHsplit) - dH = split_links( - dHsplit; - split_tags=split_tags, - split_plevs=(split_plevs[2] => split_plevs[1]), - tag_f=inv_op(tag_f), - ) - return (NoTangent(), dH, NoTangent(), NoTangent(), NoTangent()) + dimy, dimx = size(tn) + tn_vec = vec(tn) + function pullback(dtn_split::Matrix{ITensor}) + dtn_split_vec = vec(dtn_split) + dtn_vec = [] + for i in 1:(dimy * dimx) + indices = inds(tn_vec[i]) + indices_reorder = [] + for i_split in inds(dtn_split_vec[i]) + index = findall(x -> x.id == i_split.id, indices) + @assert(length(index) == 1) + push!(indices_reorder, indices[index[1]]) + end + push!(dtn_vec, setinds(dtn_split_vec[i], Tuple(indices_reorder))) + end + dtn = reshape(dtn_vec, (dimy, dimx)) + return (NoTangent(), dtn, NoTangent()) end - return split_links(H; split_tags=split_tags, split_plevs=split_plevs, tag_f=tag_f), - pullback + return split_network(tn; projector_center=projector_center), pullback end diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index 96896f2..05b9efa 100644 --- a/src/Optimizations/peps.jl +++ b/src/Optimizations/peps.jl @@ -5,6 +5,12 @@ using ITensors: setinds using ..ITensorNetworks: PEPS, inner_network, flatten, insert_projectors, split_network using ..ITensorAutoHOOT: batch_tensor_contraction +broadcast_notangent(a) = broadcast(_ -> NoTangent(), a) + +function ChainRulesCore.rrule(::typeof(ITensors.data), P::PEPS) + return P.data, d_data -> (NoTangent(), PEPS(d_data)) +end + function ChainRulesCore.rrule(::typeof(PEPS), data::Matrix{ITensor}) return PEPS(data), dpeps -> (NoTangent(), dpeps.data) end @@ -17,7 +23,7 @@ function ChainRulesCore.rrule( ::typeof(ITensors.addtags), ::typeof(linkinds), P::PEPS, args... ) function pullback(dtag_peps) - dP = ITensors.removetags(linkinds, P, args...) + dP = ITensors.removetags(linkinds, dtag_peps, args...) return (NoTangent(), NoTangent(), dP, broadcast_notangent(args)...) end return ITensors.addtags(linkinds, P, args...), pullback @@ -27,7 +33,7 @@ function ChainRulesCore.rrule( ::typeof(ITensors.removetags), ::typeof(linkinds), P::PEPS, args... ) function pullback(dtag_peps) - dP = ITensors.addtags(linkinds, P, args...) + dP = ITensors.addtags(linkinds, dtag_peps, args...) return (NoTangent(), NoTangent(), dP, broadcast_notangent(args)...) end return ITensors.removetags(linkinds, P, args...), pullback @@ -40,6 +46,14 @@ function ChainRulesCore.rrule( dprime -> (NoTangent(), NoTangent(), prime(linkinds, dprime, -n), NoTangent()) end +function ChainRulesCore.rrule( + ::typeof(ITensors.prime), indices::Array{<:Index,1}, P::PEPS, n::Integer=1 +) + primeinds = [prime(ind, n) for ind in indices] + return prime(indices, P, n), + dprime -> (NoTangent(), NoTangent(), prime(primeinds, dprime, -n), NoTangent()) +end + function ChainRulesCore.rrule(::typeof(flatten), v::Array{<:PEPS}) size_list = [size(peps.data) for peps in v] function adjoint_pullback(dt) @@ -100,12 +114,16 @@ function generate_inner_network( end @non_differentiable generate_inner_network( - peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, projectors::Array, Hlocal::Array + peps::PEPS, + peps_prime::PEPS, + peps_prime_ham::PEPS, + projectors::Array{<:ITensor,1}, + Hlocal::Array, ) function rayleigh_quotient(inners::Array) - self_inner = inners[length(inners)][] - expectations = sum(inners[1:(length(inners) - 1)])[] + self_inner = inners[end][] + expectations = sum(inners[1:(end - 1)])[] return expectations / self_inner end @@ -122,21 +140,20 @@ function loss_grad_wrap(peps::PEPS, Hlocal::Array) return loss_w_grad end -@non_differentiable insert_projectors(peps::PEPS, center) +@non_differentiable insert_projectors(peps::PEPS, center, cutoff, maxdim) @non_differentiable ITensors.commoninds(p1::PEPS, p2::PEPS) function loss_grad_wrap(peps::PEPS, Hlocal::Array, ::typeof(insert_projectors)) + center = (div(size(peps.data)[1] - 1, 2) + 1, :) function loss(peps::PEPS) - # insert projectors here - center = (div(size(peps.data)[1], 2), :) - _, projectors = insert_projectors(peps, center) + tn_split, projectors = insert_projectors(peps, center) peps_bra = addtags(linkinds, peps, "bra") peps_ket = addtags(linkinds, peps, "ket") sites = commoninds(peps_bra, peps_ket) - peps_bra_split = split_network(peps_bra) #TODO: implement rrule + peps_bra_split = split_network(peps_bra) peps_ket_split = split_network(peps_ket) - peps_ket_split_ham = prime(sites, peps_ket_split) #TODO: implement rrule + peps_ket_split_ham = prime(sites, peps_ket_split) # generate network network_list = generate_inner_network( peps_bra_split, peps_ket_split, peps_ket_split_ham, projectors, Hlocal diff --git a/src/Optimizations/run.jl b/src/Optimizations/run.jl index 401aa85..bc93e8d 100644 --- a/src/Optimizations/run.jl +++ b/src/Optimizations/run.jl @@ -2,6 +2,18 @@ using OptimKit using ..ITensorNetworks using ..ITensorNetworks: broadcast_add, broadcast_minus, broadcast_mul, broadcast_inner +function gradient_descent(peps::PEPS, loss_w_grad; stepsize::Float64, num_sweeps::Int) + # gradient descent iterations + losses = [] + for iter in 1:num_sweeps + l, g = loss_w_grad(peps) + print("The rayleigh quotient at iteraton $iter is $l\n") + peps = broadcast_minus(peps, broadcast_mul(stepsize, g)) + push!(losses, l) + end + return losses +end + """Update PEPS based on gradient descent Parameters ---------- @@ -15,21 +27,19 @@ An array containing Rayleigh quotient losses after each iteration. """ function gradient_descent(peps::PEPS, Hlocal::Array; stepsize::Float64, num_sweeps::Int) loss_w_grad = loss_grad_wrap(peps, Hlocal) - # gradient descent iterations - losses = [] - for iter in 1:num_sweeps - l, g = loss_w_grad(peps) - print("The rayleigh quotient at iteraton $iter is $l\n") - peps = broadcast_minus(peps, broadcast_mul(stepsize, g)) - push!(losses, l) - end - return losses + return gradient_descent(peps, loss_w_grad; stepsize=stepsize, num_sweeps=num_sweeps) end -function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") +function gradient_descent( + peps::PEPS, Hlocal::Array, ::typeof(insert_projectors); stepsize::Float64, num_sweeps::Int +) + loss_w_grad = loss_grad_wrap(peps, Hlocal, insert_projectors) + return gradient_descent(peps, loss_w_grad; stepsize=stepsize, num_sweeps=num_sweeps) +end + +function OptimKit.optimize(peps::PEPS, loss_w_grad; num_sweeps::Int, method="GD") @assert(method in ["GD", "LBFGS", "CG"]) inner(x, peps1, peps2) = broadcast_inner(peps1, peps2) - loss_w_grad = loss_grad_wrap(peps, Hlocal) scale(peps, alpha) = broadcast_mul(alpha, peps) add(peps1, peps2, alpha) = broadcast_add(peps1, broadcast_mul(alpha, peps2)) retract(peps1, peps2, alpha) = (add(peps1, peps2, alpha), peps2) @@ -48,3 +58,15 @@ function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="G ) return history[:, 1] end + +function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") + loss_w_grad = loss_grad_wrap(peps, Hlocal) + return optimize(peps, loss_w_grad; num_sweeps=num_sweeps, method=method) +end + +function OptimKit.optimize( + peps::PEPS, Hlocal::Array, ::typeof(insert_projectors); num_sweeps::Int, method="GD" +) + loss_w_grad = loss_grad_wrap(peps, Hlocal, insert_projectors) + return optimize(peps, loss_w_grad; num_sweeps=num_sweeps, method=method) +end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index 1a19be9..5116b79 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -24,6 +24,29 @@ using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction end end +@testset "test monotonic loss decrease of optimization with inserting projectors" begin + Nx, Ny = 3, 3 + num_sweeps = 20 + sites = siteinds("S=1/2", Ny, Nx) + peps = PEPS(sites; linkdims=2) + randn!(peps) + H_local = Models.localham(Models.Model("tfim"), sites; h=1.0) + losses_gd = gradient_descent( + peps, H_local, insert_projectors; stepsize=0.005, num_sweeps=num_sweeps + ) + losses_ls = optimize(peps, H_local, insert_projectors; num_sweeps=num_sweeps, method="GD") + losses_lbfgs = optimize( + peps, H_local, insert_projectors; num_sweeps=num_sweeps, method="LBFGS" + ) + losses_cg = optimize(peps, H_local, insert_projectors; num_sweeps=num_sweeps, method="CG") + for i in 3:(length(losses_gd) - 1) + @test losses_gd[i] >= losses_gd[i + 1] + @test losses_ls[i] >= losses_ls[i + 1] + @test losses_lbfgs[i] >= losses_lbfgs[i + 1] + @test losses_cg[i] >= losses_cg[i + 1] + end +end + @testset "test inner product gradient" begin Nx = 2 Ny = 2 @@ -45,31 +68,52 @@ end @test isapprox(g[1].data[1, 1], g_true_first_site) end -@testset "test approximate rayleigh quotient objective" begin - Nx = 3 - Ny = 3 +@testset "test split network" begin + Nx, Ny = 3, 3 sites = siteinds("S=1/2", Ny, Nx) - Hlocal = Models.localham(Models.Model("tfim"), sites; h=1.0) peps = PEPS(sites; linkdims=2) randn!(peps) + center = (div(size(peps.data)[1] - 1, 2) + 1, :) function loss(peps::PEPS) - # insert projectors here - center = (div(size(peps.data)[1] - 1, 2) + 1, :) - _, projectors = insert_projectors(peps, center) + tn_split, projectors = insert_projectors(peps, center) peps_bra = addtags(linkinds, peps, "bra") peps_ket = addtags(linkinds, peps, "ket") - sites = commoninds(peps_bra, peps_ket) peps_bra_split = split_network(peps_bra) peps_ket_split = split_network(peps_ket) - peps_ket_split_ham = prime(sites, peps_ket_split) - # generate network network_list = generate_inner_network( - peps_bra_split, peps_ket_split, peps_ket_split_ham, projectors, Hlocal + peps_bra_split, peps_ket_split, peps_ket_split, projectors, [] ) - variables = flatten([peps_bra_split, peps_ket_split, peps_ket_split_ham]) + variables = flatten([peps_bra_split, peps_ket_split]) inners = batch_tensor_contraction(network_list, variables...) - return rayleigh_quotient(inners) + return sum(inners)[] end - out = loss(peps) - @assert size(out) == () + g = gradient(loss, peps) + inner = inner_network(peps, prime(linkinds, peps)) + g_true_first_site = contract(inner[2:length(inner)]) + g_true_first_site = 2 * g_true_first_site + @test isapprox(g[1].data[1, 1], g_true_first_site) +end + +@testset "test inner product gradient with tagging" begin + Nx, Ny = 3, 3 + sites = siteinds("S=1/2", Ny, Nx) + peps = PEPS(sites; linkdims=2) + randn!(peps) + function loss(peps::PEPS) + peps_bra = addtags(linkinds, peps, "bra") + peps_ket = addtags(linkinds, peps, "ket") + sites = commoninds(peps_bra, peps_ket) + peps_ket_ham = prime(sites, peps_ket) + projectors = [ITensor(1.0)] + network_list = generate_inner_network(peps_bra, peps_ket, peps_ket_ham, projectors, []) + variables = flatten([peps_bra, peps_ket]) + inners = batch_tensor_contraction(network_list, variables...) + return sum(inners)[] + end + g = gradient(loss, peps) + inner = inner_network(peps, prime(linkinds, peps)) + g_true_first_site = contract(inner[2:length(inner)]) + g_true_first_site = 2 * g_true_first_site + @test isapprox(g[1].data[1, 1], g_true_first_site) + @test isapprox(loss(peps), contract(inner)[]) end From a4cbdec0546fb485609a556e7b7a2d75fa2daf3e Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Sun, 25 Jul 2021 23:18:06 -0500 Subject: [PATCH 4/5] Rewrite rrule for split_network --- src/ITensorNetworks/itensor_network.jl | 8 +++----- src/Optimizations/itensor_network.jl | 23 ++++++++--------------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/src/ITensorNetworks/itensor_network.jl b/src/ITensorNetworks/itensor_network.jl index 06ac641..fd9af75 100644 --- a/src/ITensorNetworks/itensor_network.jl +++ b/src/ITensorNetworks/itensor_network.jl @@ -194,9 +194,7 @@ end # Split the links of an ITensor network by changing the prime levels # or tags of pairs of links. -function split_links( - H::Union{MPS,MPO}; split_tags=("" => ""), split_plevs=(0 => 1), tag_f=addtags -) +function split_links(H::Union{MPS,MPO}; split_tags=("" => ""), split_plevs=(0 => 1)) left_tags, right_tags = split_tags left_plev, right_plev = split_plevs l = outlinkinds(H) @@ -204,8 +202,8 @@ function split_links( for bond in keys(l) n1, n2 = bond lₙ = l[bond] - left_l_n = prime(tag_f(lₙ, left_tags), left_plev) - right_l_n = prime(tag_f(lₙ, right_tags), right_plev) + left_l_n = prime(addtags(lₙ, left_tags), left_plev) + right_l_n = prime(addtags(lₙ, right_tags), right_plev) Hsplit[n1] = replaceinds(Hsplit[n1], lₙ => left_l_n) Hsplit[n2] = replaceinds(Hsplit[n2], lₙ => right_l_n) end diff --git a/src/Optimizations/itensor_network.jl b/src/Optimizations/itensor_network.jl index 627d6b4..7ba559a 100644 --- a/src/Optimizations/itensor_network.jl +++ b/src/Optimizations/itensor_network.jl @@ -2,30 +2,23 @@ using ChainRulesCore using ..ITensorNetworks using ..ITensorNetworks: split_network, default_projector_center -inv_op(::typeof(addtags)) = removetags -inv_op(::typeof(removetags)) = addtags - function ChainRulesCore.rrule( ::typeof(split_network), tn::Matrix{ITensor}; projector_center=default_projector_center(tn), ) dimy, dimx = size(tn) - tn_vec = vec(tn) function pullback(dtn_split::Matrix{ITensor}) - dtn_split_vec = vec(dtn_split) - dtn_vec = [] - for i in 1:(dimy * dimx) - indices = inds(tn_vec[i]) - indices_reorder = [] - for i_split in inds(dtn_split_vec[i]) - index = findall(x -> x.id == i_split.id, indices) - @assert(length(index) == 1) - push!(indices_reorder, indices[index[1]]) + dtn = copy(dtn_split) + for ii in 1:dimy + for jj in 1:dimx + dt, t = dtn[ii, jj], tn[ii, jj] + indices = inds(t) + get_index(i_split) = findfirst(x -> x.id == i_split.id, indices) + indices_reorder = [indices[get_index(i_split)] for i_split in inds(dt)] + dtn[ii, jj] = setinds(dt, Tuple(indices_reorder)) end - push!(dtn_vec, setinds(dtn_split_vec[i], Tuple(indices_reorder))) end - dtn = reshape(dtn_vec, (dimy, dimx)) return (NoTangent(), dtn, NoTangent()) end return split_network(tn; projector_center=projector_center), pullback From 3eee8e333d5e535ae5f6abbb887fd65c9a614d30 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 26 Jul 2021 11:07:54 -0500 Subject: [PATCH 5/5] Refactor bmps out from split_network, rewrite rrule for split_network --- src/ITensorNetworks/boundary_mps.jl | 27 +++++++++++++-------------- src/ITensorNetworks/peps.jl | 3 ++- src/Optimizations/itensor_network.jl | 12 +----------- 3 files changed, 16 insertions(+), 26 deletions(-) diff --git a/src/ITensorNetworks/boundary_mps.jl b/src/ITensorNetworks/boundary_mps.jl index 0aceb72..5aa2b38 100644 --- a/src/ITensorNetworks/boundary_mps.jl +++ b/src/ITensorNetworks/boundary_mps.jl @@ -253,19 +253,6 @@ function insert_projectors(tn, boundary_mps::BoundaryMPS; center, projector_cent return tn_split, projectors_left, projectors_right end -function insert_projectors(tn::Matrix{ITensor}; center, cutoff, maxdim) - #TODO - # Contract in every direction - combiner_gauge = combiners(linkinds, tn) - tnᶜ = insert_gauge(tn, combiner_gauge) - boundary_mpsᶜ = contract_approx(tnᶜ; maxdim=maxdim, cutoff=cutoff) - - tn_cacheᶜ = contraction_cache(tnᶜ, boundary_mpsᶜ) - tn_cache = insert_gauge(tn_cacheᶜ, combiner_gauge) - _boundary_mps = boundary_mps(tn_cache) - return insert_projectors(tn, _boundary_mps; center=center) -end - function contraction_cache_top(tn, boundary_mps::Vector{MPS}, n) tn_cache = fill(ITensor(1.0), size(tn)) for nrow in 1:size(tn, 1) @@ -352,6 +339,17 @@ function boundary_mps(tn::NamedTuple) ) end +function boundary_mps(tn::Matrix{ITensor}; cutoff, maxdim) + #TODO + # Contract in every direction + combiner_gauge = combiners(linkinds, tn) + tnᶜ = insert_gauge(tn, combiner_gauge) + boundary_mpsᶜ = contract_approx(tnᶜ; maxdim=maxdim, cutoff=cutoff) + tn_cacheᶜ = contraction_cache(tnᶜ, boundary_mpsᶜ) + tn_cache = insert_gauge(tn_cacheᶜ, combiner_gauge) + return boundary_mps(tn_cache) +end + # Return a network that when contracted equals the # squared norm of the input network. # TODO: rename norm2_network @@ -374,7 +372,8 @@ function sqnorm_approx(ψ::Matrix{ITensor}; center, cutoff, maxdim) tn = ψ′ .* ψᴴ # Insert projectors horizontally (to measure e.g. properties # in a row of the network) - tn_projected = insert_projectors(tn; center=center, cutoff=cutoff, maxdim=maxdim) + bmps = boundary_mps(tn; cutoff=cutoff, maxdim=maxdim) + tn_projected = insert_projectors(tn, bmps; center=center) tn_split, Pl, Pr = tn_projected ψᴴ_split = split_network(ψᴴ) ψ′_split = split_network(ψ′) diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index 0e6d087..14e3129 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -149,6 +149,7 @@ function insert_projectors(peps::PEPS, center, cutoff=1e-15, maxdim=100) psi_bra = addtags(linkinds, dag.(peps.data), "bra") psi_ket = addtags(linkinds, peps.data, "ket") tn = psi_bra .* psi_ket - tn_split, pl, pr = insert_projectors(tn; center=center, cutoff=cutoff, maxdim=maxdim) + bmps = boundary_mps(tn; cutoff=cutoff, maxdim=maxdim) + tn_split, pl, pr = insert_projectors(tn, bmps; center=center) return tn_split, vcat(reduce(vcat, pl), reduce(vcat, pr)) end diff --git a/src/Optimizations/itensor_network.jl b/src/Optimizations/itensor_network.jl index 7ba559a..b0f4e01 100644 --- a/src/Optimizations/itensor_network.jl +++ b/src/Optimizations/itensor_network.jl @@ -7,18 +7,8 @@ function ChainRulesCore.rrule( tn::Matrix{ITensor}; projector_center=default_projector_center(tn), ) - dimy, dimx = size(tn) function pullback(dtn_split::Matrix{ITensor}) - dtn = copy(dtn_split) - for ii in 1:dimy - for jj in 1:dimx - dt, t = dtn[ii, jj], tn[ii, jj] - indices = inds(t) - get_index(i_split) = findfirst(x -> x.id == i_split.id, indices) - indices_reorder = [indices[get_index(i_split)] for i_split in inds(dt)] - dtn[ii, jj] = setinds(dt, Tuple(indices_reorder)) - end - end + dtn = map(t -> replaceprime(t, 1 => 0), dtn_split) return (NoTangent(), dtn, NoTangent()) end return split_network(tn; projector_center=projector_center), pullback