Efficiently obtain the diagonal from a contraction

[jaredjeya]

Suppose I have a calculation such as,

result = diag(contract(dag(V'), B, C, V))

One example might be that I obtain the eigenvectors V by diagonalising some other ITensor, and want the expectation values with respect to the operator (B * C).

Is there a way to calculate this more efficiently than computing the intermediate matrix dag(V') * B * C * V and taking the diagonal?

For example, in numpy, one could do np.einsum(V.conj(), B, C, V, "ji,jk,kl,li->i"). However in ITensors, the best option seems to be looping over the eigen-index using onehot (which surprisingly makes a dense ITensor), e.g.

result = map(1:N) do k
    eigvec = onehot(i => k) * V
    contract(dag(eigvec)', B, C, eigvec))
end

but this is actually quite a bit slower (10x for moderately sized V, e.g. about 2000 * 2000), with many more total allocations.

This makes me think there must be a more efficient way, especially as it must be a fairly common use pattern, but I haven't been able to find one by digging through the documentation. Is there one I've missed? Any help much appreciated.

[mtfishman]

That's a good question. One way to do that is by contracting with a delta tensor but honestly our delta tensor contractions are not that great right now (I think some patterns are ok but others we just fall back to dense since they are hard to implement in an optimized way), and I think the best way to do some delta contractions may be something like the loop you wrote anyway (though implemented in a better way than onehot, which as you said it not designed in the best way right now). Do you know how that operation is actually handled in einsum internally?

Are the tensors dense? If so you could manually do a slice of V that is equivalent to what onehot is doing. I would be curious how that compares to einsum, which may give some idea of what strategy we should aim for in terms of implementation (and then how to provide that in a nice way in ITensor).

It doesn't help you right now, but note that in our ongoing rewrite of ITensors.jl (ITensorBase.jl and the surrounding packages) we will provide general slicing of ITensors so you'll be able to write something like eigvec = V[i => k, ..] which will give the same result as onehot(i => k) * V but much faster since it would call down to an Array slice operation.

[jaredjeya]

Thanks! I'm using blocksparse/QN-conserving tensors, which complicates things. (I also hadn't realised delta was similarly inefficient - I've been use of it for replacing indexes but perhaps I should be using replaceind then).

It looks like I might have to just accept this behaviour for now and wait for the rewrite, as for my purposes taking the diagonal is still decently fast.

I'm also not actually sure how np.einsum (or the optimised version at opt_einsum, https://github.com/dgasmith/opt_einsum) handles things, but that could be something to look into for sure.

[mtfishman]

For replacing indices it is fine and optimized (we use it for that purpose as well), for that case it basically just replaces the indices. But for delta tensor contractions that aren't just index replacements (i.e. delta tensors beyond order 2) it is a little hit-and-miss, and for QNs I don't think arbitrary delta tensors beyond order 2 are even well defined.

[mtfishman]

QNs definitely complicates things, slicing also becomes more subtle, which is honestly part of the reason behind the current design of onehot, we like to try to keep code generic between QNs and non-QNs and that was the easiest thing to do for QNs with the current design. We're rethinking the QN design and the ITensor abstraction as part of the rewrite which makes slicing easier to implement, though it is still fundamentally more complicated when QNs are involved (and becomes even more complicated for non-abelian symmetries, which is also something we are working on as part of the rewrite).