diff --git a/Project.toml b/Project.toml index 71a74a2..ac46dcc 100644 --- a/Project.toml +++ b/Project.toml @@ -7,12 +7,17 @@ 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" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" [compat] ChainRulesCore = "0.10" ITensors = "0.2" +OptimKit = "0.3" Reexport = "1" +Zygote = "0.6" ZygoteRules = "0.2" julia = "1.3" 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 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 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 new file mode 100644 index 0000000..e42627a --- /dev/null +++ b/src/ITensorNetworks/peps.jl @@ -0,0 +1,119 @@ +using Random + +""" +A finite size PEPS type. +""" +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] = 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] = 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] = 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] = ITensor( + 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...) + +function Random.randn!(P::PEPS) + randn!.(P.data) + normalize!.(P.data) + 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[], +, A.data .* B.data) + +ITensors.prime(P::PEPS, n::Integer=1) = PEPS(map(x -> prime(x, n), P.data)) + +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 +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 flatten(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..eebcb15 --- /dev/null +++ b/src/Optimizations/Optimizations.jl @@ -0,0 +1,11 @@ +module Optimizations + +using ITensors + +export gradient_descent, generate_inner_network + +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..51cc621 --- /dev/null +++ b/src/Optimizations/peps.jl @@ -0,0 +1,88 @@ +using AutoHOOT, ChainRulesCore, Zygote +using ..ITensorAutoHOOT +using ..ITensorNetworks +using ITensors: setinds +using ..ITensorNetworks: PEPS, inner_network, flatten +using ..ITensorAutoHOOT: batch_tensor_contraction + +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}) + 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 flatten(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(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...) + 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..401aa85 --- /dev/null +++ b/src/Optimizations/run.jl @@ -0,0 +1,50 @@ +using OptimKit +using ..ITensorNetworks +using ..ITensorNetworks: broadcast_add, broadcast_minus, broadcast_mul, broadcast_inner + +"""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 = broadcast_minus(peps, broadcast_mul(stepsize, g)) + push!(losses, l) + end + return losses +end + +function OptimKit.optimize(peps::PEPS, Hlocal::Array; 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) + 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, retract=retract + ) + return history[:, 1] +end diff --git a/test/ITensorNetworks/peps.jl b/test/ITensorNetworks/peps.jl new file mode 100644 index 0000000..2e560a3 --- /dev/null +++ b/test/ITensorNetworks/peps.jl @@ -0,0 +1,78 @@ +using ITensors, ITensorNetworkAD, AutoHOOT +using ITensorNetworkAD.ITensorNetworks: + PEPS, inner_network, broadcast_add, broadcast_minus, broadcast_mul, broadcast_inner +using ITensorNetworkAD.ITensorAutoHOOT: generate_optimal_tree + +@testset "test peps" begin + Nx = 4 + Ny = 5 + sites = siteinds("S=1/2", 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", Ny, Nx) + peps = PEPS(sites) + randn!(peps) + peps_prime = prime(linkinds, peps) + 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", Ny, Nx) + peps = PEPS(sites) + randn!(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(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) + # output is a scalar + @test size(out) == () +end + +@testset "test plus, minus, multiplication" begin + Nx = 2 + Ny = 3 + sites = siteinds("S=1/2", Ny, Nx) + peps1 = PEPS(sites) + 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]) + @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 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 new file mode 100644 index 0000000..d400e6f --- /dev/null +++ b/test/Optimizations/runtests.jl @@ -0,0 +1,44 @@ +using ITensors, ITensorNetworkAD, AutoHOOT, Zygote, OptimKit +using ITensorNetworkAD.ITensorNetworks: PEPS, inner_network, Models, flatten +using ITensorNetworkAD.Optimizations: gradient_descent, 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", Ny, Nx) + peps = PEPS(sites; linkdims=10) + 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") + 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", Ny, Nx) + peps = PEPS(sites; linkdims=2) + randn!(peps) + function loss(peps::PEPS) + 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(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 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)