Network solvers redesign #242
Conversation
|
Fyi, At least one test was failing on my machine, but it didn't seem to be related at all to the code in this PR. |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #242 +/- ##
==========================================
- Coverage 77.39% 0.00% -77.40%
==========================================
Files 77 86 +9
Lines 3805 4018 +213
==========================================
- Hits 2945 0 -2945
- Misses 860 4018 +3158
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
|
Is the version check meaning I should just bump the version number? Which version number should I bump to? |
The version check is checking whether or not you've bumped the package version here: ITensorNetworks.jl/Project.toml Line 4 in 2d7db49 |
| @@ -0,0 +1,103 @@ | |||
| using Printf: @printf | |||
| import ConstructionBase: setproperties | |||
There was a problem hiding this comment.
| import ConstructionBase: setproperties | |
| using ConstructionBase: setproperties |
The difference between import and using in cases like this is that import would allow you to overload setproperties without prepending with the module ConstructionBase, which is a style I avoid now anyway. So in general you should use using instead of import.
There was a problem hiding this comment.
That's a helpful reminder – I had become a bit fuzzy on the reason for preferring using. I'll switch these over and similarly in all other files.
| @preserve_graph psi[v] = C | ||
| psi = set_orthogonal_region ? set_ortho_region(psi, [v]) : psi | ||
| normalize && @preserve_graph psi[v] = psi[v] / norm(psi[v]) | ||
| return setproperties(problem; state=psi) |
There was a problem hiding this comment.
I would prefer using the style set_state(problem, psi) in general, setproperties doesn't seem to add much here. Though functions like set_state could be implemented using setproperties or as @set problem.state = psi using https://github.com/JuliaObjects/Accessors.jl.
I realize I suggested using setproperties but I was picturing it would be used to help implement more specialized setter functions like set_state, or used in cases where multiple fields are being set. If you are just setting one field I don't see it adding much value to use setproperties directly in the codebase.
There was a problem hiding this comment.
Sounds good. I think my reason was just to save a few lines of code (i.e. not even define the "setters") in the files with the problem objects, but it's only 2-3 lines of code and probably a good idea for those to have setters defined anyway.
| end | ||
|
|
||
| eigenvalue(E::EigsolveProblem) = E.eigenvalue | ||
| ITensorNetworks.state(E::EigsolveProblem) = E.state |
There was a problem hiding this comment.
| ITensorNetworks.state(E::EigsolveProblem) = E.state | |
| state(E::EigsolveProblem) = E.state |
There was a problem hiding this comment.
I was getting an error if I didn't prepend ITensorNetworks here, which is why I put it. I'll find out the exact error and maybe we can fix the issue more at the root of what's causing it.
| return E, local_state | ||
| end | ||
|
|
||
| function eigsolve_sweep_printer(region_iterator; outputlevel, sweep, nsweeps, kws...) |
There was a problem hiding this comment.
What about calling this sweep_printer and have it dispatch on the iterator or problem type?
Relatedly, if there are times when we don't want to pass both the iterator and problem, we could define an alias to dispatch on:
const EigSolveRegionIterator{Problem<:EigSolveProblem,RegionPlan} = RegionIterator{Problem,RegionPlan}
function sweep_printer(region_iterator::EigSolveRegionIterator; outputlevel, sweep, nsweeps, kws...)
# ...
endThere was a problem hiding this comment.
I see, yes I guess this makes sense in case some one (a developer or user) bypasses the eigsolve interface but still uses an EigSolveProblem and wants to get this printing behavior. Interesting pattern about the EigSolveRegionIterator - I'll see if that helps. (It has definitely been useful to have some way of accessing the region iterator in the printing functions and other implementation functions like subspace expansion.)
| # generates each tuple? | ||
| # | ||
|
|
||
| mutable struct TupleRegionIterator{RegionIter} |
There was a problem hiding this comment.
Maybe we can come up with a name that is more descriptive, like RegionIteratorWithKwargs. Tuple is a bit vague.
| current_time::Number = 0.0 | ||
| end | ||
|
|
||
| ITensorNetworks.state(A::ApplyExpProblem) = A.state |
There was a problem hiding this comment.
| ITensorNetworks.state(A::ApplyExpProblem) = A.state | |
| state(A::ApplyExpProblem) = A.state |
since we are in the ITensorNetworks module.
There was a problem hiding this comment.
I'll look into what error I was getting when I didn't prepend ITensorNetworks. I agree it shouldn't be needed in principle.
|
|
||
| if !isnothing(which) | ||
| S.region_iter = region_iterator( | ||
| problem(S.region_iter); sweep=S.which_sweep, current_sweep_kws... |
There was a problem hiding this comment.
| problem(S.region_iter); sweep=S.which_sweep, current_sweep_kws... | |
| problem(S); sweep=S.which_sweep, current_sweep_kws... |
I think, based on the definition of problem(::SweepIterator) above.
There was a problem hiding this comment.
Also, this code pattern confused me a bit, what do you think of writing it like this:
update_region_iterator!(S; current_sweep_kws...)and hiding the implementation in update_region_iterator!?
| region_plan_state = iterate(R.region_plan, which) | ||
| isnothing(region_plan_state) && return nothing | ||
| (current_region, region_kwargs), next = region_plan_state | ||
| R.problem = region_iterator_action(problem(R), R; region_kwargs...) |
There was a problem hiding this comment.
I'm not a big fan of the name region_iterator_action but I can't think of anything much better. The only thing I can think of is region_update but I don't know how I feel about reusing the term update. Maybe region_step (following the step! terminology of DifferentialEquations.jl: https://docs.sciml.ai/DiffEqDocs/stable/basics/integrator/#Initialization-and-Stepping).
| mutable struct SweepIterator | ||
| sweep_kws | ||
| region_iter | ||
| which_sweep::Int | ||
| end |
There was a problem hiding this comment.
What about parameterizing this struct by the types of the fields sweep_kws and region_iter? Is it meant to be dynamic, i.e. is it expected that those types might change?
There was a problem hiding this comment.
It seems best to have a goal of having those be concrete and not changing type, but if they do you can still parameterize the type and then set the type parameter to an abstract type when it is being constructed (as needed), for example you can explicitly construct it as SweepIterator{Any,Any}(sweep_kws, region_iter, which_sweep).
|
|
||
| problem(R::RegionIterator) = R.problem | ||
| current_region_plan(R::RegionIterator) = R.region_plan[R.which_region] | ||
| current_region(R::RegionIterator) = current_region_plan(R)[1] |
There was a problem hiding this comment.
What's the data structure for the output of current_region_plan(R)? Something doesn't feel right to me that the region is accessed with indexing by 1, maybe it should be a NamedTuple and the region could be accessed as current_region_plan(R).region or a struct where it is accessed with a function get_region(current_region_plan(R))?
| current_region(R::RegionIterator) = current_region_plan(R)[1] | ||
| region_kwargs(R::RegionIterator) = current_region_plan(R)[2] | ||
| function previous_region(R::RegionIterator) | ||
| R.which_region==1 ? nothing : R.region_plan[R.which_region - 1][1] |
There was a problem hiding this comment.
It seems like maybe we should be using function accessors rather than field accessors, i.e. which_region(R) instead of R.which_region.
| default_region_callback(problem; kws...) = nothing | ||
|
|
||
| default_sweep_callback(problem; kws...) = nothing | ||
|
|
||
| function default_sweep_printer(problem; outputlevel, sweep, nsweeps, kws...) |
There was a problem hiding this comment.
Similar to my comment above, what about calling these just region_callback, sweep_callback, and sweep_printer, since they would get overloaded on the problem type anyway?
| for (sweep, region_iter) in enumerate(sweep_iterator) | ||
| for (region, region_kwargs) in region_tuples(region_iter) | ||
| region_callback( | ||
| problem(region_iter); |
There was a problem hiding this comment.
I noticed this gets passed the problem while sweep_callback and sweep_printer get passed the region_iter, is there a logic behind that? I would lean towards all of them getting passed the region_iter, and there could be a layer that unwraps the problem as needed.
| @@ -0,0 +1,64 @@ | |||
| using Test: @test, @testset | |||
| using ITensors | |||
| using TensorOperations # Needed to use contraction order finding | |||
There was a problem hiding this comment.
| using TensorOperations # Needed to use contraction order finding | |
| using TensorOperations: TensorOperations # Needed to use contraction order finding |
That only brings TensorOperations itself into the namespace.
| import Graphs as gr | ||
| import NamedGraphs as ng |
There was a problem hiding this comment.
For the sake of keeping the style consistent throughout the library I would prefer:
using Graphs: add_edge!, ...
using NamedGraphs: NamedGraph, ...If we don't stay consistent with the style it leads to chaos where code is written in different styles and it is a lot of work to go back and make it consistent later.
There was a problem hiding this comment.
Yes, I meant to change these to the current style of ITensorNetworks. I will do that and in any other files I find.
| @@ -0,0 +1,17 @@ | |||
| import ConstructionBase: setproperties | |||
|
|
|||
| function extracter(problem, region_iterator; sweep, trunc=(;), kws...) | |||
There was a problem hiding this comment.
If we don't change this to extract as discussed elsewhere, which should change this to extractor since extracter isn't the correct spelling.
| function applyexp( | ||
| init_prob, | ||
| exponents; | ||
| extracter_kwargs=(;), |
There was a problem hiding this comment.
| extracter_kwargs=(;), | |
| extractor_kwargs=(;), |
purely spelling, since extractor is a word but extracter isn't.
| @kwdef mutable struct ApplyExpProblem{State} | ||
| state::State | ||
| operator | ||
| current_time::Number = 0.0 | ||
| end |
There was a problem hiding this comment.
Related to my comment above about SweepIterator, I think we should parameterize the problem types. Also related to my other comment I think we should switch the argument ordering so the operator is first.
Also, I'm suspicious when I see 0.0 since that is a Float64. If the time is Float64 it might then promote the element type of other objects it interacts with to double precision, which would be bad on GPU where single precision is better. We should have a constructor that determines the number type from the inputs, such as: ApplyExpProblem(operator, state) = ApplyExpProblem(operator, state, zero(promote_type(scalartype(operator), scalartype(state))).
| @@ -0,0 +1,14 @@ | |||
| default_maxdim() = typemax(Int) | |||
| default_mindim() = 1 | |||
| default_cutoff() = 0.0 | |||
There was a problem hiding this comment.
| default_cutoff() = 0.0 | |
| default_cutoff(elt::Type{<:Real}) = zero(elt) |
so that we don't assume double precision (Float64).
There was a problem hiding this comment.
We could then define default_cutoff() = default_cutoff(Float64) but I would prefer not to right now and instead make sure to propagate specific element types that are determined from the inputs.
| (norm(Y) <= 1E-15 || expand_maxdim <= 0) && return local_tensor | ||
| Ux, S, V = svd(Y, basis_inds; cutoff=1E-14, maxdim=expand_maxdim, lefttags="ux,Link") |
There was a problem hiding this comment.
We should determine an element type/scalar type elt from the inputs and replace 1E-15 and 1E-14 with something like eps(elt) * 10 or eps(elt) * 10^2 to be more agnostic about using single precision or double precision, to be friendlier for running on GPU.
| using NDTensors.BackendSelection: Backend, @Backend_str | ||
| import ConstructionBase: setproperties | ||
|
|
||
| default_expansion_factor() = 1.5 |
There was a problem hiding this comment.
| default_expansion_factor() = 1.5 | |
| default_expansion_factor(elt::Type{<:Real}) = 3 * one(elt) / 2 |
| include("utilities/simple_ed_methods.jl") | ||
| include("utilities/tree_graphs.jl") | ||
|
|
||
| @testset "Tree DMRG" begin |
There was a problem hiding this comment.
| @testset "Tree DMRG" begin | |
| @testset "Tree DMRG (eltype=elt)" for elt in (Float32, Float64, ComplexF32) |
and then propagate elt throughout the tests, for example when constructing the operator, state, cutoff, etc. I like to set up tests where we loop over element types like this from the start to make sure the code is agnostic about the element types, it is easier to do that from the start rather than after a lot of library code and tests have been written already.
Even more advanced is then looping of array backends like CPU and GPU arrays, we can discuss how to do that. Again, it is nice to set up tests like that early on rather than trying to fix those issues and test them after the fact.
There was a problem hiding this comment.
For an example of how I also test on GPU, you can see: https://github.com/ITensor/BlockSparseArrays.jl/blob/44574b808861a03c789a74620d8308f3b8c1b8d4/test/test_basics.jl#L49-L64.
|
@emstoudenmire the reason for the failing test is that the |
| ) | ||
| end | ||
| S.which_sweep += 1 | ||
| return S.region_iter, next |
There was a problem hiding this comment.
I'm a bit on the fence about this, but I would lean towards a design where iterate(::SweepIterator) actually performs the region iteration as opposed to returning the region iterator, so then for _ in sweep_iterator end actually runs the calculation. Then, we could have an iteration adapter such as region_iterators(::SweepIterator) returns an iterator that returns the region iterator at each iteration.
There was a problem hiding this comment.
An interesting question is then, in that alternative design, what iterate(::SweepIterator) should output at each iteration. The DifferentialEquations.jl design is that it outputs the iterator itself, i.e. if you call:
for x in sweep_iterator
endx will be the latest updated sweep_iterator, which I think makes sense since then you can access the information of the sweep_iterator inside the loop.
There was a problem hiding this comment.
For reference, in DifferentialEquations.jl, iterate is just defined as:
function Base.iterate(integrator::DEIntegrator, state = 0)
done(integrator) && return nothing
state += 1
step!(integrator)
return integrator, state
endwhich seems like a reasonable goal to aim for, and then the complexity of the implementation is in done(...) and step!(...).
|
|
||
| mutable struct SweepIterator | ||
| sweep_kws | ||
| region_iter |
There was a problem hiding this comment.
Maybe this should be called current_region_iter or something like that to indicate it is the latest region iterator.
I have to say this part confused me a bit. I see the region iterator is recreated from scratch from the sweep_kws at each sweep, is there a reason to store it? It seems like it could just be made and used "on the fly" at each sweep anyway. Instead it seems like we could just store the problem in the SweepIterator.
| # | ||
|
|
||
| mutable struct SweepIterator | ||
| sweep_kws |
There was a problem hiding this comment.
Maybe this could be named each_sweep_kwargs or sweep_kwargs_iterator to indicate that it itself is an iterator that returns the keyword arguments for the sweep at each iteration (rather than just the keyword arguments of the latest sweep).
Better handling of zero and vector creation Co-authored-by: Matt Fishman <[email protected]>
This PR is a rewrite of codes for "sweep solver" algorithms such as DMRG, TDVP, etc.
It introduces a simplified designs for the whole system, especially regarding the creation of "region plans", how the double loop comprising what is now called
sweep_solveis coded, and the handling of keyword arguments.The PR also adds a subspace expansion system with one currently implementation (based on a "projected density matrix perturbation" idea, basically McCulloch+White hybrid method). Having a subspace expansion is crucial to converge DMRG on certain tree lattices, not to mention other cases like 2D DMRG with QN conservation.
Internally, the core design revolves around two iterators: a sweep iterator and a region iterator. Each of these is in principle wrappable by "iteration adapters" (see for example the tuple adapter in adapters.jl). At each iteration of the RegionIterator type, it calls a
region_iterator_actionfunction which can be overloaded, but by default calls three "subactions":extracter,updater, andinserter. Currentlyextracteralso calls down into asubspace_expandfunction which can be customized through multiple backends. These "action" functions all dispatch on a "problem" type which can hold arbitrary data, making these codes more flexible and future-proof toward cases like optimizing two tensor network states at once, working with sets of operators, etc.Other improvements that may fit better into future PRs:
truncateandapplyalready in the NetworkSolvers repo (there were some issues bringing them over)dmrgwhich automatically propagate truncation arguments into more "expert" keyword argument packs