Problem when using slate Tensor

[marazzaf]

Hi everyone,

I am trying to learn how to use Slate in order to do static condensation.

But when I try to run a very simple example, I run into issues. Here is a minimal working example.

from firedrake import *
from firedrake.slate import Tensor, AssembledVector, Solve, Block
import sys

#Mesh
N = 25
mesh = UnitSquareMesh(N, N, diagonal='crossed')

#Spaces
V = FunctionSpace(mesh, "DG", 1)           # Scalar u
W = VectorFunctionSpace(mesh, "DG", 1)     # Flux sigma

#Mixed function space
Z = V * W
PETSc.Sys.Print('Nb dof: %i' % Z.dim())

#Trial and test functions
u,sigma = TrialFunctions(Z)
v,tau = TestFunctions(Z)

#Geometric functions
n = FacetNormal(mesh)
h = CellDiameter(mesh)

#RHS
f = Function(V)
x, y = SpatialCoordinate(mesh)
f.interpolate(2 * pi**2 * sin(pi * x) * sin(pi * y))  # exact solution: sin(pi x) sin(pi y)
L = f * v * dx

#LHS
a_interior = dot(sigma, tau)*dx + dot(grad(u), tau)*dx - dot(sigma, grad(v))*dx

#Interior penalty terms
alpha = Constant(10)
a_interior += alpha / avg(h) * jump(u)*jump(v)*dS - inner(dot(avg(tau), n('+')), jump(u))*dS + inner(dot(avg(sigma), n('+')), jump(v))*dS

#Boundary terms (homogeneous Dirichlet)
a_boundary = alpha / h * u * v * ds - dot(tau, n) * u * ds + dot(sigma, n) * v * ds

#Total bilinear form and RHS
a = a_interior + a_boundary

#Slate
A = Tensor(a)
b = Tensor(L)

#Solving
z = Function(Z, name="z")
solve(assemble(A), z.vector(), assemble(b))
#solve(assemble(a), z.vector(), assemble(L))

#Store result
from firedrake.output import VTKFile
file = VTKFile('sol.pvd')
file.write(z.sub(0))

If you comment the solve and use the "regular" solve without creating the Tensor objects before you get the expected solution. But if you assemble the Tensor objects and then solve, the computed solution is very wrong. I don't understand why.

Thank you for your help!