Add runnable but inefficient LinSolve fpgrad using KrylovKit.linsolve

src/algorithms/peps_opt.jl

@@ -110,8 +110,9 @@ function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::PEPSOptimize{ManualIter
         norma = norm(y.corners[NORTHWEST])
         ϵnew = norm(y′.corners[NORTHWEST] - y.corners[NORTHWEST]) / norma  # Normalize error to get comparable convergence tolerance
         Δϵ = ϵ - ϵnew
-        alg.verbosity > 1 &&
-            @printf("Gradient iter: %3d   ‖Cᵢ₊₁-Cᵢ‖: %.2e   Δ‖Cᵢ₊₁-Cᵢ‖: %.2e\n", i, ϵnew, Δϵ)
+        alg.verbosity > 1 && @printf(
+            "Gradient iter: %3d   ‖Cᵢ₊₁-Cᵢ‖: %.2e   Δ‖Cᵢ₊₁-Cᵢ‖: %.2e\n", i, ϵnew, Δϵ
+        )
         y = y′
         ϵ = ϵnew
 
@@ -123,20 +124,16 @@ function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::PEPSOptimize{ManualIter
     return ∂f∂A(y), ϵ
 end
 
-# Use proper iterative solver to solve gradient problem
+# Use KrylovKit.linsolve to solve gradient linear problem
 function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::PEPSOptimize{LinSolve})
-    # spaces = [space.(getfield(∂F∂x, f)) for f in fieldnames(CTMRGEnv)]
-    # sizes = [map(x -> size(x.data), getfield(∂F∂x, f)) for f in fieldnames(CTMRGEnv)]
-    # op = LinearMap(vecdim(∂F∂x)) do v
-    #     env = unvec(v, spaces..., sizes...)
-    #     x = env - ∂f∂x(env)
-    #     vec(x)
-    # end
-    # envvec = vec(y₀)
-    # info = gmres!(envvec, op, vec(∂F∂x); reltol=alg.grad_convtol, maxiter=alg.grad_maxiter)
-    # y = unvec(envvec, spaces..., sizes...)
-
-    # ∂f∂A(y), info
+    grad_op(env) = env - ∂f∂x(env)
+    y, info = linsolve(grad_op, ∂F∂x, y₀; rtol=alg.fpgrad_tol, maxiter=alg.fpgrad_maxiter)
+
+    if alg.verbosity > 0 && info.converged != 1
+        @warn("gradient fixed-point iteration reached maximal number of iterations:", info)
+    end
+
+    return ∂f∂A(y), info
 end
 
 # Contraction of CTMRGEnv and PEPS tensors with open physical bonds

src/environments/ctmrgenv.jl

@@ -3,7 +3,7 @@ struct CTMRGEnv{C,T}
     edges::Array{T,3}
 end
 
-# initialize ctmrg environments with some random tensors
+# Initialize ctmrg environments with some random tensors
 function CTMRGEnv(peps::InfinitePEPS{P}; Venv=oneunit(spacetype(P))) where {P}
     C_type = tensormaptype(spacetype(P), 1, 1, storagetype(P))
     T_type = tensormaptype(spacetype(P), 3, 1, storagetype(P))
@@ -28,13 +28,13 @@ function CTMRGEnv(peps::InfinitePEPS{P}; Venv=oneunit(spacetype(P))) where {P}
     return CTMRGEnv(corners, edges)
 end
 
-# TODO: Is this needed or not?
 # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation
 function ChainRulesCore.rrule(::Type{CTMRGEnv}, corners, edges)
     ctmrgenv_pullback(ē) = NoTangent(), ē.corners, ē.edges
     return CTMRGEnv(corners, edges), ctmrgenv_pullback
 end
 
+# Rotate corners & edges counter-clockwise
 function Base.rotl90(env::CTMRGEnv{C,T}) where {C,T}
     corners′ = similar(env.corners)
     edges′ = similar(env.edges)
@@ -49,10 +49,44 @@ end
 
 Base.eltype(env::CTMRGEnv) = eltype(env.corners[1])
 
+# Functions used for FP differentiation and by KrylovKit.linsolve
 function Base.:+(e₁::CTMRGEnv, e₂::CTMRGEnv)
     return CTMRGEnv(e₁.corners + e₂.corners, e₁.edges + e₂.edges)
 end
 
 function Base.:-(e₁::CTMRGEnv, e₂::CTMRGEnv)
     return CTMRGEnv(e₁.corners - e₂.corners, e₁.edges - e₂.edges)
-end
+end
+
+Base.:*(α::Number, e::CTMRGEnv) = CTMRGEnv(α * e.corners, α * e.edges)
+Base.similar(e::CTMRGEnv) = CTMRGEnv(similar(e.corners), similar(e.edges))
+
+function LinearAlgebra.mul!(edst::CTMRGEnv, esrc::CTMRGEnv, α::Number)
+    edst.corners .= α * esrc.corners
+    edst.edges .= α * esrc.edges
+    return edst
+end
+
+function LinearAlgebra.rmul!(e::CTMRGEnv, α::Number)
+    rmul!(e.corners, α)
+    rmul!(e.edges, α)
+    return e
+end
+
+function LinearAlgebra.axpy!(α::Number, e₁::CTMRGEnv, e₂::CTMRGEnv)
+    e₂.corners .+= α * e₁.corners
+    e₂.edges .+= α * e₁.edges
+    return e₂
+end
+
+function LinearAlgebra.axpby!(α::Number, e₁::CTMRGEnv, β::Number, e₂::CTMRGEnv)
+    e₂.corners .= α * e₁.corners + β * e₂.corners
+    e₂.edges .+= α * e₁.edges + β * e₂.edges
+    return e₂
+end
+
+function LinearAlgebra.dot(e₁::CTMRGEnv, e₂::CTMRGEnv)
+    dot(e₁.corners, e₂.corners) + dot(e₁.edges, e₂.edges)
+end
+
+LinearAlgebra.norm(e::CTMRGEnv) = norm(e.corners) + norm(e.edges)

src/states/infinitepeps.jl

@@ -86,7 +86,7 @@ VectorInterface.scalartype(T::InfinitePEPS) = scalartype(T.A)
 
 ## Copy
 Base.copy(T::InfinitePEPS) = InfinitePEPS(copy(T.A))
-Base.similar(T::InfinitePEPS, args...) = InfinitePEPS(similar(T.A, args...))
+# Base.similar(T::InfinitePEPS) = InfinitePEPS(similar(T.A))  # TODO: This is incompatible with inner constructor
 Base.repeat(T::InfinitePEPS, counts...) = InfinitePEPS(repeat(T.A, counts...))
 
 Base.getindex(T::InfinitePEPS, args...) = Base.getindex(T.A, args...)
@@ -99,17 +99,18 @@ Base.rotl90(t::InfinitePEPS) = InfinitePEPS(rotl90(rotl90.(t.A)))
 ## Math
 Base.:+(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = InfinitePEPS(ψ₁.A + ψ₂.A)
 Base.:-(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = InfinitePEPS(ψ₁.A - ψ₂.A)
-LinearAlgebra.norm(ψ::InfinitePEPS) = norm(ψ.A)
+Base.:*(α::Number, ψ::InfinitePEPS) = InfinitePEPS(α * ψ.A)
 LinearAlgebra.dot(ψ₁::InfinitePEPS, ψ₂::InfinitePEPS) = dot(ψ₁.A, ψ₂.A)
+LinearAlgebra.norm(ψ::InfinitePEPS) = norm(ψ.A)
 
-# For _scale! during optimization
-function LinearAlgebra.rmul!(ψ::InfinitePEPS, β::Number)
-    rmul!(ψ.A, β)
+# Used in _scale during OptimKit.optimize
+function LinearAlgebra.rmul!(ψ::InfinitePEPS, α::Number)
+    rmul!(ψ.A, α)
     return ψ
 end
 
-# For _add! during optimization
-function LinearAlgebra.axpy!(β::Number, ψ::InfinitePEPS, η::InfinitePEPS)
-    ψ.A .+= η.A .* β
-    return ψ
-end
+# Used in _add during OptimKit.optimize
+function LinearAlgebra.axpy!(α::Number, ψ₁::InfinitePEPS, ψ₂::InfinitePEPS)
+    ψ₂.A .+= α * ψ₁.A
+    return ψ₂
+end