Clean up symmetrization and add optimization callback

[pbrehmer]

This PR cleans up the symmetrization functionality a bit and adds a callback keyword argument for fixedpoint which allows to call a user-specified function after each optimization iteration. Also, I added the J1-J2 model Hamiltonian and a small test since I anyway had that code lying around.

[pbrehmer]

@lkdvos What do you think about the callback function in fixedpoint? I started to need such a thing, e.g. to symmetrize the optimized PEPS and its gradient, but it can also be used for diagnostic purposes, advanced printing, etc. However, there might be a better way do this.

Also I didn't really spend time to rewrite the symmetrization, so I just refactored it a bit to make it more readable. It is definitely usable for now so I'd be fine with just merging it like this. I probably don't have the time currently to actually improve it and add the rotation symmetrization.

[codecov]

Codecov Report

Attention: Patch coverage is 81.57895% with 14 lines in your changes missing coverage. Please review.

Files	Patch %	Lines
src/utility/symmetrization.jl	77.08%	11 Missing ⚠️
src/states/infinitepeps.jl	50.00%	3 Missing ⚠️

Files	Coverage Δ
src/PEPSKit.jl	100.00% <ø> (ø)
src/algorithms/peps_opt.jl	95.23% <100.00%> (+0.04%)	⬆️
src/operators/models.jl	100.00% <100.00%> (ø)
src/states/infinitepeps.jl	69.62% <50.00%> (+2.49%)	⬆️
src/utility/symmetrization.jl	80.85% <77.08%> (+13.80%)	⬆️

... and 1 file with indirect coverage changes

[lkdvos]

I think this all looks okay to me. I would call the callback kwarg finalize, in analogy to how OptimKit and MPSKit already have this.
This does give a somewhat interesting question: do you want to have the callback/finalizer run every iteration -- optimization step, or every iteration -- function/gradient evaluation. In principle, both could be useful and are definitely distinct.

As a side-note, I think we will have to reconsider our optimization interface again, the current approach has to many options, some of which are very non-obvious, and some are incompatible, so it is not that straightforward to get this going. This adds another option, which I think is fine, but we should probably think a bit about that at some point.

Finally, @leburgel might not be too happy with the disappearance of _make_it_fit, I have a feeling that was his favorite function.

[pbrehmer]

As a side-note, I think we will have to reconsider our optimization interface again

I definitely agree. Currently, I'm using a wrapper struct around all the options which hides away some stuff and that makes launching optimizations definitely more usable. But on the other hand, PEPS optimizations just have a lot of parameter you have to be able to tweak in some way, so there should be a possibility to get fine-grained control over everything. Still, in its current form the interface is not super user-friendly. We should probably meet up at some point and discuss this in more detail, once we have the time.

This does give a somewhat interesting question: do you want to have the callback/finalizer run every iteration -- optimization step, or every iteration -- function/gradient evaluation.

I would also think that both scenarios are useful. There are probably cases where having either one or the other is fine. That should be a thing to think about when we come up with a new optimization interface at some point.

Finally, @leburgel might not be too happy with the disappearance of _make_it_fit, I have a feeling that was his favorite function.

Haha, I just gave it a more descriptive name, the function is still there. (I do admit that it's quite useful sometimes :P)

[leburgel]

Looks like a strict improvement to me!

Part of the motivation for calling it _make_it_fit was to make it sound as silly as possible, since this is absolutely not something you want to be doing when working with actual symmetries. The new name sounding less shady is more a con than a pro for me, but I can definitely live with it. All of the symmetrization things are a bit janky anyway in the sense that they work fine until they don't. I've been meaning to redo things such that spaces never need to be forced to match and we also have normal reflection symmetry aside from the weird hermitian reflections that are in there now. But for now this looks great.

[pbrehmer]

Part of the motivation for calling it _make_it_fit was to make it sound as silly as possible, since this is absolutely not something you want to be doing when working with actual symmetries.

You definitely have a point. While it should be fine for now, it's good to keep in mind that we should revise that part of the code at some point :-)

I've been meaning to redo things such that spaces never need to be forced to match and we also have normal reflection symmetry aside from the weird hermitian reflections that are in there now.

That would be super nice and definitely a good thing to have in the future!

By the way, I slightly changed the previous implementation of the finalize function such that it now uses that finalize keyword argument from OptimKit and that it can be assumed mutating. It seemed a bit redundant to implement an extra finalize and not use what's already there - especially since it does the same thing. I've also added a rudimentary symmetrization test.

[lkdvos]

If @leburgel is okay with the naming thing, you have my blessing. I left one remark in the code and have two more that I'll add here:

For the finalize! kwarg, I am a bit torn between having that as a separate keyword, and having it part of the PEPSOptimize alg. For consistency, it is nice that everything is in the algorithm struct, (and it might help with type stability issues), as I somehow like to have the following pattern:

run_algorithm(args...; kwargs...)
run_algorithm(args..., alg_struct) # expert mode

On the other hand, I might be okay with getting that fixed later, I really want to invest some time in getting a good grip on all the different kwargs with some good automated options there, so that could also just fit in that PR.

The second comment is about the j1j2 example as a test file. I think we could probably have either heisenberg or j1j2 as a test, and invest some time in getting a proper examples thing set up to add both. From a CI point of view, I don't think that testing both really makes that big of a difference, except increase the waiting times

[lkdvos]

Do we want to just define == for pepses? I am also not entirely convinced that unpacking the tensor data is a solid approach, that really is reaching into the tensor internals which might break at some point.

[pbrehmer]

Yes, that's definitely the better way to do it. I now defined == and isapprox methods for InfinitePEPS which just use the TensorKit methods element-wise.

[pbrehmer]

On the other hand, I might be okay with getting that fixed later, I really want to invest some time in getting a good grip on all the different kwargs with some good automated options there, so that could also just fit in that PR.

Maybe I would suggest that we do it that way since it's not quite clear yet to me how to best automate all the options. Also, I'm always a bit hesitant to put functions into an algorithm struct since that makes it hard to save them to disk (which you probably want to do somewhere in the simulation workflow). So perhaps we can tackle that in a separate PR?

I think we could probably have either heisenberg or j1j2 as a test, and invest some time in getting a proper examples thing set up to add both. From a CI point of view, I don't think that testing both really makes that big of a difference, except increase the waiting times

Sounds like a good idea. But would you in general like to keep the square_lattice_heisenberg and square_lattice_j1j2 functions separate? Technically, the Heisenberg model is just a special case of the J1-J2 model. However, I thought it would be useful to have both functions because with the Heisenberg one you can tweak $(J_x, J_y, J_z)$ whereas with the J1-J2 one you only tweak $(J_1, J_2)$ - of course you could combine all of that into a single function.

One reason to maybe keep both tests would be to check that the PEPS optimization runs fine for different parameter ranges of the same Hamiltonian (which is not always guaranteed I guess). If we decide to test only one of them it's probably best to test the J1-J2 model since that also includes next-nearest neighbor terms.

[lkdvos]

I guess the bigger question is if the different parameter regimes would fail because of an implementation detail, or just because the algorithms are not suited. The CI really is just there to check that the code does what it is advertising, and it might be that the implementations are "perfect", and still the optimization does not converge.
If it does not take too long, it's completely fine to keep them in, but we should not be using the test-suite to explore too many models I guess. I like that J1J2 tests the next-nearest neighbor terms, so that actually increases "test coverage", but tweaking parameters does not :)

[lkdvos]

left some final small comments, otherwise good to go for me

[lkdvos]

Suggested change

	g = only(g) # withgradient returns tuple of gradients `g`
	return E, g
	return E, only(g) # withgradient returns tuple of gradients `g`

While it might be a little more readable, you should be careful with re-using the same variable name for things like this. Sometimes, the compiler gets confused and can no longer figure out the type of g, especially in contexts like these with anonymous functions etc.
So I would either keep it as it was, or use E, gs for the return values of withgradient and then g = only(gs).
I don't have a preference about that :)

[lkdvos]

we can probably just define this in PEPSKit ? It feels a bit strange to overload a function within the tests

[pbrehmer]

Thanks for the comments and also for completing the PR :-) I didn't manage to do it yesterday anymore (was on a super long night train right without wi-fi).