Merge lp:~fluidity-core/fluidity/firedrake-mixed-space-expressions into lp:fluidity/firedrake

[Lawrence Mitchell]

Adds minimal infrastructure for mixed function spaces, along with support for expression evaluation.

LHS form assembly does not work correctly, because of a mismatch in what pyop2 wants and ffc generates. I believe that is on someone's roadmap.

So this is more an early review request to see if this the correct approach.

As a bonus, expression evaluation is fixed for vector function spaces. And errors are correctly raised for augmented assignment on mismatched function spaces.

[David Ham (david-ham)]

Please add docstrings, especially in core_types.pyx

Please add tests for such functionality as exists. For example firedrake_expressions could be expanded or duplicated to have some mixed functionspace cases in it.

Does any of the PyOP2 interaction work? For example if RHS assembly works then a simple integral test could be included.

review: Needs Fixing

[Lawrence Mitchell (wence)]

I believe RHS assembly doesn't work at all yet, but Florian suggests that may be here soon.

I'll fix up the rest.

[lp:~fluidity-core/fluidity/firedrake-mixed-space-expressions]

4297. By Lawrence Mitchell on 2013-09-02

Add docstrings to MixedFunctionSpace

4298. By Lawrence Mitchell on 2013-09-02

Add some tests of mixed space expressions

[Lawrence Mitchell (wence)]

docstrings and some tests added.

Note that for reasons I don't profess to understand the following does not work:

U = FunctionSpace(mesh, 'DG', 0)
W = U*U

w = Function(W)
v = Function(W)
w.assign(2*v)
 => AttributeError: 'ComponentTensor' object has no attribute '_operands'

Equally:

w.assign(v**2)
 => generated code is bad
w.assign(v**v)
 => UFLException: Cannot take the power of a non-scalar expression.

So this still needs some work.

[Lawrence Mitchell (wence)]

So, foo**bar doesn't work at all for mixed spaces. This is the same as dolfin. I think it would be possible to make it work, but I'm not sure if it actually makes sense. The error is deep in UFL.

The component tensor error I think needs one to unpick more UFL classes.

The other option is to lie about the ufl_element in the mixed function spaces. If I say that

fs._ufl_element = self._function_spaces[0]._ufl_element then everything works. But that's tremendously icky.

[lp:~fluidity-core/fluidity/firedrake-mixed-space-expressions]

4299. By Lawrence Mitchell on 2013-09-03

Implement more expressions on mixed spaces

Instead of carrying around a slot on every class we need to index, just
keep a global index object that we look at if trying to generate an
expression over a mixed space.

This simplifies the code, and makes implementing correct expression
assembly for Indexed and ComponentTensor expressions much easier. So
do that, and add appropriate tests.

4300. By Lawrence Mitchell on 2013-09-03

Add more docstrings

[Lawrence Mitchell (wence)]

So I fixed more of the things, so that now:

a.assign(3*b + c)

on mixed spaces should work.

[Florian Rathgeber (florian-rathgeber)]

I've merged this into the firedrake-mixed branch for now since it depends on the mixed data types.

[Lawrence Mitchell (wence)]

This will come in with the mixed space stuff since it now lives there, so this merge proposal is obsoleted.

Unmerged revisions

4300. By Lawrence Mitchell on 2013-09-03

Add more docstrings

4299. By Lawrence Mitchell on 2013-09-03

Implement more expressions on mixed spaces

Instead of carrying around a slot on every class we need to index, just
keep a global index object that we look at if trying to generate an
expression over a mixed space.

This simplifies the code, and makes implementing correct expression
assembly for Indexed and ComponentTensor expressions much easier. So
do that, and add appropriate tests.

4298. By Lawrence Mitchell on 2013-09-02

Add some tests of mixed space expressions

4297. By Lawrence Mitchell on 2013-09-02

Add docstrings to MixedFunctionSpace

4296. By Lawrence Mitchell on 2013-09-02

Support expressions on mixed spaces

As a bonus, this fixes expressions on VectorFunctionSpaces and
correctly raises an error for augmented assignment on mismatching
FunctionSpaces.

4295. By Lawrence Mitchell on 2013-09-02

Add rudimentary support for mixed function spaces

=== modified file 'python/firedrake/assemble_expressions.py'
--- python/firedrake/assemble_expressions.py	2013-08-27 15:14:32 +0000
+++ python/firedrake/assemble_expressions.py	2013-09-03 15:11:37 +0000
@@ -5,6 +5,11 @@
 import core_types
 import cgen
 
+# Global state telling us which function space in a mixed space we're indexing into
+# while generating the kernel.
+global _fs_idx
+_fs_idx = None
+
 class DummyFunction(ufl.Coefficient):
     """A dummy object to take the place of a Function in the
 expression. This has the sole role of producing the right strings when
@@ -21,16 +26,29 @@
         self.intent = intent
 
     def __str__(self):
-        if (self.function.rank()==0):
+        # To generate the correct kernel body for a mixed function
+        # space we need to know which of the sub spaces we're
+        # currently generating for.  The only way to do this is to
+        # peek at some hidden state.
+        if _fs_idx is not None:
+            fs = self.function.function_space()[_fs_idx]
+        else:
+            fs = self.function.function_space()
+        if isinstance(fs, core_types.VectorFunctionSpace):
+            return "fn_%d[dim]" % self.argnum
+        else:
             return "fn_%d[0]" % self.argnum
-        else:
-            return "fn_%d[0][dim]" % self.argnum
 
     @property
     def arg(self):
-        argtype = self.function.dat.ctype+"*"
-        if (self.function.rank()>0):
-           arg += "*"
+        if _fs_idx is not None:
+            dat = self.function.dat[_fs_idx]
+            fs = self.function.function_space()[_fs_idx]
+        else:
+            dat = self.function.dat
+            fs = self.function.function_space()
+        argtype = dat.ctype+"*"
+
         name = " fn_" + `self.argnum`
 
         return cgen.Value(argtype, name)
@@ -69,8 +87,13 @@
             new_lhs.intent = op2.RW
 
         except KeyError:
+            if transformer._function_space is None:
+                transformer._function_space = lhs._function_space
+            elif transformer._function_space != lhs._function_space:
+                raise ValueError("Expression has incompatible function spaces")
             transformer._args[lhs] = DummyFunction(lhs, len(transformer._args),
                                                    intent=op2.WRITE)
+            transformer._function_space = lhs._function_space
             new_lhs = transformer._args[lhs]
 
         return [new_lhs, self._operands[1]]
@@ -89,6 +112,10 @@
         try:
             new_lhs = transformer._args[lhs]
         except KeyError:
+            if transformer._function_space is None:
+                transformer._function_space = lhs._function_space
+            elif transformer._function_space != lhs._function_space:
+                raise ValueError("Expression has incompatible function spaces")
             transformer._args[lhs] = DummyFunction(lhs, len(transformer._args))
             new_lhs = transformer._args[lhs]
 
@@ -139,6 +166,20 @@
     def __str__(self):
         return "log(%s)" % str(self._argument)
 
+
+class Indexed(ufl.indexed.Indexed):
+    """Subclass of :class:`ufl.indexed.Indexed` which prints the unindexed expression."""
+
+    def __str__(self):
+        return "%s" % self._expression
+
+class ComponentTensor(ufl.tensors.ComponentTensor):
+    """Subclass of :class:`ufl.tensors.ComponentTensor` which prints the expression
+without any indexed set notation."""
+
+    def __str__(self):
+        return "%s" % self._expression
+
 class ExpressionWalker(ReuseTransformer):
     def __init__(self):
         ReuseTransformer.__init__(self)
@@ -148,14 +189,11 @@
         self._result = None
 
     def coefficient(self, o):
-
         if isinstance(o, core_types.Function):
             if self._function_space is None:
                 self._function_space = o._function_space
-            else:
-                if self._function_space != o._function_space:
-                    raise ValueError("Expression has incompatible function spaces")
-
+            elif self._function_space != o._function_space:
+                raise ValueError("Expression has incompatible function spaces")
             try:
                 arg = self._args[o]
                 if arg.intent == op2.WRITE:
@@ -179,6 +217,12 @@
     condition = ReuseTransformer.reuse_if_possible
     math_function = ReuseTransformer.reuse_if_possible
 
+    def component_tensor(self, o, *operands):
+        return ComponentTensor(*operands)
+
+    def indexed(self, o, *operands):
+        return Indexed(*operands)
+
     def power(self, o, *operands):
         # Need to convert notation to c for exponents.
         return Power(*operands)
@@ -198,7 +242,6 @@
             operands = [operands[0], self.visit(operands[1])]
 
         else:
-
             # For all other operators, just visit the children.
             operands = map(self.visit, o._operands)
 
@@ -208,41 +251,60 @@
     """Produce a :class:`PyOP2.Kernel` from the processed UFL expression
 expr and the corresponding args.
     """
-
-    fs = args[0].function.function_space()
-
-    body=cgen.Block()
-
-    if isinstance(fs, core_types.VectorFunctionSpace):
-        body.append(
-            [cgen.Value("int", "dim"),
-             cgen.For("dim = 0",
-                      "dim < %s"%fs.dof_dset.cdim,
-                      "++i",
-                      cgen.Line(str(expr)+";")
+    def gen(fs, idx=None):
+        body=cgen.Block()
+        global _fs_idx
+        _fs_idx = idx
+        if isinstance(fs, core_types.VectorFunctionSpace):
+            body.extend([cgen.Value("int", "dim"),
+                 cgen.For("dim = 0",
+                          "dim < %s"%fs.dof_dset.cdim,
+                          "++dim",
+                          cgen.Line(str(expr)+";")
                       )
              ]
             )
+        else:
+            body.append(cgen.Line(str(expr)+";"))
+
+        fdecl = cgen.FunctionDeclaration(
+            cgen.Value("void", "expression"),
+            [arg.arg for arg in args])
+
+        kernel = str(cgen.FunctionBody(fdecl, body))
+        _fs_idx = None
+        return op2.Kernel(kernel, "expression")
+
+    space = args[0].function.function_space()
+
+    if isinstance(space, core_types.MixedFunctionSpace):
+        for i, subspace in enumerate(space.split()):
+            yield gen(subspace, i)
     else:
-        body.append(cgen.Line(str(expr)+";"))
-
-    fdecl = cgen.FunctionDeclaration(
-        cgen.Value("void", "expression"),
-        [arg.arg for arg in args])
-
-    return op2.Kernel(str(cgen.FunctionBody(fdecl, body)), "expression")
+        yield gen(space)
+            
 
 def evaluate_preprocessed_expression(expr, args):
 
-    kernel = expression_kernel(expr, args)
+    kernels = expression_kernel(expr, args)
 
     fs = args[0].function._function_space
 
-    parloop_args = [kernel, fs.node_set] +\
-        [arg.function.dat(arg.intent)
-         for arg in args ]
-
-    op2.par_loop(*parloop_args)
+    if isinstance(fs, core_types.MixedFunctionSpace):
+        spaces = fs.split()
+    else:
+        spaces = [fs]
+
+    for i, (kernel, space) in enumerate(zip(kernels, spaces)):
+        parloop_args = [kernel, space.node_set]
+        for arg in args:
+            if isinstance(fs, core_types.MixedFunctionSpace):
+                dat = arg.function.dat[i]
+            else:
+                dat = arg.function.dat
+            parloop_args.append(dat(arg.intent))
+
+        op2.par_loop(*parloop_args)
 
 def evaluate_expression(expr):
     """Evaluates UFL expressions on Functions.

=== modified file 'python/firedrake/core_types.pyx'
--- python/firedrake/core_types.pyx	2013-08-30 10:55:57 +0000
+++ python/firedrake/core_types.pyx	2013-09-03 15:11:37 +0000
@@ -443,12 +443,11 @@
         this :class:`FunctionSpace`."""
         return op2.DataSet(self.node_set, self.dim)
 
-    def make_dat(self):
+    def make_dat(self, val=None, uid=None):
         """Return a newly allocated :class:`pyop2.Dat` defined on the
         :attr:`dof.dset` of this :class:`Function`."""
-
-        return op2.Dat(self.dof_dset, np.zeros(self.dof_count),
-                       valuetype)
+        data = val if val is not None else np.zeros(self.dof_count)
+        return op2.Dat(self.dof_dset, data, valuetype, uid=uid)
 
 
     @utils.cached_property
@@ -475,6 +474,12 @@
         """The :class:`Mesh` used to construct this :class:`FunctionSpace`."""
         return self._mesh
 
+    def __mul__(self, other):
+        """Create a :class:`MixedFunctionSpace` composed of this
+        :class:`FunctionSpace` and other"""
+        return MixedFunctionSpace(self, other)
+
+
 class VectorFunctionSpace(FunctionSpace):
     def __init__(self, mesh, family, degree, dim = None, name = None):
         super(VectorFunctionSpace, self).__init__(mesh, family, degree, name)
@@ -486,6 +491,125 @@
         else:
             self._dim = dim
 
+
+class MixedFunctionSpace(FunctionSpace):
+    """A representation of a mixed function space."""
+    def __init__(self, *args, **kwargs):
+        """Build a :class:`MixedFunctionSpace` from some number of
+    :class:`FunctionSpaces`.
+
+    :arg args: base :class:`FunctionSpace`\s to be composed.
+    This can also be a list or tuple of :class:`FunctionSpace`\s (for
+    Dolfin compatibilty).
+    :arg name: set the name of the :class:`MixedFunctionSpace`."""
+        tmp = []
+        def flatten(a):
+            for arg in a:
+                if isinstance(arg, FunctionSpace):
+                    tmp.append(arg)
+                elif isinstance(arg, MixedFunctionSpace):
+                    flatten(arg)
+                else:
+                    raise RuntimeError("Don't know how to build a MixedFunctionSpace using a %s" % type(arg))
+        if len(args) == 1:
+            args = args[0]
+        flatten(args)
+        self._function_spaces = tuple(tmp)
+        mesh = self._function_spaces[0].mesh()
+        for fs in self._function_spaces:
+            assert fs.mesh() == mesh, \
+                "Mismatching meshes in mixed function space definition"
+        self._mesh = mesh
+        self._ufl_element = ufl.MixedElement(*[fs.ufl_element() \
+                                               for fs in self._function_spaces])
+        self.name = kwargs.get('name')
+
+    def split(self):
+        """The :class:`FunctionSpace`\s this
+    :class:`MixedFunctionSpace` is composed of"""
+        return self._function_spaces
+
+    def sub(self, i):
+        """Return the `i`th :class:`FunctionSpace` in this :class:`MixedFunctionSpace`."""
+        return self[i]
+
+    def num_sub_spaces(self):
+        """Return the number of :class:`FunctionSpace`\s this
+        :class:`MixedFunctionSpace` is composed of."""
+        return len(self)
+
+    def __len__(self):
+        """Return the number of :class:`FunctionSpace`\s this
+    :class:`MixedFunctionSpace` is composed of."""
+        return len(self._function_spaces)
+
+    def __getitem__(self, i):
+        """Return the `i`th :class:`FunctionSpace` in this
+    :class:`MixedFunctionSpace`."""
+        return self._function_spaces[i]
+
+    @property
+    def dim(self):
+        """Return a tuple of the :attr:`FunctionSpace.dim`\s of the
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return tuple(fs.dim for fs in self._function_spaces)
+
+    @property
+    def node_count(self):
+        """Return a tuple of the :attr:`FunctionSpace.node_count`\s of the
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return tuple(fs.node_count for fs in self._function_spaces)
+
+    @property
+    def dof_count(self):
+        """Return a tuple of the :attr:`FunctionSpace.dof_count`\s of the
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return tuple(fs.dof_count for fs in self._function_spaces)
+
+    @utils.cached_property
+    def node_set(self):
+        """A :class:`pyop2.MixedSet` containing the nodes of this
+    :class:`MixedFunctionSpace`.  This is composed of the
+    :attr:`FunctionSpace.node_set`\s of the underlying
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return op2.MixedSet([fs.node_set for fs in self._function_spaces])
+
+    @utils.cached_property
+    def dof_dset(self):
+        """A :class:`pyop2.MixedDataSet` containing the degrees of freedom of
+    this :class:`MixedFunctionSpace`.  This is composed of the
+    :attr:`FunctionSpace.dof_dset`\s of the underlying
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return op2.MixedDataSet([fs.dof_dset for fs in self._function_spaces])
+
+    @utils.cached_property
+    def cell_node_map(self):
+        """A :class:`pyop2.MixedMap` from the :attr:`Mesh.cell_set` of the
+    underlying mesh to the :attr:`node_set` of this
+    :class:`MixedFunctionSpace`.  This is composed of the
+    :attr:`FunctionSpace.cell_node_map`\s of the underlying
+    :class:`FunctionSpace`\s this :class:`MixedFunctionSpace` is
+    composed of."""
+        return op2.MixedMap([fs.cell_node_map for fs in self._function_spaces])
+    
+    def make_dat(self, vals=None, uid=None):
+        """Build a :class:`pyop2.MixedDat` on this
+    :class:`MixedFunctionSpace`.  If `vals` is provided, it should
+    be an iterable of numpy data arrays of the same length as the
+    :class:`MixedFunctionSpace`."""
+        if vals is not None:
+            assert len(vals) == len(self)
+        else:
+            vals = [None] * len(self)
+        return op2.MixedDat([fs.make_dat(val, uid) \
+                             for fs, val in zip(self._function_spaces, vals)])
+
+
 class Function(ufl.Coefficient):
     """A :class:`Function` represents a discretised field over the
 domain defined by the underlying :class:`Mesh`. Functions are
@@ -507,7 +631,7 @@
 
         if isinstance(function_space, Function):
             self._function_space = function_space._function_space
-        elif isinstance(function_space, FunctionSpace):
+        elif isinstance(function_space, (MixedFunctionSpace, FunctionSpace)):
             self._function_space = function_space
         else:
             raise NotImplementedError("Can't make a Function defined on a "
@@ -518,13 +642,7 @@
         self.name = name
         self.uid = _new_uid()
 
-        if val is not None:
-            self.dat = op2.Dat(self.dof_dset, val, valuetype, self.name,
-                               uid=self.uid)
-        else:
-            self.dat = op2.Dat(self.dof_dset, 
-                               np.zeros(self._function_space.dof_count),
-                               valuetype, self.name, uid=self.uid)
+        self.dat = self._function_space.make_dat(val, self.uid)
 
         self._repr = None
 
@@ -656,4 +774,3 @@
             assemble_expressions.IDiv(self, expr))
 
         return self
-

=== modified file 'python/firedrake/ufl_expr.py'
--- python/firedrake/ufl_expr.py	2013-08-27 15:14:32 +0000
+++ python/firedrake/ufl_expr.py	2013-09-03 15:11:37 +0000
@@ -2,6 +2,7 @@
 import ufl.argument
 from ufl.assertions import ufl_assert
 from ufl.finiteelement import FiniteElementBase
+from ufl.split_functions import split
 
 class Argument(ufl.argument.Argument):
 
@@ -41,3 +42,9 @@
 
 def TrialFunction(function_space):
     return Argument(function_space.ufl_element(), function_space, -1)
+
+def TestFunctions(function_space):
+    return split(TestFunction(function_space))
+
+def TrialFunctions(function_space):
+    return split(TrialFunction(function_space))

=== modified file 'tests/firedrake_expressions/test_firedrake_expressions.py'
--- tests/firedrake_expressions/test_firedrake_expressions.py	2013-08-28 09:58:18 +0000
+++ tests/firedrake_expressions/test_firedrake_expressions.py	2013-09-03 15:11:37 +0000
@@ -5,8 +5,12 @@
 mesh = fd.Mesh("square-2d.node")
 
 fs = fd.FunctionSpace(mesh, "Lagrange", 1)
-    
+
+mixed = fs * fd.FunctionSpace(mesh, "DG", 1)
+
 f = fd.Function(fs, name="f")
+mf = fd.Function(mixed, name="mixed")
+
 one = fd.Function(fs, name="one")
 two = fd.Function(fs, name="two")
 minusthree =  fd.Function(fs, name="minusthree")
@@ -48,6 +52,35 @@
     return (str(f)+" /= "+str(expr)+", "+str(f), x,
             np.all(f.dat.data == x))
 
+def mixed_assign(f, expr, x):
+    f.assign(expr)
+    return (str(f)+" = "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
+
+def mixed_iadd(f, expr, x):
+    f += expr
+    return (str(f)+" += "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
+
+def mixed_isub(f, expr, x):
+    f -= expr
+    return (str(f)+" -= "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
+
+def mixed_imul(f, expr, x):
+    f *= expr
+    return (str(f)+" *= "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
+
+def mixed_idiv(f, expr, x):
+    f /= expr
+    return (str(f)+" /= "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
+
+def mixed_assign(f, expr, x):
+    f.assign(expr)
+    return (str(f)+" = "+str(expr)+", "+str(f), x,
+            np.all([np.all(data == x) for data in f.dat.data]))
 tests = [
     test(one + two, 3),
     test(ufl.ln(one), 0),
@@ -59,7 +92,16 @@
     iaddtest(f, 2, 7),
     isubtest(f, 2, 5),
     imultest(f, 2, 10),
-    idivtest(f, 2, 5)
+    idivtest(f, 2, 5),
+    mixed_assign(mf, 1, 1),
+    mixed_iadd(mf, 2, 3),
+    mixed_iadd(mf, mf, 6),
+    mixed_isub(mf, 10, -4),
+    mixed_idiv(mf, -4, 1),
+    mixed_imul(mf, 10, 10),
+    mixed_idiv(mf, mf, 1),
+    mixed_assign(mf, mf + mf, 2),
+    mixed_assign(mf, 2*mf, 4)
     ]