Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement sparse SVD with function handles #72

Merged
merged 21 commits into from
Oct 29, 2024
Merged

Implement sparse SVD with function handles #72

merged 21 commits into from
Oct 29, 2024

Conversation

pbrehmer
Copy link
Collaborator

@pbrehmer pbrehmer commented Oct 2, 2024

In this PR we will use function handles to perform the iterative SVD of the half-infinite environment in CTMRG. To that end, a struct will be implemented that stores all necessary tensors for the environment contraction and on which tsvd! can dispatch.

Before doing that, a long-standing issue will be resolved where IterSVD would lead to singular vectors which, when supplied to FixedSVD, would not lead to element-wise converged environments and thereby making it not possible to use the iterative SVD during PEPS optimization.

Copy link

codecov bot commented Oct 2, 2024

Codecov Report

Attention: Patch coverage is 60.00000% with 50 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
src/algorithms/ctmrg/sparse_environments.jl 45.45% 24 Missing ⚠️
src/algorithms/contractions/ctmrg_contractions.jl 47.36% 20 Missing ⚠️
src/algorithms/ctmrg/ctmrg.jl 80.00% 4 Missing ⚠️
src/utility/svd.jl 91.30% 2 Missing ⚠️
Files with missing lines Coverage Δ
src/PEPSKit.jl 100.00% <ø> (ø)
src/algorithms/peps_opt.jl 96.29% <ø> (+0.84%) ⬆️
src/utility/svd.jl 88.52% <91.30%> (+0.19%) ⬆️
src/algorithms/ctmrg/ctmrg.jl 91.39% <80.00%> (-3.37%) ⬇️
src/algorithms/contractions/ctmrg_contractions.jl 69.23% <47.36%> (-30.77%) ⬇️
src/algorithms/ctmrg/sparse_environments.jl 45.45% <45.45%> (ø)

... and 1 file with indirect coverage changes

@lkdvos
Copy link
Member

lkdvos commented Oct 2, 2024

I think if you change the initial guess to be a slight deviation from the first singular vector, you'll have better performance. I think you need a small component orthogonal to that vector to start generating new vectors in the krylov subspaces.

@pbrehmer
Copy link
Collaborator Author

pbrehmer commented Oct 4, 2024

So in principle this should be a relatively complete implementation of "function handle CTMRG", however it's still untested and there are some things to discuss/fix:

  • somehow the way the adjoint action of HalfInfiniteEnv is defined leads to errors when calling tsvd! and I haven't figured out why
  • implementing this took quite a bit of extra compatibility methods and some rewriting, so I'm sure one can improve on some of the design choices

Maybe we could discuss, review and fix these things next week or so @lkdvos.

Copy link
Member

@lkdvos lkdvos left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I left some specific remarks in the code, but maybe let me summarize:

  • We should probably be a little careful with trying to push too many things into the type domain. The runtime benefit is not obvious, as it would just compile away a couple if checks, and the compiletime will increase, as well as the code complexity.
  • The main source of misery here seems to be that we should rethink a good interface for svdsolve with both TensorMaps and Functions. I will attempt to come up with something better (probably tomorrow, otherwise coming week)

As a more general remark, I do really like the sparse environment structs, and would probably structure the code to always construct them first, and then, depending on the algorithm, either convert to a dense environment TensorMap or not. This would probably (hopefully) organise a bit of the code a bit more, such that the responsability of dealing with these objects is centralized to these structs.

src/algorithms/ctmrg/ctmrg.jl Outdated Show resolved Hide resolved
Comment on lines 205 to 233
TensorKit.InnerProductStyle(::HalfInfiniteEnv) = EuclideanProduct()
TensorKit.sectortype(::HalfInfiniteEnv) = Trivial
TensorKit.storagetype(env::HalfInfiniteEnv) = storagetype(env.ket_1)
TensorKit.spacetype(env::HalfInfiniteEnv) = spacetype(env.ket_1)

function TensorKit.domain(env::HalfInfiniteEnv)
return domain(env.E_4) * domain(env.ket_2)[3] * domain(env.bra_2)[3]'
end
function TensorKit.codomain(env::HalfInfiniteEnv)
return codomain(env.E_1)[1] * domain(env.ket_1)[3]' * domain(env.bra_1)[3]
end
function TensorKit.space(env::HalfInfiniteEnv)
return codomain(env) ← domain(env)
end
function TensorKit.blocks(env::HalfInfiniteEnv)
return TensorKit.SingletonDict(Trivial() => env)
end
function TensorKit.blocksectors(::HalfInfiniteEnv)
return TensorKit.OneOrNoneIterator{Trivial}(true, Trivial())
end
function TensorKit.block(env::HalfInfiniteEnv, c::Sector)
return env
end
function TensorKit.tsvd!(f::HalfInfiniteEnv; trunc=NoTruncation(), p::Real=2, alg=IterSVD())
return _tsvd!(f, alg, trunc, p)
end
function TensorKit.MatrixAlgebra.svd!(env::HalfInfiniteEnv, args...)
return TensorKit.MatrixAlgebra.svd!(env(), args...) # Construct environment densely as fallback
end
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This somehow confuses me, I would not have guessed that if we dispatch to KrylovKit, we would need this TensorKit interop...

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I looked at these changes in a bit more detail, and I think we will probably have to commit to a better svd interface before we can make this work, as this definitely seems a bit unwieldy.
I think the main issue comes from the fact that we ended up merging the "compute pullback of svd block" step with the "merge pullbacks back into TensorMap form" step. I have some time on the flight tomorrow, I'll try and see if I cannot make this work a bit cleaner.

The idea would be something like: if we can have a stack-like function for TensorMap's, which is compatible with AD, it should be possible to use the KrylovKit rrule for the first step described above, while letting the AD engine combine the pullbacks automatically. This should then avoid this code-duplication, and make it more robust for function handles as well.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, I definitely agree. The way the interface is designed right now makes things tricky but there needs to be some compatibility layer that translates between the svdsolve calls and the block TensorMap structure that also needs to be generic enough to somehow work for function handles. I'd really welcome a better solution than what we have right now.

@@ -0,0 +1,265 @@
const SparseCTMRG = CTMRG{M,true} where {M} # TODO: Is this really a good way to dispatch on the sparse CTMRG methods?
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would probably put this as a regular field. I am starting to feel less and less happy with all our type parameters, as it makes things a lot less transparent...

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As a more general remark, I do really like the sparse environment structs, and would probably structure the code to always construct them first, and then, depending on the algorithm, either convert to a dense environment TensorMap or not. This would probably (hopefully) organise a bit of the code a bit more, such that the responsability of dealing with these objects is centralized to these structs.

When I was writing these struct I was also thinking that it would be more useful to center the CTMRG function around them. That would definitely organize the code a bit more naturally. I'll get to it next week!

@pbrehmer
Copy link
Collaborator Author

pbrehmer commented Oct 7, 2024

Had to add some extra renormalize_corner methods to make EnlargedCorner work with CTMRG. For some reason CTMRG doesn't want to be differentiated anymore (...well when does AD ever work right away anyway?) Tried to find the problem, but no success yet

Edit: Is this related to the rrule of the EnlargedCorner functor? If so, how does the functor rrule actually look like?

@pbrehmer
Copy link
Collaborator Author

Wow, only now I figured out that the problem with the iterative SVD and the :fixed mode leading to non-converging environments was due to wrong daggers in the relative gauge fixing contractions. (Big thanks to Bram on this one!)

Weirdly enough, the TensorKit tsvd and iterative SVDs with weird initial guesses retain some sort of symmetry, where the relative phases are just $\pm 1$; if, however, the initial guess is random, the relative phases will be complex which led the gauge fixing procedure to fail since the daggers were in the wrong place... Anyways, this is quite an important fix!

@pbrehmer
Copy link
Collaborator Author

This should be finished up to the function handle SVD aspects - but those we'll anyway implement later on. Once the KrylovKit/TensorKit SVD interface is implemented, the function handles should also be very easy to add.

@pbrehmer pbrehmer requested a review from lkdvos October 25, 2024 08:45
Copy link
Member

@lkdvos lkdvos left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have just a few small comments, otherwise I think I'm okay with merging.

src/algorithms/ctmrg/ctmrg.jl Outdated Show resolved Hide resolved
src/algorithms/ctmrg/sparse_environments.jl Outdated Show resolved Hide resolved
src/utility/svd.jl Outdated Show resolved Hide resolved
@pbrehmer pbrehmer merged commit c8c4f11 into master Oct 29, 2024
8 of 9 checks passed
@pbrehmer pbrehmer deleted the pb-sparse-svd branch October 29, 2024 08:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants