Fix assemble(slate.Tensor, diagonal=True)

[pbrubeck]

Description

[connorjward]

We have V, here to assert a singleton. Why is that now violated? What are the contents of V[1:]?

[pbrubeck]

With slate we get the two copies of the same space in Vs. For some reason diagonal=True will cause len(indices) == 1 only for ufl.Form, but not for slate.Tensor. On both cases len(form.arguments()) == 2.

firedrake/firedrake/assemble.py

Line 2214 in 32611a7

return tuple(a.ufl_function_space()[i] for i, a in zip(indices, form.arguments()))

[pbrubeck]

And this is where we get the indices from

firedrake/firedrake/assemble.py

Line 1653 in 32611a7

self._indices, self._kinfo = local_knl

Which means that we should trim the lenght of indices somewhere in slate/slac/compiler.py, but it gets hard to track down because diagonal is not mentioned at all in that script.

[connorjward]

Could you use V = pytools.single_valued(Vs) then?

[pbrubeck]

As I was adding tests, I saw that there is support for assemble(Tensor(form, diagonal=True)), which assembles directly the diagonal. But this PR aims to enable assemble(Tensor(form), diagonal=True) to assemble the diagonal from the Matrix entries. Unfortunately the latter gives incorrect results.

In general we cannot rely on Tensor(form, diagonal=True), as we would like to support diagonals of more complicate SLATE expressions (e.g. Schur complements). So this PR would require more work.