quantumkithub / pepskit.jl Goto Github PK

function northwest_corner(E4, C1, E1, peps_above, peps_below=peps_above)
    @planar opt=true corner[-1 -2 -3; -4 -5 -6] := 
        E4[-1 1 2; 3] *
        C1[3; 4] *
        E1[4 5 6; -4] *
        peps_above[21; 11 10 18 17] *
        conj(peps_below[21; 13 14 15 16])*
        τ[5 7;6 8]*τ[-5 9;-6 10]*τ[8 12;9 11]*τ[7 14;12 13]*τ[15 2;16 19]*τ[19 1;17 20]*τ[20 -2;18 -3]
end

function planarcontract!(C::AbstractTensorMap{S,N₁,N₂},
                         A::BraidingTensor{S}   , (oindA, cindA)::Index2Tuple{2,2}, 
                         B::AbstractTensorMap{S}, (cindB, oindB)::Index2Tuple{2,N₃},
                         (p1, p2)::Index2Tuple{N₁,N₂},
                         α::Number, β::Number,
                         backend::Backend...) where {S,N₁,N₂,N₃}

using Revise
using TensorKit, PEPSKit, KrylovKit, OptimKit
########################################
#Define the symmetry gorup of the model#
########################################
sym = (SU2Irrep ⊠ U1Irrep ⊠ FermionParity);  #mixing colors * PN_conservation * it's fermions
#define some elementary IRREPS of this symmetry group.
empty  = sym.(0   ,0,0);  #empty empty
single = sym.(1//2,1,1);  #full empty + empty full
double = sym.(0   ,2,0);  #full full
#define the physical Hilbert space of this model :
ph = Vect[sym](empty => 1, single => 1, double => 1);
    
#############
#Find the GS#
#############
#define IC
D = Vect[sym](empty => 1, single=>1, double=>1)  #intial space of PEPS
intitial_tensors = [TensorMap(rand, ComplexF64, ph, D*D'*D'*D)  TensorMap(rand, ComplexF64, ph, D*D'*D'*D) ;
                    TensorMap(rand, ComplexF64, ph, D*D'*D'*D)  TensorMap(rand, ComplexF64, ph, D*D'*D'*D) ]  ;
initial_state = InfinitePEPS(intitial_tensors);

#Define the GS_search algorithm 
χ = Vect[sym](empty => 3, single=>3, double=>3)  #intial space of the boundary MPS
χ_max = 20                                       #maximal bond dimension of the boundary MPS
ctm_alg = CTMRG(; tol=1e-10, miniter=4, maxiter=100, verbosity=3, trscheme=truncdim( χ_max )); #verbosity can be lowered if you want less output.

#use the above to actually get the GS
state = initial_state
environments = leading_boundary(CTMRGEnv(state, χ), state, ctm_alg);  #initial boundaries

ERROR: MethodError: Cannot `convert` an object of type
  CTMRGEnv{TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing},TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing}} to an object of type
  CTMRGEnv{TensorMap{ComplexSpace,1,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing},TensorMap{ComplexSpace,3,1,Trivial,Array{Complex{Float64},2},Nothing,Nothing}}

Closest candidates are:
  convert(::Type{T}, ::T) where T
   @ Base Base.jl:84
  (::Type{CTMRGEnv{C, T}} where {C, T})(::Any, ::Any)
   @ PEPSKit ~/repos/PEPSKit.jl/src/environments/ctmrgenv.jl:7

Stacktrace:
  [1] setindex!(A::Vector{CTMRGEnv{…}}, x::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}}, i1::Int64)
    @ Base ./array.jl:1021
  [2] setindex!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, i::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}}, q::Int64)
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/orthonormal.jl:42
  [3] _rmul!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, G::LinearAlgebra.Givens{ComplexF64})
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/dense/givens.jl:34
  [4] rmul!(b::KrylovKit.OrthonormalBasis{CTMRGEnv{…}}, G::LinearAlgebra.Givens{ComplexF64})
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/dense/givens.jl:16
  [5] linsolve(operator::Function, b::CTMRGEnv{…}, x₀::CTMRGEnv{…}, alg::GMRES{…}, a₀::Int64, a₁::Int64)
    @ KrylovKit ~/.julia/packages/KrylovKit/wPepJ/src/linsolve/gmres.jl:114
  [6] fpgrad(∂F∂x::CTMRGEnv{…}, ∂f∂x::Function, ∂f∂A::PEPSKit.var"#∂f∂A#174"{…}, y₀::CTMRGEnv{…}, alg::GMRES{…})
    @ PEPSKit ~/repos/PEPSKit.jl/src/algorithms/peps_opt.jl:201
  [7] (::PEPSKit.var"#leading_boundary_pullback#172"{…})(Δenvs′::CTMRGEnv{…})
    @ PEPSKit ~/repos/PEPSKit.jl/src/algorithms/peps_opt.jl:128
  [8] #28
    @ ~/repos/PEPSKit.jl/src/utility/hook_pullback.jl:34 [inlined]
  [9] ZBack
    @ ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:211 [inlined]
 [10] (::Zygote.var"#kw_zpullback#53"{PEPSKit.var"#28#29"{…}})(dy::CTMRGEnv{TrivialTensorMap{…}, TrivialTensorMap{…}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/chainrules.jl:237
 [11] #2
    @ ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:63 [inlined]
 [12] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface2.jl:0
 [13] (::Zygote.var"#75#76"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:91
 [14] withgradient(f::Function, args::InfinitePEPS{TrivialTensorMap{ComplexSpace, 1, 4, Matrix{ComplexF64}}})
    @ Zygote ~/.julia/packages/Zygote/nsBv0/src/compiler/interface.jl:213
 [15] compute_grad(peps::InfinitePEPS{…}, envs::CTMRGEnv{…}, boundary_alg::CTMRG, alg_rrule::GMRES{…})
    @ Main ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:62
 [16] top-level scope
    @ ~/repos/PEPSKit.jl/test/ctmrg/gradients.jl:70

using Random
using PEPSKit
using TensorKit
using Zygote
using KrylovKit

Random.seed!(42039482038)
χbond = 2
χenv = 8
boundary_alg = CTMRG(;
    trscheme=truncdim(χenv), tol=1e-10, miniter=20, maxiter=100, verbosity=1
)
psi = InfinitePEPS(2, χbond)
env = leading_boundary(CTMRGEnv(psi; Venv=ℂ^χenv), psi, boundary_alg)

function compute_grad(peps, envs, boundary_alg, alg_rrule)
    E, g = Zygote.withgradient(peps) do psi
        envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, boundary_alg; alg_rrule)
        return costfun(psi, envs2, H)
    end
    return only(g)
end

g = compute_grad(psi, env, boundary_alg, KrylovKit.GMRES(; tol=1e-16, maxiter=2, verbosity=3))