From 15153ff693278aa46151b4e06643408b77ec696d Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 15 Jul 2021 18:08:43 -0500 Subject: [PATCH 01/10] Moving remaining implementations in PEPSAD.jl to the lib --- Project.toml | 2 + src/ITensorNetworkAD.jl | 1 + src/ITensorNetworks/ITensorNetworks.jl | 1 + src/ITensorNetworks/peps.jl | 147 +++++++++++++++++++++++++ src/Optimizations/Optimizations.jl | 14 +++ src/Optimizations/optimizers.jl | 1 + src/Optimizations/peps.jl | 95 ++++++++++++++++ src/Optimizations/run.jl | 47 ++++++++ test/ITensorNetworks/peps.jl | 81 ++++++++++++++ test/ITensorNetworks/runtests.jl | 2 +- test/Optimizations/runtests.jl | 47 ++++++++ test/runtests.jl | 1 + 12 files changed, 438 insertions(+), 1 deletion(-) create mode 100644 src/ITensorNetworks/peps.jl create mode 100644 src/Optimizations/Optimizations.jl create mode 100644 src/Optimizations/optimizers.jl create mode 100644 src/Optimizations/peps.jl create mode 100644 src/Optimizations/run.jl create mode 100644 test/ITensorNetworks/peps.jl create mode 100644 test/Optimizations/runtests.jl diff --git a/Project.toml b/Project.toml index 71a74a2..338aeb1 100644 --- a/Project.toml +++ b/Project.toml @@ -7,7 +7,9 @@ version = "0.0.1" AutoHOOT = "0bdf4f06-f75d-419a-8d6b-4ad79dc3c0f5" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5" +OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" [compat] diff --git a/src/ITensorNetworkAD.jl b/src/ITensorNetworkAD.jl index 7cbb940..79cf89d 100644 --- a/src/ITensorNetworkAD.jl +++ b/src/ITensorNetworkAD.jl @@ -5,5 +5,6 @@ using Reexport include("ITensorNetworks/ITensorNetworks.jl") include("ITensorChainRules/ITensorChainRules.jl") include("ITensorAutoHOOT/ITensorAutoHOOT.jl") +include("Optimizations/Optimizations.jl") end diff --git a/src/ITensorNetworks/ITensorNetworks.jl b/src/ITensorNetworks/ITensorNetworks.jl index b732cf6..3c24300 100644 --- a/src/ITensorNetworks/ITensorNetworks.jl +++ b/src/ITensorNetworks/ITensorNetworks.jl @@ -5,5 +5,6 @@ using ITensors include("models/models.jl") include("tensor_networks.jl") include("boundary_mps_projectors.jl") +include("peps.jl") end diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl new file mode 100644 index 0000000..7ef28b8 --- /dev/null +++ b/src/ITensorNetworks/peps.jl @@ -0,0 +1,147 @@ +""" +A finite size PEPS type. +""" +mutable struct PEPS + data::Matrix{ITensor} +end + +PEPS(Nx::Int, Ny::Int) = PEPS(Matrix{ITensor}(undef, Nx, Ny)) + +""" + PEPS([::Type{ElT} = Float64, sites; linkdims=1) +Construct an PEPS filled with Empty ITensors of type `ElT` from a collection of indices. +Optionally specify the link dimension with the keyword argument `linkdims`, which by default is 1. +""" +function PEPS(::Type{T}, sites::Matrix{<:Index}; linkdims::Integer=1) where {T<:Number} + Ny, Nx = size(sites) + tensor_grid = Matrix{ITensor}(undef, Ny, Nx) + # we assume the PEPS at least has size (2,2). Can generalize if necessary + @assert(Nx >= 2 && Ny >= 2) + + lh = Matrix{Index}(undef, Ny, Nx - 1) + for ii in 1:(Nx - 1) + for jj in 1:(Ny) + lh[jj, ii] = Index(linkdims, "Lh,$jj,$ii") + end + end + lv = Matrix{Index}(undef, Ny - 1, Nx) + for ii in 1:(Nx) + for jj in 1:(Ny - 1) + lv[jj, ii] = Index(linkdims, "Lv,$jj,$ii") + end + end + + # boundary cases + tensor_grid[1, 1] = emptyITensor(T, lh[1, 1], lv[1, 1], sites[1, 1]) + tensor_grid[1, Nx] = emptyITensor(T, lh[1, Nx - 1], lv[1, Nx], sites[1, Nx]) + tensor_grid[Ny, 1] = emptyITensor(T, lh[Ny, 1], lv[Ny - 1, 1], sites[Ny, 1]) + tensor_grid[Ny, Nx] = emptyITensor(T, lh[Ny, Nx - 1], lv[Ny - 1, Nx], sites[Ny, Nx]) + for ii in 2:(Nx - 1) + tensor_grid[1, ii] = emptyITensor(T, lh[1, ii], lh[1, ii - 1], lv[1, ii], sites[1, ii]) + tensor_grid[Ny, ii] = emptyITensor( + T, lh[Ny, ii], lh[Ny, ii - 1], lv[Ny - 1, ii], sites[Ny, ii] + ) + end + + # inner sites + for jj in 2:(Ny - 1) + tensor_grid[jj, 1] = emptyITensor(T, lh[jj, 1], lv[jj, 1], lv[jj - 1, 1], sites[jj, 1]) + tensor_grid[jj, Nx] = emptyITensor( + T, lh[jj, Nx - 1], lv[jj, Nx], lv[jj - 1, Nx], sites[jj, Nx] + ) + for ii in 2:(Nx - 1) + tensor_grid[jj, ii] = emptyITensor( + T, lh[jj, ii], lh[jj, ii - 1], lv[jj, ii], lv[jj - 1, ii], sites[jj, ii] + ) + end + end + + return PEPS(tensor_grid) +end + +PEPS(sites::Matrix{<:Index}, args...; kwargs...) = PEPS(Float64, sites, args...; kwargs...) + +PEPS(data::Array, size::Tuple{Int64,Int64}) = PEPS(reshape(data, size[1], size[2])) + +function randomizePEPS!(P::PEPS) + dimy, dimx = size(P.data) + for ii in 1:dimx + for jj in 1:dimy + randn!(P.data[jj, ii]) + normalize!(P.data[jj, ii]) + end + end +end + +function Base.:+(A::PEPS, B::PEPS) + @assert(size(A.data) == size(B.data)) + out = [t1 + t2 for (t1, t2) in zip(vcat(A.data...), vcat(B.data...))] + return PEPS(out, size(A.data)) +end + +function Base.:-(A::PEPS, B::PEPS) + @assert(size(A.data) == size(B.data)) + out = [t1 - t2 for (t1, t2) in zip(vcat(A.data...), vcat(B.data...))] + return PEPS(out, size(A.data)) +end + +function Base.:*(c::Number, A::PEPS) + out = [c * t for t in vcat(A.data...)] + return PEPS(out, size(A.data)) +end + +function Base.:*(A::PEPS, B::PEPS) + @assert(size(A.data) == size(B.data)) + out_scalar = 0 + for (t1, t2) in zip(vcat(A.data...), vcat(B.data...)) + out_scalar += ITensors.scalar(t1 * t2) + end + return out_scalar +end + +function ITensors.prime(peps::PEPS; ham=true) + prime_peps = [] + for tensor in vcat(peps.data...) + indices = inds(tensor) + if ham == true + bonds = indices + else + bonds = [indices[i] for i in 1:(length(indices) - 1)] + end + prime_peps = vcat(prime_peps, [ITensors.prime(tensor, bonds)]) + end + return PEPS(prime_peps, size(peps.data)) +end + +# Get the tensor network of +function inner_network(peps::PEPS, peps_prime::PEPS) + return vcat(vcat(peps.data...), vcat(peps_prime.data...)) +end + +# Get the tensor network of +# The local MPO specifies the 2-site term of the Hamiltonian +function inner_network( + peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, mpo::MPO, coordinates::Array +) + @assert(length(mpo) == length(coordinates)) + network = vcat(peps.data...) + dimy, dimx = size(peps.data) + for ii in 1:dimx + for jj in 1:dimy + if (jj => ii) in coordinates + index = findall(x -> x == (jj => ii), coordinates) + @assert(length(index) == 1) + network = vcat(network, [mpo.data[index[1]]]) + network = vcat(network, [peps_prime_ham.data[jj, ii]]) + else + network = vcat(network, [peps_prime.data[jj, ii]]) + end + end + end + return network +end + +function extract_data(v::Array{<:PEPS}) + tensor_list = [vcat(peps.data...) for peps in v] + return vcat(tensor_list...) +end diff --git a/src/Optimizations/Optimizations.jl b/src/Optimizations/Optimizations.jl new file mode 100644 index 0000000..a494ab4 --- /dev/null +++ b/src/Optimizations/Optimizations.jl @@ -0,0 +1,14 @@ +module Optimizations + +using ITensors + +# peps and models +export PEPS, randomizePEPS!, inner_network, mpo, localham, checklocalham, Model, prime +# optimizations +export gradient_descent, optimize, generate_inner_network, extract_data + +include("peps.jl") +include("run.jl") +include("optimizers.jl") + +end diff --git a/src/Optimizations/optimizers.jl b/src/Optimizations/optimizers.jl new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/src/Optimizations/optimizers.jl @@ -0,0 +1 @@ + diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl new file mode 100644 index 0000000..fdf79e4 --- /dev/null +++ b/src/Optimizations/peps.jl @@ -0,0 +1,95 @@ +using AutoHOOT, ChainRulesCore, Zygote +using ..ITensorAutoHOOT +using ..ITensorNetworks +using ITensors: setinds +using ..ITensorNetworks: PEPS, randomizePEPS!, inner_network, extract_data +using ..ITensorAutoHOOT: batch_tensor_contraction + +function ChainRulesCore.rrule(::typeof(prime), peps::PEPS; ham=true) + dimy, dimx = size(peps.data) + peps_vec = reshape(peps.data, dimy * dimx) + function adjoint_pullback(dpeps_prime::PEPS) + dpeps_prime_vec = reshape(dpeps_prime.data, dimy * dimx) + dpeps_vec = [] + for i in 1:(dimy * dimx) + indices = inds(peps_vec[i]) + indices_reorder = [] + for i_prime in inds(dpeps_prime_vec[i]) + index = findall(x -> x.id == i_prime.id, indices) + @assert(length(index) == 1) + push!(indices_reorder, indices[index[1]]) + end + push!(dpeps_vec, setinds(dpeps_prime_vec[i], Tuple(indices_reorder))) + end + dpeps = PEPS(dpeps_vec, (dimy, dimx)) + return (NoTangent(), dpeps, NoTangent()) + end + return prime(peps; ham=ham), adjoint_pullback +end + +function ChainRulesCore.rrule(::typeof(extract_data), v::Array{<:PEPS}) + size_list = [size(peps.data) for peps in v] + function adjoint_pullback(dt) + dt = [t for t in dt] + index = 0 + dv = [] + for (dimy, dimx) in size_list + size = dimy * dimx + d_peps = PEPS(reshape(dt[(index + 1):(index + size)], dimy, dimx)) + index += size + push!(dv, d_peps) + end + return (NoTangent(), dv) + end + return extract_data(v), adjoint_pullback +end + +"""Generate an array of networks representing inner products, , ..., , +Parameters +---------- +peps: a peps network with datatype PEPS +peps_prime: prime of peps used for inner products +peps_prime_ham: prime of peps used for calculating expectation values +Hlocal: An array of MPO operators with datatype LocalMPO +Returns +------- +An array of networks. +""" +function generate_inner_network( + peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, Hlocal::Array +) + network_list = [] + for H_term in Hlocal + inner = inner_network( + peps, peps_prime, peps_prime_ham, H_term.mpo, [H_term.coord1, H_term.coord2] + ) + network_list = vcat(network_list, [inner]) + end + inner = inner_network(peps, peps_prime) + network_list = vcat(network_list, [inner]) + return network_list +end + +# gradient of this function returns nothing. +@non_differentiable generate_inner_network( + peps::PEPS, peps_prime::PEPS, peps_prime_ham::PEPS, Hlocal::Array +) + +function rayleigh_quotient(inners::Array) + self_inner = inners[length(inners)][] + expectations = sum(inners[1:(length(inners) - 1)])[] + return expectations / self_inner +end + +function loss_grad_wrap(peps::PEPS, Hlocal::Array) + function loss(peps::PEPS) + peps_prime = prime(peps; ham=false) + peps_prime_ham = prime(peps; ham=true) + network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, Hlocal) + variables = extract_data([peps, peps_prime, peps_prime_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/src/Optimizations/run.jl b/src/Optimizations/run.jl new file mode 100644 index 0000000..5bc1295 --- /dev/null +++ b/src/Optimizations/run.jl @@ -0,0 +1,47 @@ +using OptimKit + +"""Update PEPS based on gradient descent +Parameters +---------- +peps: a peps network with datatype PEPS +Hlocal: An array of MPO operators with datatype LocalMPO +stepsize: step size used in the gradient descent +num_sweeps: number of gradient descent sweeps/iterations +Returns +------- +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 = peps - stepsize * g + push!(losses, l) + end + return losses +end + +function optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") + @assert(method in ["GD", "LBFGS", "CG"]) + inner(x, peps1, peps2) = peps1 * peps2 + loss_w_grad = loss_grad_wrap(peps, Hlocal) + scale(peps, alpha) = alpha * peps + add(peps1, peps2, alpha) = peps1 + alpha * peps2 + linesearch = HagerZhangLineSearch() + if method == "GD" + alg = GradientDescent(num_sweeps, 1e-8, linesearch, 2) + elseif method == "LBFGS" + alg = LBFGS(16; maxiter=num_sweeps, gradtol=1e-8, linesearch=linesearch, verbosity=2) + elseif method == "CG" + alg = ConjugateGradient(; + maxiter=num_sweeps, gradtol=1e-8, linesearch=linesearch, verbosity=2 + ) + end + _, _, _, _, history = OptimKit.optimize( + loss_w_grad, peps, alg; inner=inner, (scale!)=scale, (add!)=add + ) + return history[:, 1] +end diff --git a/test/ITensorNetworks/peps.jl b/test/ITensorNetworks/peps.jl new file mode 100644 index 0000000..ba0e5fa --- /dev/null +++ b/test/ITensorNetworks/peps.jl @@ -0,0 +1,81 @@ +using ITensors, ITensorNetworkAD, AutoHOOT +using ITensorNetworkAD.ITensorNetworks: PEPS, randomizePEPS!, inner_network +using ITensorNetworkAD.ITensorAutoHOOT: generate_optimal_tree + +@testset "test peps" begin + Nx = 4 + Ny = 5 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps = PEPS(sites) + + for ii in 1:(Ny - 1) + for jj in 1:(Nx - 1) + inds1 = inds(peps.data[ii, jj]) + inds2 = inds(peps.data[ii, jj + 1]) + inds3 = inds(peps.data[ii + 1, jj]) + inds4 = inds(peps.data[ii + 1, jj + 1]) + @test length(intersect(inds1, inds2)) == 1 + @test length(intersect(inds1, inds3)) == 1 + @test length(intersect(inds1, inds4)) == 0 + end + end +end + +@testset "test inner product" begin + Nx = 2 + Ny = 3 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps = PEPS(sites) + randomizePEPS!(peps) + peps_prime = prime(peps; ham=false) + inner = inner_network(peps, peps_prime) + + opt_inner = generate_optimal_tree(inner) + out = contract(opt_inner) + # output is a scalar + @test size(out) == () +end + +@testset "test inner product with hamiltonian" begin + Nx = 3 + Ny = 4 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps = PEPS(sites) + randomizePEPS!(peps) + + opsum = OpSum() + opsum += 0.5, "S+", 1, "S-", 2 + opsum += 0.5, "S-", 1, "S+", 2 + opsum += "Sz", 1, "Sz", 2 + mpo = MPO(opsum, [sites[2, 2], sites[2, 3]]) + + peps_prime = prime(peps; ham=false) + peps_prime_ham = prime(peps; ham=true) + inner = inner_network(peps, peps_prime, peps_prime_ham, mpo, [2 => 2, 2 => 3]) + opt_inner = generate_optimal_tree(inner) + out = contract(opt_inner) + # output is a scalar + @test size(out) == () +end + +@testset "test plus, minus, multiplication" begin + Nx = 2 + Ny = 3 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps1 = PEPS(sites) + randomizePEPS!(peps1) + peps2 = 1.5 * peps1 + peps3 = peps1 + peps2 + peps4 = peps1 - peps2 + for i in 1:Nx + for j in 1:Ny + @test norm(peps2.data[j, i]) == norm(1.5 * peps1.data[j, i]) + @test norm(peps3.data[j, i]) == norm(peps1.data[j, i] + peps2.data[j, i]) + @test norm(peps4.data[j, i]) == norm(peps1.data[j, i] - peps2.data[j, i]) + end + end +end diff --git a/test/ITensorNetworks/runtests.jl b/test/ITensorNetworks/runtests.jl index dc2b764..bb43d58 100644 --- a/test/ITensorNetworks/runtests.jl +++ b/test/ITensorNetworks/runtests.jl @@ -2,7 +2,7 @@ using ITensorNetworkAD using Test @testset "ITensorNetworks.jl" begin - for filename in ["models.jl", "projectors.jl"] + for filename in ["peps.jl", "models.jl", "projectors.jl", "peps.jl"] println("Running $filename in ITensorNetworks.jl") include(filename) end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl new file mode 100644 index 0000000..13662c0 --- /dev/null +++ b/test/Optimizations/runtests.jl @@ -0,0 +1,47 @@ +using ITensors, ITensorNetworkAD, AutoHOOT, Zygote +using ITensorNetworkAD.ITensorNetworks: + PEPS, randomizePEPS!, inner_network, Models, extract_data +using ITensorNetworkAD.Optimizations: gradient_descent, optimize, generate_inner_network +using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction + +@testset "test monotonic loss decrease of optimization" begin + Nx, Ny = 2, 3 + num_sweeps = 20 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps = PEPS(sites; linkdims=10) + randomizePEPS!(peps) + H_local = Models.localham(Models.Model("tfim"), sites; h=1.0) + losses_gd = gradient_descent(peps, H_local; stepsize=0.005, num_sweeps=num_sweeps) + losses_ls = optimize(peps, H_local; num_sweeps=num_sweeps, method="GD") + losses_lbfgs = optimize(peps, H_local; num_sweeps=num_sweeps, method="LBFGS") + losses_cg = optimize(peps, H_local; 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 + sites = siteinds("S=1/2", Nx * Ny) + sites = reshape(sites, Ny, Nx) + peps = PEPS(sites; linkdims=2) + randomizePEPS!(peps) + function loss(peps::PEPS) + peps_prime = prime(peps; ham=false) + peps_prime_ham = prime(peps; ham=true) + network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, []) + variables = extract_data([peps, peps_prime]) + inners = batch_tensor_contraction(network_list, variables...) + return sum(inners)[] + end + g = gradient(loss, peps) + inner = inner_network(peps, prime(peps; ham=false)) + g_true_first_site = contract(inner[2:length(inner)]) + g_true_first_site = 2 * g_true_first_site + @test isapprox(norm(g[1].data[1, 1]), norm(g_true_first_site)) +end diff --git a/test/runtests.jl b/test/runtests.jl index 884afa4..5a0e1be 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -6,6 +6,7 @@ using Test "ITensorAutoHOOT/runtests.jl", "ITensorChainRules/runtests.jl", "ITensorNetworks/runtests.jl", + "Optimizations/runtests.jl", ] println("Running $filename") include(filename) From d815623a21ae312a46709359e217cea7e6bedc0e Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 15 Jul 2021 18:31:18 -0500 Subject: [PATCH 02/10] Change redefinition of optimize to OptimKit.optimize --- src/Optimizations/Optimizations.jl | 2 +- src/Optimizations/run.jl | 2 +- test/Optimizations/runtests.jl | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/Optimizations/Optimizations.jl b/src/Optimizations/Optimizations.jl index a494ab4..235aa1b 100644 --- a/src/Optimizations/Optimizations.jl +++ b/src/Optimizations/Optimizations.jl @@ -5,7 +5,7 @@ using ITensors # peps and models export PEPS, randomizePEPS!, inner_network, mpo, localham, checklocalham, Model, prime # optimizations -export gradient_descent, optimize, generate_inner_network, extract_data +export gradient_descent, generate_inner_network, extract_data include("peps.jl") include("run.jl") diff --git a/src/Optimizations/run.jl b/src/Optimizations/run.jl index 5bc1295..4bbff03 100644 --- a/src/Optimizations/run.jl +++ b/src/Optimizations/run.jl @@ -24,7 +24,7 @@ function gradient_descent(peps::PEPS, Hlocal::Array; stepsize::Float64, num_swee return losses end -function optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") +function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") @assert(method in ["GD", "LBFGS", "CG"]) inner(x, peps1, peps2) = peps1 * peps2 loss_w_grad = loss_grad_wrap(peps, Hlocal) diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index 13662c0..be38a38 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -1,7 +1,7 @@ using ITensors, ITensorNetworkAD, AutoHOOT, Zygote using ITensorNetworkAD.ITensorNetworks: PEPS, randomizePEPS!, inner_network, Models, extract_data -using ITensorNetworkAD.Optimizations: gradient_descent, optimize, generate_inner_network +using ITensorNetworkAD.Optimizations: gradient_descent, generate_inner_network using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction @testset "test monotonic loss decrease of optimization" begin From 5b6119109cf475543934c461ef09a95ebb0e69a9 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 16 Jul 2021 00:07:21 +0000 Subject: [PATCH 03/10] CompatHelper: add new compat entry for "OptimKit" at version "0.3" --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 338aeb1..23390e2 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" [compat] ChainRulesCore = "0.10" ITensors = "0.2" +OptimKit = "0.3" Reexport = "1" ZygoteRules = "0.2" julia = "1.3" From 0db1aa0edc90b01d32af53b8e4a0bb1ec7973919 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 16 Jul 2021 00:07:26 +0000 Subject: [PATCH 04/10] CompatHelper: add new compat entry for "Zygote" at version "0.6" --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 338aeb1..226d1dd 100644 --- a/Project.toml +++ b/Project.toml @@ -16,5 +16,6 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" ChainRulesCore = "0.10" ITensors = "0.2" Reexport = "1" +Zygote = "0.6" ZygoteRules = "0.2" julia = "1.3" From cc2f57fc6ac7e1345c2333703729f612eaa913f1 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 15 Jul 2021 20:32:02 -0500 Subject: [PATCH 05/10] Minor changes --- src/Optimizations/Optimizations.jl | 5 +---- test/ITensorNetworks/peps.jl | 18 +++++++----------- test/ITensorNetworks/runtests.jl | 2 +- test/Optimizations/runtests.jl | 10 ++++------ 4 files changed, 13 insertions(+), 22 deletions(-) diff --git a/src/Optimizations/Optimizations.jl b/src/Optimizations/Optimizations.jl index 235aa1b..eebcb15 100644 --- a/src/Optimizations/Optimizations.jl +++ b/src/Optimizations/Optimizations.jl @@ -2,10 +2,7 @@ module Optimizations using ITensors -# peps and models -export PEPS, randomizePEPS!, inner_network, mpo, localham, checklocalham, Model, prime -# optimizations -export gradient_descent, generate_inner_network, extract_data +export gradient_descent, generate_inner_network include("peps.jl") include("run.jl") diff --git a/test/ITensorNetworks/peps.jl b/test/ITensorNetworks/peps.jl index ba0e5fa..c7a2a35 100644 --- a/test/ITensorNetworks/peps.jl +++ b/test/ITensorNetworks/peps.jl @@ -5,8 +5,7 @@ using ITensorNetworkAD.ITensorAutoHOOT: generate_optimal_tree @testset "test peps" begin Nx = 4 Ny = 5 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) for ii in 1:(Ny - 1) @@ -25,8 +24,7 @@ end @testset "test inner product" begin Nx = 2 Ny = 3 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) randomizePEPS!(peps) peps_prime = prime(peps; ham=false) @@ -41,8 +39,7 @@ end @testset "test inner product with hamiltonian" begin Nx = 3 Ny = 4 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) randomizePEPS!(peps) @@ -64,8 +61,7 @@ end @testset "test plus, minus, multiplication" begin Nx = 2 Ny = 3 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps1 = PEPS(sites) randomizePEPS!(peps1) peps2 = 1.5 * peps1 @@ -73,9 +69,9 @@ end peps4 = peps1 - peps2 for i in 1:Nx for j in 1:Ny - @test norm(peps2.data[j, i]) == norm(1.5 * peps1.data[j, i]) - @test norm(peps3.data[j, i]) == norm(peps1.data[j, i] + peps2.data[j, i]) - @test norm(peps4.data[j, i]) == norm(peps1.data[j, i] - peps2.data[j, i]) + @test isapprox(peps2.data[j, i], 1.5 * peps1.data[j, i]) + @test isapprox(peps3.data[j, i], peps1.data[j, i] + peps2.data[j, i]) + @test isapprox(peps4.data[j, i], peps1.data[j, i] - peps2.data[j, i]) end end end diff --git a/test/ITensorNetworks/runtests.jl b/test/ITensorNetworks/runtests.jl index bb43d58..e116e8d 100644 --- a/test/ITensorNetworks/runtests.jl +++ b/test/ITensorNetworks/runtests.jl @@ -2,7 +2,7 @@ using ITensorNetworkAD using Test @testset "ITensorNetworks.jl" begin - for filename in ["peps.jl", "models.jl", "projectors.jl", "peps.jl"] + for filename in ["peps.jl", "models.jl", "projectors.jl"] println("Running $filename in ITensorNetworks.jl") include(filename) end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index be38a38..d7797c8 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -1,4 +1,4 @@ -using ITensors, ITensorNetworkAD, AutoHOOT, Zygote +using ITensors, ITensorNetworkAD, AutoHOOT, Zygote, OptimKit using ITensorNetworkAD.ITensorNetworks: PEPS, randomizePEPS!, inner_network, Models, extract_data using ITensorNetworkAD.Optimizations: gradient_descent, generate_inner_network @@ -7,8 +7,7 @@ using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction @testset "test monotonic loss decrease of optimization" begin Nx, Ny = 2, 3 num_sweeps = 20 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites; linkdims=10) randomizePEPS!(peps) H_local = Models.localham(Models.Model("tfim"), sites; h=1.0) @@ -27,8 +26,7 @@ end @testset "test inner product gradient" begin Nx = 2 Ny = 2 - sites = siteinds("S=1/2", Nx * Ny) - sites = reshape(sites, Ny, Nx) + sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites; linkdims=2) randomizePEPS!(peps) function loss(peps::PEPS) @@ -43,5 +41,5 @@ end inner = inner_network(peps, prime(peps; ham=false)) g_true_first_site = contract(inner[2:length(inner)]) g_true_first_site = 2 * g_true_first_site - @test isapprox(norm(g[1].data[1, 1]), norm(g_true_first_site)) + @test isapprox(g[1].data[1, 1], g_true_first_site) end From 84785d060d1c5195364c56f4c66d5a47459a5085 Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Thu, 15 Jul 2021 22:38:57 -0500 Subject: [PATCH 06/10] Rewrite PEPS prime function --- src/ITensorNetworks/peps.jl | 17 ++++++----------- src/Optimizations/peps.jl | 7 ++++--- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index 7ef28b8..ac42b65 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -99,18 +99,13 @@ function Base.:*(A::PEPS, B::PEPS) return out_scalar end +get_prime_bonds(indices; ham=true) = ham ? indices : indices[1:(end - 1)] + function ITensors.prime(peps::PEPS; ham=true) - prime_peps = [] - for tensor in vcat(peps.data...) - indices = inds(tensor) - if ham == true - bonds = indices - else - bonds = [indices[i] for i in 1:(length(indices) - 1)] - end - prime_peps = vcat(prime_peps, [ITensors.prime(tensor, bonds)]) - end - return PEPS(prime_peps, size(peps.data)) + prime_peps = map( + tensor -> prime(tensor, get_prime_bonds(inds(tensor); ham=ham)), peps.data + ) + return PEPS(prime_peps) end # Get the tensor network of diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index fdf79e4..943da23 100644 --- a/src/Optimizations/peps.jl +++ b/src/Optimizations/peps.jl @@ -5,11 +5,12 @@ using ITensors: setinds using ..ITensorNetworks: PEPS, randomizePEPS!, inner_network, extract_data using ..ITensorAutoHOOT: batch_tensor_contraction -function ChainRulesCore.rrule(::typeof(prime), peps::PEPS; ham=true) +# TODO: rewrite this function +function ChainRulesCore.rrule(::typeof(ITensors.prime), peps::PEPS; ham=true) dimy, dimx = size(peps.data) - peps_vec = reshape(peps.data, dimy * dimx) + peps_vec = vec(peps.data) function adjoint_pullback(dpeps_prime::PEPS) - dpeps_prime_vec = reshape(dpeps_prime.data, dimy * dimx) + dpeps_prime_vec = vec(dpeps_prime.data) dpeps_vec = [] for i in 1:(dimy * dimx) indices = inds(peps_vec[i]) From 57f0abd1cb0625f430c4f4578e4acd95d039cbaa Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 16 Jul 2021 17:31:23 -0500 Subject: [PATCH 07/10] Rewrite several functions for PEPS based on comments --- Project.toml | 1 + src/ITensorNetworks/peps.jl | 62 +++++++++++----------------------- src/Optimizations/peps.jl | 4 +-- src/Optimizations/run.jl | 8 +++-- test/ITensorNetworks/peps.jl | 15 ++++---- test/Optimizations/runtests.jl | 7 ++-- 6 files changed, 39 insertions(+), 58 deletions(-) diff --git a/Project.toml b/Project.toml index 009bfea..ac46dcc 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ AutoHOOT = "0bdf4f06-f75d-419a-8d6b-4ad79dc3c0f5" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ITensors = "9136182c-28ba-11e9-034c-db9fb085ebd5" OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index ac42b65..5cde3f2 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -1,7 +1,9 @@ +using Random + """ A finite size PEPS type. """ -mutable struct PEPS +struct PEPS data::Matrix{ITensor} end @@ -32,25 +34,25 @@ function PEPS(::Type{T}, sites::Matrix{<:Index}; linkdims::Integer=1) where {T<: end # boundary cases - tensor_grid[1, 1] = emptyITensor(T, lh[1, 1], lv[1, 1], sites[1, 1]) - tensor_grid[1, Nx] = emptyITensor(T, lh[1, Nx - 1], lv[1, Nx], sites[1, Nx]) - tensor_grid[Ny, 1] = emptyITensor(T, lh[Ny, 1], lv[Ny - 1, 1], sites[Ny, 1]) - tensor_grid[Ny, Nx] = emptyITensor(T, lh[Ny, Nx - 1], lv[Ny - 1, Nx], sites[Ny, Nx]) + tensor_grid[1, 1] = ITensor(T, lh[1, 1], lv[1, 1], sites[1, 1]) + tensor_grid[1, Nx] = ITensor(T, lh[1, Nx - 1], lv[1, Nx], sites[1, Nx]) + tensor_grid[Ny, 1] = ITensor(T, lh[Ny, 1], lv[Ny - 1, 1], sites[Ny, 1]) + tensor_grid[Ny, Nx] = ITensor(T, lh[Ny, Nx - 1], lv[Ny - 1, Nx], sites[Ny, Nx]) for ii in 2:(Nx - 1) - tensor_grid[1, ii] = emptyITensor(T, lh[1, ii], lh[1, ii - 1], lv[1, ii], sites[1, ii]) - tensor_grid[Ny, ii] = emptyITensor( + tensor_grid[1, ii] = ITensor(T, lh[1, ii], lh[1, ii - 1], lv[1, ii], sites[1, ii]) + tensor_grid[Ny, ii] = ITensor( T, lh[Ny, ii], lh[Ny, ii - 1], lv[Ny - 1, ii], sites[Ny, ii] ) end # inner sites for jj in 2:(Ny - 1) - tensor_grid[jj, 1] = emptyITensor(T, lh[jj, 1], lv[jj, 1], lv[jj - 1, 1], sites[jj, 1]) - tensor_grid[jj, Nx] = emptyITensor( + tensor_grid[jj, 1] = ITensor(T, lh[jj, 1], lv[jj, 1], lv[jj - 1, 1], sites[jj, 1]) + tensor_grid[jj, Nx] = ITensor( T, lh[jj, Nx - 1], lv[jj, Nx], lv[jj - 1, Nx], sites[jj, Nx] ) for ii in 2:(Nx - 1) - tensor_grid[jj, ii] = emptyITensor( + tensor_grid[jj, ii] = ITensor( T, lh[jj, ii], lh[jj, ii - 1], lv[jj, ii], lv[jj - 1, ii], sites[jj, ii] ) end @@ -61,43 +63,19 @@ end PEPS(sites::Matrix{<:Index}, args...; kwargs...) = PEPS(Float64, sites, args...; kwargs...) -PEPS(data::Array, size::Tuple{Int64,Int64}) = PEPS(reshape(data, size[1], size[2])) - -function randomizePEPS!(P::PEPS) - dimy, dimx = size(P.data) - for ii in 1:dimx - for jj in 1:dimy - randn!(P.data[jj, ii]) - normalize!(P.data[jj, ii]) - end - end +function Random.randn!(P::PEPS) + randn!.(P.data) + normalize!.(P.data) + return P end -function Base.:+(A::PEPS, B::PEPS) - @assert(size(A.data) == size(B.data)) - out = [t1 + t2 for (t1, t2) in zip(vcat(A.data...), vcat(B.data...))] - return PEPS(out, size(A.data)) -end +broadcast_add(A::PEPS, B::PEPS) = PEPS(A.data .+ B.data) -function Base.:-(A::PEPS, B::PEPS) - @assert(size(A.data) == size(B.data)) - out = [t1 - t2 for (t1, t2) in zip(vcat(A.data...), vcat(B.data...))] - return PEPS(out, size(A.data)) -end +broadcast_minus(A::PEPS, B::PEPS) = PEPS(A.data .- B.data) -function Base.:*(c::Number, A::PEPS) - out = [c * t for t in vcat(A.data...)] - return PEPS(out, size(A.data)) -end +broadcast_mul(c::Number, A::PEPS) = PEPS(c .* A.data) -function Base.:*(A::PEPS, B::PEPS) - @assert(size(A.data) == size(B.data)) - out_scalar = 0 - for (t1, t2) in zip(vcat(A.data...), vcat(B.data...)) - out_scalar += ITensors.scalar(t1 * t2) - end - return out_scalar -end +broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> (v[1] * v[2])[], sum, (A.data, B.data)) get_prime_bonds(indices; ham=true) = ham ? indices : indices[1:(end - 1)] diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index 943da23..5b28ed7 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, randomizePEPS!, inner_network, extract_data +using ..ITensorNetworks: PEPS, inner_network, extract_data using ..ITensorAutoHOOT: batch_tensor_contraction # TODO: rewrite this function @@ -22,7 +22,7 @@ function ChainRulesCore.rrule(::typeof(ITensors.prime), peps::PEPS; ham=true) end push!(dpeps_vec, setinds(dpeps_prime_vec[i], Tuple(indices_reorder))) end - dpeps = PEPS(dpeps_vec, (dimy, dimx)) + dpeps = PEPS(reshape(dpeps_vec, (dimy, dimx))) return (NoTangent(), dpeps, NoTangent()) end return prime(peps; ham=ham), adjoint_pullback diff --git a/src/Optimizations/run.jl b/src/Optimizations/run.jl index 4bbff03..2134d89 100644 --- a/src/Optimizations/run.jl +++ b/src/Optimizations/run.jl @@ -1,4 +1,6 @@ using OptimKit +using ..ITensorNetworks +using ..ITensorNetworks: broadcast_add, broadcast_minus, broadcast_mul, broadcast_inner """Update PEPS based on gradient descent Parameters @@ -26,10 +28,10 @@ end function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="GD") @assert(method in ["GD", "LBFGS", "CG"]) - inner(x, peps1, peps2) = peps1 * peps2 + inner(x, peps1, peps2) = broadcast_inner(peps1, peps2) loss_w_grad = loss_grad_wrap(peps, Hlocal) - scale(peps, alpha) = alpha * peps - add(peps1, peps2, alpha) = peps1 + alpha * peps2 + scale(peps, alpha) = broadcast_mul(alpha, peps) + add(peps1, peps2, alpha) = broadcast_add(peps1, broadcast_mul(alpha, pep2)) linesearch = HagerZhangLineSearch() if method == "GD" alg = GradientDescent(num_sweeps, 1e-8, linesearch, 2) diff --git a/test/ITensorNetworks/peps.jl b/test/ITensorNetworks/peps.jl index c7a2a35..afdb931 100644 --- a/test/ITensorNetworks/peps.jl +++ b/test/ITensorNetworks/peps.jl @@ -1,5 +1,6 @@ using ITensors, ITensorNetworkAD, AutoHOOT -using ITensorNetworkAD.ITensorNetworks: PEPS, randomizePEPS!, inner_network +using ITensorNetworkAD.ITensorNetworks: + PEPS, inner_network, broadcast_add, broadcast_minus, broadcast_mul, broadcast_inner using ITensorNetworkAD.ITensorAutoHOOT: generate_optimal_tree @testset "test peps" begin @@ -26,7 +27,7 @@ end Ny = 3 sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) - randomizePEPS!(peps) + randn!(peps) peps_prime = prime(peps; ham=false) inner = inner_network(peps, peps_prime) @@ -41,7 +42,7 @@ end Ny = 4 sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) - randomizePEPS!(peps) + randn!(peps) opsum = OpSum() opsum += 0.5, "S+", 1, "S-", 2 @@ -63,10 +64,10 @@ end Ny = 3 sites = siteinds("S=1/2", Ny, Nx) peps1 = PEPS(sites) - randomizePEPS!(peps1) - peps2 = 1.5 * peps1 - peps3 = peps1 + peps2 - peps4 = peps1 - peps2 + randn!(peps1) + peps2 = broadcast_mul(1.5, peps1) + peps3 = broadcast_add(peps1, peps2) + peps4 = broadcast_minus(peps1, peps2) for i in 1:Nx for j in 1:Ny @test isapprox(peps2.data[j, i], 1.5 * peps1.data[j, i]) diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index d7797c8..d12dcac 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -1,6 +1,5 @@ using ITensors, ITensorNetworkAD, AutoHOOT, Zygote, OptimKit -using ITensorNetworkAD.ITensorNetworks: - PEPS, randomizePEPS!, inner_network, Models, extract_data +using ITensorNetworkAD.ITensorNetworks: PEPS, inner_network, Models, extract_data using ITensorNetworkAD.Optimizations: gradient_descent, generate_inner_network using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction @@ -9,7 +8,7 @@ using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction num_sweeps = 20 sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites; linkdims=10) - randomizePEPS!(peps) + randn!(peps) H_local = Models.localham(Models.Model("tfim"), sites; h=1.0) losses_gd = gradient_descent(peps, H_local; stepsize=0.005, num_sweeps=num_sweeps) losses_ls = optimize(peps, H_local; num_sweeps=num_sweeps, method="GD") @@ -28,7 +27,7 @@ end Ny = 2 sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites; linkdims=2) - randomizePEPS!(peps) + randn!(peps) function loss(peps::PEPS) peps_prime = prime(peps; ham=false) peps_prime_ham = prime(peps; ham=true) From f986206b4d5e7bc14bf4c2e6f665df5411f7526b Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Fri, 16 Jul 2021 22:15:08 -0500 Subject: [PATCH 08/10] Fix bugs w.r.t. change API for PEPS broadcast operations --- src/ITensorNetworks/peps.jl | 6 ++++-- src/Optimizations/peps.jl | 8 ++++---- src/Optimizations/run.jl | 7 ++++--- test/Optimizations/runtests.jl | 4 ++-- 4 files changed, 14 insertions(+), 11 deletions(-) diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index 5cde3f2..91a90b6 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -69,13 +69,15 @@ function Random.randn!(P::PEPS) return P end +Base.:+(A::PEPS, B::PEPS) = broadcast_add(A, B) + broadcast_add(A::PEPS, B::PEPS) = PEPS(A.data .+ B.data) broadcast_minus(A::PEPS, B::PEPS) = PEPS(A.data .- B.data) broadcast_mul(c::Number, A::PEPS) = PEPS(c .* A.data) -broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> (v[1] * v[2])[], sum, (A.data, B.data)) +broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> v[], +, A.data .* B.data) get_prime_bonds(indices; ham=true) = ham ? indices : indices[1:(end - 1)] @@ -114,7 +116,7 @@ function inner_network( return network end -function extract_data(v::Array{<:PEPS}) +function flatten(v::Array{<:PEPS}) tensor_list = [vcat(peps.data...) for peps in v] return vcat(tensor_list...) end diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index 5b28ed7..19b14c3 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, extract_data +using ..ITensorNetworks: PEPS, inner_network, flatten using ..ITensorAutoHOOT: batch_tensor_contraction # TODO: rewrite this function @@ -28,7 +28,7 @@ function ChainRulesCore.rrule(::typeof(ITensors.prime), peps::PEPS; ham=true) return prime(peps; ham=ham), adjoint_pullback end -function ChainRulesCore.rrule(::typeof(extract_data), v::Array{<:PEPS}) +function ChainRulesCore.rrule(::typeof(flatten), v::Array{<:PEPS}) size_list = [size(peps.data) for peps in v] function adjoint_pullback(dt) dt = [t for t in dt] @@ -42,7 +42,7 @@ function ChainRulesCore.rrule(::typeof(extract_data), v::Array{<:PEPS}) end return (NoTangent(), dv) end - return extract_data(v), adjoint_pullback + return flatten(v), adjoint_pullback end """Generate an array of networks representing inner products, , ..., , @@ -87,7 +87,7 @@ function loss_grad_wrap(peps::PEPS, Hlocal::Array) peps_prime = prime(peps; ham=false) peps_prime_ham = prime(peps; ham=true) network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, Hlocal) - variables = extract_data([peps, peps_prime, peps_prime_ham]) + variables = flatten([peps, peps_prime, peps_prime_ham]) inners = batch_tensor_contraction(network_list, variables...) return rayleigh_quotient(inners) end diff --git a/src/Optimizations/run.jl b/src/Optimizations/run.jl index 2134d89..401aa85 100644 --- a/src/Optimizations/run.jl +++ b/src/Optimizations/run.jl @@ -20,7 +20,7 @@ function gradient_descent(peps::PEPS, Hlocal::Array; stepsize::Float64, num_swee for iter in 1:num_sweeps l, g = loss_w_grad(peps) print("The rayleigh quotient at iteraton $iter is $l\n") - peps = peps - stepsize * g + peps = broadcast_minus(peps, broadcast_mul(stepsize, g)) push!(losses, l) end return losses @@ -31,7 +31,8 @@ function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="G 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, pep2)) + add(peps1, peps2, alpha) = broadcast_add(peps1, broadcast_mul(alpha, peps2)) + retract(peps1, peps2, alpha) = (add(peps1, peps2, alpha), peps2) linesearch = HagerZhangLineSearch() if method == "GD" alg = GradientDescent(num_sweeps, 1e-8, linesearch, 2) @@ -43,7 +44,7 @@ function OptimKit.optimize(peps::PEPS, Hlocal::Array; num_sweeps::Int, method="G ) end _, _, _, _, history = OptimKit.optimize( - loss_w_grad, peps, alg; inner=inner, (scale!)=scale, (add!)=add + loss_w_grad, peps, alg; inner=inner, (scale!)=scale, (add!)=add, retract=retract ) return history[:, 1] end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index d12dcac..06111f4 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -1,5 +1,5 @@ using ITensors, ITensorNetworkAD, AutoHOOT, Zygote, OptimKit -using ITensorNetworkAD.ITensorNetworks: PEPS, inner_network, Models, extract_data +using ITensorNetworkAD.ITensorNetworks: PEPS, inner_network, Models, flatten using ITensorNetworkAD.Optimizations: gradient_descent, generate_inner_network using ITensorNetworkAD.ITensorAutoHOOT: batch_tensor_contraction @@ -32,7 +32,7 @@ end peps_prime = prime(peps; ham=false) peps_prime_ham = prime(peps; ham=true) network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, []) - variables = extract_data([peps, peps_prime]) + variables = flatten([peps, peps_prime]) inners = batch_tensor_contraction(network_list, variables...) return sum(inners)[] end From 94dfc584b6a4ba07afe3f38d6256ea5e9facbf8e Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 19 Jul 2021 14:59:39 -0500 Subject: [PATCH 09/10] Fix a bug brought by resolving conflicts --- src/ITensorNetworks/ITensorNetworks.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ITensorNetworks/ITensorNetworks.jl b/src/ITensorNetworks/ITensorNetworks.jl index 58e53c1..8d6903b 100644 --- a/src/ITensorNetworks/ITensorNetworks.jl +++ b/src/ITensorNetworks/ITensorNetworks.jl @@ -10,5 +10,6 @@ include("lattices.jl") include("inds_network.jl") include("itensor_network.jl") include("boundary_mps.jl") +include("peps.jl") end From 5ee0bdc8a3cc7d6b150af4e22bdf54622e3a768a Mon Sep 17 00:00:00 2001 From: LinjianMa Date: Mon, 19 Jul 2021 18:20:32 -0500 Subject: [PATCH 10/10] Rewrite prime functions for peps --- src/ITensorNetworks/itensor_network.jl | 1 + src/ITensorNetworks/peps.jl | 9 ++---- src/Optimizations/peps.jl | 38 ++++++++++---------------- test/ITensorNetworks/peps.jl | 6 ++-- test/ITensorNetworks/runtests.jl | 2 +- test/Optimizations/runtests.jl | 6 ++-- 6 files changed, 26 insertions(+), 36 deletions(-) diff --git a/src/ITensorNetworks/itensor_network.jl b/src/ITensorNetworks/itensor_network.jl index a942f43..04af495 100644 --- a/src/ITensorNetworks/itensor_network.jl +++ b/src/ITensorNetworks/itensor_network.jl @@ -145,6 +145,7 @@ end function ITensors.prime(::typeof(linkinds), tn, args...) return mapinds(x -> prime(x, args...), linkinds, tn) end + function ITensors.addtags(::typeof(linkinds), tn, args...) return mapinds(x -> addtags(x, args...), linkinds, tn) end diff --git a/src/ITensorNetworks/peps.jl b/src/ITensorNetworks/peps.jl index 91a90b6..e42627a 100644 --- a/src/ITensorNetworks/peps.jl +++ b/src/ITensorNetworks/peps.jl @@ -79,13 +79,10 @@ broadcast_mul(c::Number, A::PEPS) = PEPS(c .* A.data) broadcast_inner(A::PEPS, B::PEPS) = mapreduce(v -> v[], +, A.data .* B.data) -get_prime_bonds(indices; ham=true) = ham ? indices : indices[1:(end - 1)] +ITensors.prime(P::PEPS, n::Integer=1) = PEPS(map(x -> prime(x, n), P.data)) -function ITensors.prime(peps::PEPS; ham=true) - prime_peps = map( - tensor -> prime(tensor, get_prime_bonds(inds(tensor); ham=ham)), peps.data - ) - return PEPS(prime_peps) +function ITensors.prime(::typeof(linkinds), P::PEPS, n::Integer=1) + return PEPS(mapinds(x -> prime(x, n), linkinds, P.data)) end # Get the tensor network of diff --git a/src/Optimizations/peps.jl b/src/Optimizations/peps.jl index 19b14c3..51cc621 100644 --- a/src/Optimizations/peps.jl +++ b/src/Optimizations/peps.jl @@ -5,27 +5,19 @@ using ITensors: setinds using ..ITensorNetworks: PEPS, inner_network, flatten using ..ITensorAutoHOOT: batch_tensor_contraction -# TODO: rewrite this function -function ChainRulesCore.rrule(::typeof(ITensors.prime), peps::PEPS; ham=true) - dimy, dimx = size(peps.data) - peps_vec = vec(peps.data) - function adjoint_pullback(dpeps_prime::PEPS) - dpeps_prime_vec = vec(dpeps_prime.data) - dpeps_vec = [] - for i in 1:(dimy * dimx) - indices = inds(peps_vec[i]) - indices_reorder = [] - for i_prime in inds(dpeps_prime_vec[i]) - index = findall(x -> x.id == i_prime.id, indices) - @assert(length(index) == 1) - push!(indices_reorder, indices[index[1]]) - end - push!(dpeps_vec, setinds(dpeps_prime_vec[i], Tuple(indices_reorder))) - end - dpeps = PEPS(reshape(dpeps_vec, (dimy, dimx))) - return (NoTangent(), dpeps, NoTangent()) - end - return prime(peps; ham=ham), adjoint_pullback +function ChainRulesCore.rrule(::typeof(PEPS), data::Matrix{ITensor}) + return PEPS(data), dpeps -> (NoTangent(), dpeps.data) +end + +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.prime), ::typeof(linkinds), P::PEPS, n::Integer=1 +) + return prime(linkinds, P, n), + dprime -> (NoTangent(), NoTangent(), prime(linkinds, dprime, -n), NoTangent()) end function ChainRulesCore.rrule(::typeof(flatten), v::Array{<:PEPS}) @@ -84,8 +76,8 @@ end function loss_grad_wrap(peps::PEPS, Hlocal::Array) function loss(peps::PEPS) - peps_prime = prime(peps; ham=false) - peps_prime_ham = prime(peps; ham=true) + peps_prime = prime(linkinds, peps) + peps_prime_ham = prime(peps) network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, Hlocal) variables = flatten([peps, peps_prime, peps_prime_ham]) inners = batch_tensor_contraction(network_list, variables...) diff --git a/test/ITensorNetworks/peps.jl b/test/ITensorNetworks/peps.jl index afdb931..2e560a3 100644 --- a/test/ITensorNetworks/peps.jl +++ b/test/ITensorNetworks/peps.jl @@ -28,7 +28,7 @@ end sites = siteinds("S=1/2", Ny, Nx) peps = PEPS(sites) randn!(peps) - peps_prime = prime(peps; ham=false) + peps_prime = prime(linkinds, peps) inner = inner_network(peps, peps_prime) opt_inner = generate_optimal_tree(inner) @@ -50,8 +50,8 @@ end opsum += "Sz", 1, "Sz", 2 mpo = MPO(opsum, [sites[2, 2], sites[2, 3]]) - peps_prime = prime(peps; ham=false) - peps_prime_ham = prime(peps; ham=true) + peps_prime = prime(linkinds, peps) + peps_prime_ham = prime(peps) inner = inner_network(peps, peps_prime, peps_prime_ham, mpo, [2 => 2, 2 => 3]) opt_inner = generate_optimal_tree(inner) out = contract(opt_inner) diff --git a/test/ITensorNetworks/runtests.jl b/test/ITensorNetworks/runtests.jl index a786dc9..2cf9fd6 100644 --- a/test/ITensorNetworks/runtests.jl +++ b/test/ITensorNetworks/runtests.jl @@ -2,7 +2,7 @@ using ITensorNetworkAD using Test @testset "ITensorNetworks.jl" begin - for filename in ["models.jl", "projectors.jl", "itensor_network.jl"] + for filename in ["peps.jl", "models.jl", "projectors.jl", "itensor_network.jl"] println("Running $filename in ITensorNetworks.jl") include(filename) end diff --git a/test/Optimizations/runtests.jl b/test/Optimizations/runtests.jl index 06111f4..d400e6f 100644 --- a/test/Optimizations/runtests.jl +++ b/test/Optimizations/runtests.jl @@ -29,15 +29,15 @@ end peps = PEPS(sites; linkdims=2) randn!(peps) function loss(peps::PEPS) - peps_prime = prime(peps; ham=false) - peps_prime_ham = prime(peps; ham=true) + peps_prime = prime(linkinds, peps) + peps_prime_ham = prime(peps) network_list = generate_inner_network(peps, peps_prime, peps_prime_ham, []) variables = flatten([peps, peps_prime]) inners = batch_tensor_contraction(network_list, variables...) return sum(inners)[] end g = gradient(loss, peps) - inner = inner_network(peps, prime(peps; ham=false)) + 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)