How to write user-defined JIT-compiled Expressions with changeable parameters?

Question

I have written a complicated expression as a subclass-expression with changeable parameters following examples found on the forum. Unfortunately, the approach is prohibitively slow. In the reference on expressions it is stated that writing a JIT-compiled expression on the form

code = '''class MyFunc : public Expression { 
};''' 
cell_data = CellFunction('uint', V.mesh())
f = Expression(code)
f.cell_data = cell_data

is much faster. However, in the documentation there are no examples of how to actually write a user-defined JIT-compiled Expressions with changeable parameters. Does anyone know how to write a JIT-compiled expression in the above form of the simple subclass defined expression given below ?

class MyExpression(Expression):
def __init__(self, a, b,element):
    self.a = a
    self.b = b
    self._ufl_element = element
def eval(self, value, x):
    dx = x[0] - self.a
    dy = x[1] - self.b
    value[0] = exp(-(dx*dx + dy*dy)/0.02)
    value[1] = exp(-(dx*dx + dy*dy)/0.02)
def value_shape(self):
    return (2,)

mesh = UnitCircle(20)
V = FunctionSpace(mesh, "BDM", 1)
f = MyExpression(a=0.5,b=0.5,element=V.ufl_element())

Answer 1

Hi, consider

from dolfin import *

class MyExpression(Expression):
    def __init__(self, a=0, b=0):
        self.a = a
        self.b = b

    def eval(self, value, x):
        dx = x[0] - self.a
        dy = x[1] - self.b
        value[0] = exp(-(dx*dx + dy*dy)/0.02)
        value[1] = exp(-(dx*dx + dy*dy)/0.02)

    def value_shape(self):
        return (2,)

cpp_code = '''
class MyCppExpression : public Expression
{
public:
  MyCppExpression() : Expression(2), a(0), b(0) { }

  void eval(Array<double>& values, const Array<double>& x) const
  {
    const double dx = x[0] - a;
    const double dy = x[1] - b;
    values[0] = exp(-(dx*dx + dy*dy)/0.02);
    values[1] = exp(-(dx*dx + dy*dy)/0.02);
  }
public:
    double a, b;
};'''

mesh = UnitSquareMesh(20, 20)

f = MyExpression(a=0.5, b=0.5)
f.a = 0.5
f.b = 0.25

f_cpp = Expression(cpp_code)
f_cpp.a = 0.5
f_cpp.b = 0.25

# Compare the two
V = VectorFunctionSpace(mesh, 'CG', 1)
F = interpolate(f, V).vector()
F_cpp = interpolate(f_cpp, V).vector()
F.axpy(-1, F_cpp)
print '|F-F_cpp|', F.norm('l2')

You might want to have a look into the source code to see that the python Expression class does not simply mirror the cpp Expression class.

Comments

It works, thanks!