Cleanup and more testing

Author: mtfishman

src/AlgorithmsInterfaceExtensions/AlgorithmsInterfaceExtensions.jl

@@ -196,7 +196,8 @@ An algorithm that consists of running an algorithm at each iteration
 from a list of stored algorithms.
 =#
 @kwdef struct DefaultNestedAlgorithm{
-        Algorithms <: AbstractVector{<:Algorithm},
+        ChildAlgorithm <: Algorithm,
+        Algorithms <: AbstractVector{ChildAlgorithm},
         StoppingCriterion <: AI.StoppingCriterion,
     } <: NestedAlgorithm
     algorithms::Algorithms
@@ -256,7 +257,8 @@ function AI.step!(
 end
 
 @kwdef struct DefaultFlattenedAlgorithm{
-        Algorithms <: AbstractVector{<:Algorithm},
+        ChildAlgorithm <: Algorithm,
+        Algorithms <: AbstractVector{ChildAlgorithm},
         StoppingCriterion <: AI.StoppingCriterion,
     } <: FlattenedAlgorithm
     algorithms::Algorithms
@@ -281,18 +283,19 @@ end
 
 # Algorithm that only performs a single step.
 abstract type NonIterativeAlgorithm <: Algorithm end
+abstract type NonIterativeAlgorithmState <: State end
 
-function Base.getproperty(algorithm::NonIterativeAlgorithm, name::Symbol)
-    if name ≡ :stopping_criterion
-        return AI.StopAfterIteration(1)
-    else
-        return getfield(algorithm, name)
-    end
+function AI.initialize_state(problem::Problem, algorithm::NonIterativeAlgorithm; kwargs...)
+    return DefaultNonIterativeAlgorithmState(; kwargs...)
+end
+function AI.solve!(
+        problem::Problem, algorithm::NonIterativeAlgorithm, state::State; kwargs...
+    )
+    return throw(MethodError(AI.solve!, (problem, algorithm, state)))
 end
 
-abstract type NonIterativeAlgorithmState <: State end
-
-mutable struct DefaultNonIterativeAlgorithmState{Iterate} <: NonIterativeAlgorithmState
+@kwdef mutable struct DefaultNonIterativeAlgorithmState{Iterate} <:
+    NonIterativeAlgorithmState
     iterate::Iterate
 end
 

src/ITensorNetworksNext.jl

@@ -6,7 +6,6 @@ include("abstracttensornetwork.jl")
 include("tensornetwork.jl")
 include("TensorNetworkGenerators/TensorNetworkGenerators.jl")
 include("contract_network.jl")
-include("sweeping/sweeping.jl")
 include("sweeping/eigenproblem.jl")
 
 end

src/sweeping/eigenproblem.jl

@@ -1,63 +1,62 @@
 import AlgorithmsInterface as AI
 import .AlgorithmsInterfaceExtensions as AIE
 
-maybe_fill(value, len::Int) = fill(value, len)
-function maybe_fill(v::AbstractVector, len::Int)
-    @assert length(v) == len
-    return v
-end
+#=
+    EigenProblem(operator)
 
-function dmrg_sweep(operator, algorithm, state)
-    problem = select_problem(dmrg_sweep, operator, algorithm, state)
-    return AI.solve(problem, algorithm; iterate = state).iterate
+Represents the problem we are trying to solve and minimal algorithm-independent
+information, so for an eigenproblem it is the operator we want the eigenvector of.
+=#
+struct EigenProblem{Operator} <: AIE.Problem
+    operator::Operator
 end
-function dmrg_sweep(operator, state; kwargs...)
-    algorithm = select_algorithm(dmrg_sweep, operator, state; kwargs...)
-    return dmrg_sweep(operator, algorithm, state)
+
+struct EigsolveRegion{R, Kwargs <: NamedTuple} <: AIE.NonIterativeAlgorithm
+    region::R
+    kwargs::Kwargs
 end
+EigsolveRegion(region; kwargs...) = EigsolveRegion(region, (; kwargs...))
 
-function select_problem(::typeof(dmrg_sweep), operator, algorithm, state)
-    return EigenProblem(operator)
+function AI.solve!(
+        problem::EigenProblem, algorithm::EigsolveRegion, state::AIE.State; kwargs...
+    )
+    return error("EigsolveRegion step for EigenProblem not implemented yet.")
 end
-function select_algorithm(::typeof(dmrg_sweep), operator, state; regions, region_kwargs)
-    region_kwargs′ = maybe_fill(region_kwargs, length(regions))
-    return Sweep(length(regions)) do i
-        return Returns(Region(regions[i]; region_kwargs′[i]...))
-    end
+
+maybe_fill(value, len::Int) = fill(value, len)
+function maybe_fill(v::AbstractVector, len::Int)
+    @assert length(v) == len
+    return v
 end
 
 function dmrg(operator, algorithm, state)
-    problem = select_problem(dmrg, operator, algorithm, state)
+    problem = EigenProblem(operator)
     return AI.solve(problem, algorithm; iterate = state).iterate
 end
 function dmrg(operator, state; kwargs...)
+    problem = EigenProblem(operator)
     algorithm = select_algorithm(dmrg, operator, state; kwargs...)
-    return dmrg(operator, algorithm, state)
+    return AI.solve(problem, algorithm; iterate = state).iterate
 end
 
-function select_problem(::typeof(dmrg), operator, algorithm, state)
-    return EigenProblem(operator)
+function repeat_last(v::AbstractVector, len::Int)
+    length(v) ≥ len && return v
+    return [v; fill(v[end], len - length(v))]
 end
-function select_algorithm(::typeof(dmrg), operator, state; nsweeps, regions, region_kwargs)
-    region_kwargs′ = maybe_fill(region_kwargs, nsweeps)
-    return Sweeping(nsweeps) do i
-        return select_algorithm(
-            dmrg_sweep, operator, state;
-            regions, region_kwargs = region_kwargs′[i],
-        )
-    end
+repeat_last(v, len::Int) = fill(v, len)
+function extend_columns(nt::NamedTuple, len::Int)
+    return NamedTuple{keys(nt)}(map(v -> repeat_last(v, len), values(nt)))
 end
-
-#=
-    EigenProblem(operator)
-
-Represents the problem we are trying to solve and minimal algorithm-independent
-information, so for an eigenproblem it is the operator we want the eigenvector of.
-=#
-struct EigenProblem{Operator} <: AIE.Problem
-    operator::Operator
+function eachrow(nt::NamedTuple, len::Int)
+    return [NamedTuple{keys(nt)}(map(v -> v[i], values(nt))) for i in 1:len]
 end
 
-function AI.step!(problem::EigenProblem, algorithm::Region, state::AIE.State; kwargs...)
-    return error("Region step for EigenProblem not implemented.")
+function select_algorithm(::typeof(dmrg), operator, state; nsweeps, regions, kwargs...)
+    extended_kwargs = extend_columns((; kwargs...), nsweeps)
+    region_kwargs = eachrow(extended_kwargs, nsweeps)
+    return AIE.nested_algorithm(nsweeps) do i
+        return AIE.nested_algorithm(length(regions)) do j
+            return EigsolveRegion(regions[j]; region_kwargs[i]...)
+        end
+    end
 end

src/sweeping/sweeping.jl

@@ -0,0 +1,34 @@
+import AlgorithmsInterface as AI
+using ITensorNetworksNext: EigsolveRegion, dmrg, select_algorithm
+import ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
+using Test: @test, @testset
+
+@testset "select_algorithm(dmrg, ...)" begin
+    operator = "operator"
+    init = "init"
+    nsweeps = 3
+    regions = ["region1", "region2"]
+    maxdim = [10, 20]
+    cutoff = 1.0e-7
+    algorithm = select_algorithm(dmrg, operator, init; nsweeps, regions, maxdim, cutoff)
+    @test algorithm isa AIE.NestedAlgorithm
+    @test length(algorithm.algorithms) == nsweeps
+
+    maxdims = [10, 20, 20]
+    cutoffs = [1.0e-7, 1.0e-7, 1.0e-7]
+    algorithm′ = AIE.nested_algorithm(nsweeps) do i
+        return AIE.nested_algorithm(length(regions)) do j
+            return EigsolveRegion(
+                regions[j];
+                maxdim = maxdims[i],
+                cutoff = cutoffs[i],
+            )
+        end
+    end
+    for i in 1:nsweeps
+        for j in 1:length(regions)
+            @test algorithm.algorithms[i].algorithms[j] ==
+                algorithm′.algorithms[i].algorithms[j]
+        end
+    end
+end

test/test_dmrg.jl

@@ -1,33 +1,65 @@
 import AlgorithmsInterface as AI
-using ITensorNetworksNext: Region, Sweep, Sweeping
 import ITensorNetworksNext.AlgorithmsInterfaceExtensions as AIE
 using Test: @test, @testset
 
 struct TestProblem <: AIE.Problem
 end
 
-function AI.step!(problem::TestProblem, algorithm::Region, state::AIE.State; kwargs...)
-    state.iterate = algorithm.region
+struct TestRegion{R, Kwargs <: NamedTuple} <: AIE.NonIterativeAlgorithm
+    region::R
+    kwargs::Kwargs
+end
+TestRegion(region; kwargs...) = TestRegion(region, (; kwargs...))
+
+function AI.solve!(problem::TestProblem, algorithm::TestRegion, state::AIE.State; kwargs...)
+    new_iterate = (; algorithm.region, algorithm.kwargs.foo, algorithm.kwargs.bar)
+    state.iterate = [state.iterate; [new_iterate]]
     return state
 end
 
 @testset "Sweeping" begin
-    @testset "Region" begin
-        algorithm = Region("region"; foo = 1, bar = 2)
+    @testset "TestRegion" begin
+        algorithm = TestRegion("region"; foo = 1, bar = 2)
         @test algorithm isa AIE.NonIterativeAlgorithm
         @test algorithm isa AIE.Algorithm
         @test algorithm isa AI.Algorithm
         @test algorithm.region == "region"
         @test algorithm.kwargs == (; foo = 1, bar = 2)
-        @test Region(; region = "region", foo = 1, bar = 2) == algorithm
 
         problem = TestProblem()
-        iterate = ""
+        iterate = []
         state = AI.solve(problem, algorithm; iterate)
-        @test state.iterate == "region"
+        @test state.iterate == [(; region = "region", foo = 1, bar = 2)]
     end
     @testset "Sweep" begin
+        algorithm = AIE.nested_algorithm(3) do i
+            return TestRegion("region$i"; foo = i, bar = 2i)
+        end
+        problem = TestProblem()
+        iterate = []
+        state = AI.solve(problem, algorithm; iterate)
+        @test state.iterate == [
+            (; region = "region1", foo = 1, bar = 2),
+            (; region = "region2", foo = 2, bar = 4),
+            (; region = "region3", foo = 3, bar = 6),
+        ]
     end
     @testset "Sweeping" begin
+        algorithm = AIE.nested_algorithm(2) do i
+            AIE.nested_algorithm(3) do j
+                return TestRegion("sweep$i, region$j"; foo = (i, j), bar = (2i, 2j))
+            end
+        end
+        problem = TestProblem()
+        iterate = []
+        state = AI.solve(problem, algorithm; iterate)
+        @test state.iterate == [
+            (; region = "sweep1, region1", foo = (1, 1), bar = (2, 2)),
+            (; region = "sweep1, region2", foo = (1, 2), bar = (2, 4)),
+            (; region = "sweep1, region3", foo = (1, 3), bar = (2, 6)),
+            (; region = "sweep2, region1", foo = (2, 1), bar = (4, 2)),
+            (; region = "sweep2, region2", foo = (2, 2), bar = (4, 4)),
+            (; region = "sweep2, region3", foo = (2, 3), bar = (4, 6)),
+        ]
     end
 end