Moving remaining implementations in PEPSAD.jl to the lib

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"

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

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

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

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 <peps|peps'>
+function inner_network(peps::PEPS, peps_prime::PEPS)
+  return vcat(vcat(peps.data...), vcat(peps_prime.data...))
+end
+
+# Get the tensor network of <peps|mpo|peps'>
+# 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

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

src/Optimizations/optimizers.jl

@@ -0,0 +1 @@
+

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, <p|H_1|p>, ..., <p|H_n|p>, <p|p>
+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

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

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