UFL implicit summation with power

Hello! I have an issue defining a variational formulation in UFL. I know the title is not very descriptive, sorry about that!

I want to define a certain quantity

Q = \sum_i \sum_j b_{ij} E_{ij}^2

where b is a matrix and E = \frac{\nabla u + \nabla u^T}{2} is the symmetric gradient of a vector function u.

In my ufl file I defined

u = Coefficient(element)
b = as_matrix([[ , , ],
               [ , , ],
               [ , , ]])
E_ = sym(grad(u))
Q = b[i, j] * E_[i, j]**2

If I try to compile it, I get an error saying

ufl.log.UFLException: Cannot take the power of a non-scalar expression <Indexed id=140030292038096>.

It’s not really clear to me what the issue is: E_ is a matrix, which means that E_[i, j] is one of its components, so it’s a scalar and I’m not taking the power of a non-scalar.
If I write down the whole expression expanded as b[0, 0] * E_[0, 0]**2 + b[0, 1] * E_[0, 1]**2 ..., then it works fine, which makes me think that maybe there’s something going on with the implicit summation.

Thanks for any help!
Massimiliano

Hi Massimiliano,

This error occurs because E_[i,j] is not a “true” UFL scalar. Because it has free indices, it represents a tensor-like quantity. Specifically, it fails the test is_true_ufl_scalar in ufl.checks because its attribute ufl_free_indices is a non-empty tuple. Thus, the only way to achieve what you want in this case is to use the expanded expression using fixed (int) indices.

In general this could become quite cumbersome, so I’ve written a function element_wise(expr, func) that takes an expression expr (with arbitrary shape) and a function func (having one argument) and applies that function element-wise to all elements of expr. See below for this function and examples of its use:

from fenics import (UnitSquareMesh, FunctionSpace, VectorFunctionSpace, 
                    TensorFunctionSpace, Function, as_tensor, sym, grad)
from ufl import (i, j, exp, sqrt)

mesh = UnitSquareMesh(2,2)

U = FunctionSpace(mesh, 'Lagrange', 2)
V = VectorFunctionSpace(mesh, 'Lagrange', 2)
W = TensorFunctionSpace(mesh, 'Lagrange', 2, shape=(2,4,3))

u = Function(U)
v = Function(V)
w = Function(W)

def element_wise(expr, func):
    if not expr.ufl_shape:
        return as_tensor(func(expr))
    else:
        def build_matrix(expr, func, idx=tuple(), axis=0):
            if axis == len(expr.ufl_shape) - 1:
                return [func(expr[idx+(i,)]) for i in range(expr.ufl_shape[axis])]
            else:
                return [build_matrix(expr, func, idx+(i,), axis+1) for i in range(expr.ufl_shape[axis])]
        
        return as_tensor(build_matrix(expr, func))

# Example 1: `exp` of a scalar Function
u_exp = element_wise(u, exp)
print(u_exp.ufl_shape)

# Example 2: element-wise `sqrt` of a vector Function
v_sqrt = element_wise(v, sqrt)
print(v_sqrt.ufl_shape)

# Example 3: tensor Function squared element-wise
w_squared = element_wise(w, lambda x: x**2)
print(w_squared.ufl_shape)

# Example 4: element-wise square of symmetric gradient of a vector Function
b = as_tensor([[1,2,3],
               [4,5,6],
               [7,8,9]])
E = sym(grad(v))
E_squared = element_wise(E, lambda x: x**2)
print(E_squared.ufl_shape)

Q = b[i, j] * element_wise(E_squared, lambda x: x**2)[i,j]
print(Q.ufl_shape)

producing

()
(2,)
(2, 4, 3)
(2, 2)
()

Note that you should apply free indexing after using element_wise, i.e.

Q = b[i, j] * element_wise(E_squared, lambda x: x**2)[i,j]

not

Q = b[i, j] * element_wise(E_squared[i,j], lambda x: x**2)
2 Likes

Hello conpierce, thanks for the exhaustive reply and the explanation. I understand this is a limit of UFL then, not a mistake on my part.
Thanks for the snippet you provided, but it looks to me a bit like killing a fly with a sledgehammer :smiley:
I can replace my faulty line with Q = sum(b[i, j] * E_[i, j]**2 for i in range(3) for j in range(3)) and Python will do its magic just nicely.

2 Likes