[WIP] Environment initialization for CTMRG

[leburgel]

Attempt at some better environment initialization routines for CTMRG. Should solve #255 once it's finished and documented.

This implements an initialize_environment method which initializes a CTMRGEnv for a given network according to a specific InitializationStyle. I specified three kinds of InitializationStyles: RandomInitialization for initializing a random environment, ProductStateInitialization for initializing a (potentially embedded) product state environment, and ApplicationInitialization which grows an initial environment from a product state according to a given truncation scheme. Names, arguments and which ones are optional for discussion, but at least this already handles the test case from #255 in a fairly straightforward way.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
src/PEPSKit.jl	100.00% <ø> (ø)
src/algorithms/ctmrg/ctmrg.jl	90.62% <100.00%> (+0.14%)	⬆️
src/algorithms/ctmrg/initialization.jl	100.00% <100.00%> (ø)

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lkdvos]

I think I would imagine this being a part of the ApplicationInitialization struct, and the same for the trscheme?

[leburgel]

For the trscheme, I kind of think of this as equivalent to the virtual space specification in the other schemes, which is why I kept it separate. See the comment below for some more explanation on that.

For the algorithm being part of the ApplicationInitialization struct, that was also my first thought, but it really wasn't clear to me how to do this in a simple way. Part of the reason is that I would like to reuse this for other kinds of initializations (e.g. PEPS fixed point initialization for Hamiltonians or PEPOs, using different kinds of algorithms such as SimpleUpdate or variational truncation). So it didn't really make sense to have constructors for ApplicationInitialization which put in a default algorithm since the possible algorithm really depend on the objects we're dealing with. But at the same time I really wanted a simple calling signature with a default algorithm that users wouldn't have to manually construct first. So I kind of like the black box behavior now, but I'm also not too happy with the implementation

If we make the algorithm a part of the struct,

struct ApplicationInitialization{A} <: InitializationStyle
    alg::A
end

we could get in the algorithm based on the object here through something like

function ApplicationInitialization(network::InfiniteSquareNetwork; kwargs...)
    return ApplicationIntialization(CTMRGAlgorithm(; kwargs...))
end

Then I just move the defaults there, and it cleans up the initialize_environment signature. Would that be better?

[lkdvos]

Having a look at the argument order, I think something feels a little off with the virtual_spaces as final argument to me.
(Disclaimer that everything that follows is subjective and therefore please do tell me to shut up if you disagree)

In general, most of our "algorithm functions" actually follow a very similar type of interface, based on the purpose of the arguments:

algorithm_function(init, operator, algorithm, [additional_optional_structures])

With this in mind, I think I would try to make something like this work:

initialize_environment([spaces], network, alg)
initialize_environment([(T, spaces)], network, alg)

but I wonder if there is a more elegant solution for the T, possibly just putting it in the keyword arguments. I don't think there are many occasions where T != scalartype(network) is desired, and we can try to make this at least @constinferred?

[leburgel]

I wouldn't mind changing T to a keyword argument, the only reason I added it like this was to stay closer to the environment and state constructor signatures.

The choice of position for the virtual_spaces argument was purely pragmatic, as this is basically the only way I could make it compatible with the way the space specification in the CTMRGEnv constructor works. Since the virtual space specification for the different directions works with optional positional arguments, putting virtual_spaces as the last positional argument was the only way to allow for more elaborate virtual space specifications as far as I could tell.

If we want to move it, we would probably need some kind of CTMRGEnvSpaces structure which contains all the information on the virtual spaces in a single object that we can then put wherever we want. This is anyway something we might want to get rid of signatures like CTMRGEnv(ComplexF64, randn, network, virtual_spaces...) in favor of something like randn(ComplexF64, spaces::CTMRGEnvSpaces). That's also part of the reason I introduced these _fill_environment_virtual_spaces methods in #266, since these would come in handy when defining CTMRGEnvSpaces(network, virtual_spaces...). We can decide to do this first, and at the same time change how we internally specify all other space specification.

Starting from the assumption that virtual_spaces needs (at least right now) to be the last positional argument, my reasoning for the argument order was to have something like

initialize_environment([T,] network, alg, virtual_space_specification)

The way you would specify a space depends on the algorithm. For product state or random initialization, the virtual space specification is just an actual space specification as it would be passed to the constructor, hence the whole reason the space specification is positioned last. For an ApplicationInitialization, a full space specification doesn't really make sense, what matters then is how the virtual spaces are truncated during the application iterations. That's why I chose to specify spaces here using a TruncationScheme, and also why the trscheme was the last positional argument in the initialize_environment signature for ApplicationInitialization.

[Yue-Zhengyuan]

I just noticed one more way to initialize the environment (I already have some code on this in my neighborhood tensor update branch), which we may call Neighbo(u)rhoodInitialization. This was mentioned in 1104.5463:

Initially, the corner and edge tensors at position [x,y] are constructed similarly as a reduced tensor a[x,y], by multiplying the tensor A[x,y] to its conjugate and fusing the bond indices,9 where we trace over the legs directed toward an open boundary.

For example, an edge tensor is constructed like this:

           2    
         ╱      
┌-----ket----- 4
|    ╱ |        
|   6  |   1    
|      | ╱      
└-----bra----- 3
     ╱          
    5           

In this way, the initialized environment directly have boundary dimension $\chi = D^2$.

Recently I realized that having a good initialization is also important during full update. Right now, I'm using the environment from the last iteration as initialization for the next step. But sometimes even this doesn't work well. So I kind of hope that we can finish this PR soon.

[lkdvos]

Would this give a different result from starting from environments that are the identity and doing a single step of ctmrg without truncating?

[Yue-Zhengyuan]

Nice insight, they may indeed be the same thing... for fermions we may need to further ensure that the twists are correctly added.

[Yue-Zhengyuan]

We definitely need a test with fermionic iPEPS on this, in case parity-odd elements become -1 due to twists.

[leburgel]

Given the fact that this PR might be useful for #321, I'll try to pick this up again now. However, considering the original purpose, the review comments, and the new use case in #321, it's not entirely clear to me where to go with this. Let me try to summarize.

The original motivation for this PR was to solve #255, by making it 'easy' to initialize a CTMRG environment with a given bond dimension/virtual space by growing it from a smaller environment (with a product state being the default starting point. Basically, package up what people do all the time in a single method. The easiest way to do this would be to write a specialized initialize_ctmrg_environment.

However, since initializing environments is something very common, I tried to make it a bit more generic so that the same signature and initialization flavors could be reused in other places. This came out a bit awkward right away, as can be seen from the review comments. A first obvious problem is that the specification of virtual environment spaces really depends on what exactly we're doing, which gives rise to some awkward argument orderings. In addition, to actually 'grow' an environment we need to apply a few iterations of an algorithm (see ApplicationInitialization), which requires some algorithm settings to be specified. Again, this all made sense in the context of the original application to CTMRG, but looks very arbitrary when seen in a generic context.

Now in #321 there's an additional aspect, where we actually want to initialize a different kind of environment depending on the contraction algorithm we use. For now the C4v implementation uses the normal CTMRGEnv struct, but it would probably make more sense to use a custom C4vCTMRGEnv. Anticipating that we also want to support boundary MPS contractions soon™, this gives more different environment types for more contraction algorithms.

Trying to consider all of this feels a bit annoying, which is also why this PR got completely stalled in the first place. To get this moving, I think the first thing to decide is whether we want to solve the whole (potentially still growing) problem here, or we first just address the original issue #255.

To address the original problem, I would switch to just writing an initialize_ctmrg_environment with the minimal functionality required to address the issue.

To go for a full generic initialize_environment, I think it becomes inevitable to include an actual contraction algorithm instance in the arguments. Either this, or an equivalent named tuple that can identify multiple different contraction algorithms (and not just flavors of CTMRG like we support now). That would lead to something like

initialize_environment([T,] network, contraction_algorithm, initialization_style, virtual_space_specification)

for the generic case withing PEPSKit.jl alone. The question then becomes if this is compatible with other kinds of environment initializations (e.g. in MPSKit.jl), and if we still want it to be.

Asking for any and all opinions from @lkdvos, @pbrehmer and @Yue-Zhengyuan, after which I'll have a go at getting this to a point where it's usable.