Refactor groundstate algorithms to use find_groundstate! and GeometryStyle dispatch

src/algorithms/groundstate/dmrg.jl

@@ -32,7 +32,7 @@ function DMRG(;
     return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize)
 end
 
-function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environments(ψ, H))
+function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG, envs = environments(ψ, H))
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
     log = IterLog("DMRG")
@@ -109,7 +109,7 @@ function DMRG2(;
     return DMRG2(tol, maxiter, verbosity, alg_eigsolve′, alg_svd, trscheme, finalize)
 end
 
-function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG2, envs = environments(ψ, H))
+function find_groundstate!(::FiniteChainStyle, ψ, H, alg::DMRG2, envs = environments(ψ, H))
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
     log = IterLog("DMRG2")
@@ -165,7 +165,3 @@ function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG2, envs = environm
     end
     return ψ, envs, ϵ
 end
-
-function find_groundstate(ψ, H, alg::Union{DMRG, DMRG2}, envs...; kwargs...)
-    return find_groundstate!(copy(ψ), H, alg, envs...; kwargs...)
-end

src/algorithms/groundstate/find_groundstate.jl

@@ -1,9 +1,13 @@
 """
     find_groundstate(ψ₀, H, [environments]; kwargs...) -> (ψ, environments, ϵ)
-    find_groundstate(ψ₀, H, algorithm, environments) -> (ψ, environments, ϵ)
+    find_groundstate(ψ₀, H, algorithm, [environments]) -> (ψ, environments, ϵ)
+    find_groundstate!(ψ₀, H, [environments]; kwargs...) -> (ψ, environments, ϵ)
+    find_groundstate!(ψ₀, H, algorithm, [environments]) -> (ψ, environments, ϵ)
 
 Compute the ground state for Hamiltonian `H` with initial guess `ψ`. If not specified, an
 optimization algorithm will be attempted based on the supplied keywords.
+The `!` variant updates `ψ` in-place, reusing its storage; the non-mutating variant 
+operates on a copy.
 
 ## Arguments
 - `ψ₀::AbstractMPS`: initial guess
@@ -22,25 +26,43 @@ optimization algorithm will be attempted based on the supplied keywords.
 - `ϵ::Float64`: final convergence error upon terminating the algorithm
 """
 function find_groundstate(
+        ψ::AbstractMPS, H, envs::AbstractMPSEnvironments = environments(ψ, H);
+        kwargs...
+    )
+    return find_groundstate!(copy(ψ), H, envs; kwargs...)
+end
+function find_groundstate(
+        ψ::AbstractMPS, H, alg,
+        envs::AbstractMPSEnvironments = environments(ψ, H)
+    )
+    return find_groundstate!(copy(ψ), H, alg, envs)
+end
+function find_groundstate!(
         ψ::AbstractMPS, H, envs::AbstractMPSEnvironments = environments(ψ, H);
         tol = Defaults.tol, maxiter = Defaults.maxiter,
         verbosity = Defaults.verbosity, trscheme = nothing
     )
-    if isa(ψ, InfiniteMPS)
+    if GeometryStyle(ψ, H) isa InfiniteChainStyle
         alg = VUMPS(; tol = max(1.0e-4, tol), verbosity, maxiter)
         if tol < 1.0e-4
             alg = alg & GradientGrassmann(; tol = tol, maxiter, verbosity)
         end
         if !isnothing(trscheme)
             alg = IDMRG2(; tol = min(1.0e-2, 100tol), verbosity, trscheme) & alg
         end
-    elseif isa(ψ, AbstractFiniteMPS)
+    elseif GeometryStyle(ψ, H) isa FiniteChainStyle
         alg = DMRG(; tol, maxiter, verbosity)
         if !isnothing(trscheme)
             alg = DMRG2(; tol = min(1.0e-2, 100tol), verbosity, trscheme) & alg
         end
     else
         throw(ArgumentError("Unknown input state type"))
     end
-    return find_groundstate(ψ, H, alg, envs)
+    return find_groundstate!(ψ, H, alg, envs)
+end
+function find_groundstate!(
+        ψ::AbstractMPS, H, alg::Algorithm,
+        envs::AbstractMPSEnvironments = environments(ψ, H)
+    )
+    return find_groundstate!(GeometryStyle(ψ, H), ψ, H, alg, envs)
 end

src/algorithms/groundstate/gradient_grassmann.jl

@@ -55,10 +55,11 @@ struct GradientGrassmann{O <: OptimKit.OptimizationAlgorithm, F} <: Algorithm
     end
 end
 
-function find_groundstate(
-        ψ::S, H, alg::GradientGrassmann, envs::P = environments(ψ, H)
+function find_groundstate!(
+        style::GeometryStyle, ψ::S, H, alg::GradientGrassmann,
+        envs::P = environments(ψ, H)
     )::Tuple{S, P, Float64} where {S, P}
-    !isa(ψ, FiniteMPS) || dim(ψ.C[end]) == 1 ||
+    !(style isa FiniteChainStyle) || dim(ψ.C[end]) == 1 ||
         @warn "This is not fully supported - split the mps up in a sum of mps's and optimize separately"
     normalize!(ψ)
 

src/algorithms/groundstate/idmrg.jl

@@ -70,10 +70,12 @@ function IDMRGState{T}(mps::S, operator::O, envs::E, iter::Int, ϵ::Float64, ene
     return IDMRGState{S, O, E, T}(mps, operator, envs, iter, ϵ, T(energy))
 end
 
-function find_groundstate(mps, operator, alg::alg_type, envs = environments(mps, operator)) where {alg_type <: Union{<:IDMRG, <:IDMRG2}}
+function find_groundstate!(
+        ::InfiniteChainStyle, mps, operator, alg::alg_type,
+        envs = environments(mps, operator)
+    ) where {alg_type <: Union{<:IDMRG, <:IDMRG2}}
     (length(mps) ≤ 1 && alg isa IDMRG2) && throw(ArgumentError("unit cell should be >= 2"))
     log = alg isa IDMRG ? IterLog("IDMRG") : IterLog("IDMRG2")
-    mps = copy(mps)
     iter = 0
     ϵ = calc_galerkin(mps, operator, mps, envs)
     E = zero(TensorOperations.promote_contract(scalartype(mps), scalartype(operator)))

src/algorithms/groundstate/vumps.jl

@@ -45,8 +45,9 @@ struct VUMPSState{S, O, E}
     which::Symbol
 end
 
-function find_groundstate(
-        mps::InfiniteMPS, operator, alg::VUMPS, envs = environments(mps, operator)
+function find_groundstate!(
+        ::InfiniteChainStyle, mps, operator, alg::VUMPS,
+        envs = environments(mps, operator)
     )
     return dominant_eigsolve(operator, mps, alg, envs; which = :SR)
 end
@@ -57,7 +58,6 @@ function dominant_eigsolve(
     )
     log = IterLog("VUMPS")
     iter = 0
-    mps = copy(mps)
     ϵ = calc_galerkin(mps, operator, mps, envs)
     alg_environments = updatetol(alg.alg_environments, iter, ϵ)
     recalculate!(envs, mps, operator, mps; alg_environments.tol)

src/algorithms/unionalg.jl

@@ -32,3 +32,8 @@ function find_groundstate(state, H, alg::UnionAlg, envs = environments(state, H)
     state, envs = find_groundstate(state, H, alg.alg1, envs)
     return find_groundstate(state, H, alg.alg2, envs)
 end
+
+function find_groundstate!(state, H, alg::UnionAlg, envs = environments(state, H))
+    state, envs = find_groundstate!(state, H, alg.alg1, envs)
+    return find_groundstate!(state, H, alg.alg2, envs)
+end