Interpolation of some indexed expressions onto vector-valued tensor-product elements fails

[glumpshikt]

Hello. I've noticed that the interpolation function fails for some indexed expressions. In the code below, lines 11-14 work as intended but 15-16 throw an error (shown below). When line 4 mesh = ExtrudedMesh(mesh,5) is commented out, all expressions work as intended.

  1 from firedrake import *
  2
  3 mesh = UnitSquareMesh(5,5, quadrilateral=True)
  4 mesh = ExtrudedMesh(mesh, 5)
  5
  6 elem = FiniteElement("Lagrange", cell=mesh.ufl_cell(), degree=1)
  7 V = VectorFunctionSpace(mesh, elem, dim=2)
  8 u = Function(V)
  9
 10 x = SpatialCoordinate(mesh)
 11 u.interpolate(as_vector([dot(x,x), 0]))
 12 u.interpolate(as_vector([1,0])*(x[0]**2+x[1]**2))
 13 u.interpolate(Constant([1,0])*(x[0]**2+x[1]**2))
 14 u.interpolate(as_vector([1,0])*exp(dot(x,x)))
 15 u.interpolate(as_vector([1,0])*dot(x,x))
 16 u.interpolate(Constant([1,0])*dot(x,x))                                                   

Traceback (most recent call last):
  File "sumfactorisationerror.py", line 15, in <module>
    u.interpolate(as_vector([1,0])*dot(x,x))
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/function.py", line 365, in interpolate
    return interpolation.interpolate(expression, self, subset=subset, ad_block_tag=ad_block_tag)
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/interpolation.py", line 61, in interpolate
    return Interpolator(expr, V, subset=subset, access=access).interpolate(ad_block_tag=ad_block_tag)
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/interpolation.py", line 87, in __init__
    self.callable, arguments = make_interpolator(expr, V, subset, access)
  File "PETSc/Log.pyx", line 115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "PETSc/Log.pyx", line 116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/interpolation.py", line 218, in make_interpolator
    loops.extend(_interpolator(V, tensor, expr, subset, arguments, access))
  File "<decorator-gen-28>", line 2, in _interpolator
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/utils.py", line 74, in wrapper
    return f(*args, **kwargs)
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/interpolation.py", line 291, in _interpolator
    kernel = compile_expression(cell_set.comm, expr, to_element, V.ufl_element(),
  File "/home/glumpshikt/firedrake/src/PyOP2/pyop2/caching.py", line 300, in wrapper
    v = func(*args, **kwargs)
  File "/home/glumpshikt/firedrake/src/firedrake/firedrake/interpolation.py", line 406, in compile_expression
    return compile_expression_dual_evaluation(*args, **kwargs)
  File "/home/glumpshikt/firedrake/src/tsfc/tsfc/driver.py", line 253, in compile_expression_dual_evaluation
    evaluation, basis_indices = to_element.dual_evaluation(fn)
  File "/home/glumpshikt/firedrake/src/FInAT/finat/tensorfiniteelement.py", line 200, in dual_evaluation
    evaluation = gem.optimise.contraction(evaluation)
  File "/home/glumpshikt/firedrake/src/tsfc/gem/optimise.py", line 584, in contraction
    return rebuild(expression)
  File "/home/glumpshikt/firedrake/src/tsfc/gem/optimise.py", line 558, in rebuild
    return sum_factorise(sum_indices, factors)
  File "/home/glumpshikt/firedrake/src/tsfc/gem/optimise.py", line 358, in sum_factorise
    raise NotImplementedError("Too many indices for sum factorisation!")
NotImplementedError: Too many indices for sum factorisation!

[glumpshikt]

The issue can be worked around if the line evaluation=gem.optimise.contraction(evaluation) in FInAT/finat/tensorfiniteelement.py is commented out

[wence-]

Thanks. @ReubenHill this could be treated as part of the continuing work on dual evaluation? When we implemented it for tensor product elements, we used a "big hammer" to optimise the contractions, but really needed to do a better job (as indicated here).

[ReubenHill]

What do you think the right thing to do would be? I wonder how many free indices it's finding...

[wence-]

More than six iirc (since 6! = 720 which was considered to be "big enough for anyone").

It's been a while since I thought about this, but I think the right way to try and optimise the contractions is to first evaluate F_q := f(x_q) and then optimise the change of basis contraction Q F_q. This, I think, gives at least an asymptotically correct algorithmic complexity. But right now we just hand over Q f(x_q) which has many indices.

[ReubenHill]

Shouldn't GEM be clever enough to do the pre-evaluation like that itself? It also won't work when x_q is to be specified at runtime no?

[wence-]

Shouldn't GEM be clever enough to do the pre-evaluation like that itself? It also won't work when x_q is to be specified at runtime no?

GEM is declarative, so doesn't do any optimisation (except some minimal amount of constant folding).

What I mean is that the interpolation is two steps:

1. evaluate the to-be-interpolated expression at the output space node points
2. Hit the evaluated thing with the big Q matrix.

Each of those steps needs to separately have contractions optimised I think.