How about subclass-ing your k_2 as well? Much like how you are doing for K. Of course, there may be a smarter way to do this, but, for instance
from dolfin import *
class k2(UserExpression):
def __init__(self, **kwargs):
super().__init__(**kwargs)
def eval_cell(self, values, x, cell):
values[0] = x[0]**2
class K(UserExpression):
def __init__(self, materials, k_0, k_1, k_2, **kwargs):
super().__init__(**kwargs)
self.materials = materials
self.k_0 = k_0
self.k_1 = k_1
self.k_2 = k_2
def eval_cell(self, values, x, cell):
if self.materials.array()[cell.index] == 3:
self.k_2.eval_cell(values,x,cell)
elif self.materials.array()[cell.index] == 1:
values[0] = self.k_0
else:
values[0] = self.k_1
mesh = UnitSquareMesh(10, 10)
top = CompiledSubDomain("x[1] >= 0.7")
bottom = CompiledSubDomain("x[1] < 0.5")
middle = CompiledSubDomain("x[1] >=0.5 && x[1] < 0.7")
materials = cpp.mesh.MeshFunctionSizet(mesh, mesh.topology().dim())
top.mark(materials, 1)
middle.mark(materials, 2)
bottom.mark(materials, 3)
V = FunctionSpace(mesh, "DG", 2)
E = K(materials, 10., 20., k2(element=V.ufl_element()), degree=1, element=V.ufl_element())
u = Function(V)
u.interpolate(E)
should work too. Alternatively you can directly use functions in DG spaces to assign different material parameters. See this post for an example