Very large steps towards piotr stress formulation

python/dune/perftool/loopy/transformer.py

@@ -31,11 +31,11 @@ from pymbolic.primitives import Subscript, Variable
 
 
 class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapper):
-    def __init__(self, measure, subdomain_id, dimension_index_aliases):
+    def __init__(self, measure, subdomain_id, dimension_indices):
         # Some variables describing the integral measure of this integral
         self.measure = measure
         self.subdomain_id = subdomain_id
-        self.dimension_index_aliases = dimension_index_aliases
+        self.dimension_indices = dimension_indices
 
         # Call base class constructors
         super(UFL2LoopyVisitor, self).__init__()
@@ -96,7 +96,8 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
             if len(subel.value_shape()) != 0:
                 from dune.perftool.pdelab.geometry import dimension_iname
                 from dune.perftool.pdelab.basis import lfs_child
-                lfs = lfs_child(lfs, dimension_iname(context='arg'))
+                idims = tuple(dimension_iname(context='arg', count=i) for i in range(len(subel.value_shape())))
+                lfs = lfs_child(lfs, idims, shape=subel.value_shape())
                 residual_shape[icount] = name_lfs_bound(name_leaf_lfs(subel.sub_elements()[0], ma.restriction))
                 subel = subel.sub_elements()[0]
             else:
@@ -181,8 +182,8 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
         # Now treat the case of this being a vector finite element
         if element.num_sub_elements() > 0:
             # I cannot handle general mixed elements here...
-            from ufl import VectorElement
-            assert isinstance(element, VectorElement)
+            from ufl import VectorElement, TensorElement
+            assert isinstance(element, (VectorElement, TensorElement))
 
             # Determine whether this is a non-scalar subargument. This information is later
             # used by the Indexed-Mapper to make some indices loop indices!
@@ -191,7 +192,7 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
             # If this is a vector element, we need add an additional accumulation loop iname
             from dune.perftool.pdelab.geometry import dimension_iname
             for i in range(self.argshape):
-                self.inames.append(dimension_iname(context='arg'))
+                self.inames.append(dimension_iname(context='arg', count=i))
 
             # For the purpose of basis evaluation, we need to take the leaf element
             leaf_element = element.sub_elements()[0]
@@ -286,7 +287,7 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
 
             # Recurse to get the summation expression
             term = self.call(o.ufl_operands[0])
-            self.redinames = tuple(i for i in self.redinames if i not in self.dimension_index_aliases)
+            self.redinames = tuple(i for i in self.redinames if i not in self.dimension_indices)
             if len(self.redinames) > 0:
                 ret = Reduction("sum", tuple(name_index(ind) for ind in self.redinames), term)
                 name = get_temporary_name()
@@ -314,10 +315,10 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker, UFL2PymbolicMapper, GeometryMapp
         else:
             from pymbolic.primitives import Variable
             from dune.perftool.pdelab import name_index
-            if index in self.dimension_index_aliases:
+            if index in self.dimension_indices:
                 from dune.perftool.pdelab.geometry import dimension_iname
-                self.inames.append(dimension_iname(context='arg'))
-                return Variable(dimension_iname(context='arg'))
+                self.inames.append(self.dimension_indices[index])
+                return Variable(self.dimension_indices[index])
             else:
                 return Variable(name_index(index))
 

python/dune/perftool/pdelab/basis.py

@@ -75,11 +75,26 @@ def define_lfs(name, father, child):
     return "auto {} = child({}, _{});".format(name, father, child)
 
 
-def lfs_child(lfs, child):
-    from pymbolic.primitives import Call
+def lfs_child(lfs, children, shape=None):
+    from pymbolic.primitives import Call, Product, Sum
+    # Old pre-TensorElement implementation kept for comaptibility
+    if shape is None:
+        indices = (Variable(children[0]),)
+    else:
+        children = tuple(Variable(c) for c in children)
+        offset = Product(tuple())
+        indices = (Sum(
+                      tuple(
+                            Product(
+                                    (Product(tuple(shape[j] for j in range(i))
+                                             ),
+                                     child)
+                                    )
+                            for i, child in enumerate(children))
+                      ),)
+
     from dune.perftool.loopy.functions import LFSChild
-    return Call(LFSChild(lfs), (Variable(child),))
-#     return "{}.child({})".format(lfs, child)
+    return Call(LFSChild(lfs), indices)
 
 
 @generator_factory(cache_key_generator=lambda e, r, **kw: (e, r))
@@ -307,6 +322,14 @@ def name_basis_gradient(leaf_element, restriction):
     return name
 
 
+def shape_as_pymbolic(shape):
+    def _shape_as_pymbolic(s):
+        if isinstance(s, str):
+            return Variable(s)
+        else:
+            return s
+    return tuple(_shape_as_pymbolic(s) for s in shape)
+
 @cached
 def evaluate_trialfunction(element, name, restriction, component):
     from ufl.functionview import select_subelement
@@ -314,28 +337,28 @@ def evaluate_trialfunction(element, name, restriction, component):
 
     # Determine the rank of the trialfunction tensor
     rank = len(sub_element.value_shape())
-    assert rank in (0, 1)
 
-    shape = (name_dimension(),) * rank
-    shape_impl = ('fv', ) * rank
+    shape = sub_element.value_shape()
+    shape_impl = ('arr', ) * rank
     idims = tuple(dimension_iname(count=i) for i in range(rank))
     leaf_element = sub_element
-    from ufl import VectorElement
-    if isinstance(sub_element, VectorElement):
+    from ufl import VectorElement, TensorElement
+    if isinstance(sub_element, (VectorElement, TensorElement)):
         leaf_element = sub_element.sub_elements()[0]
 
     temporary_variable(name, shape=shape, shape_impl=shape_impl)
     lfs = name_lfs(element, restriction, component)
     index = lfs_iname(leaf_element, restriction, context='trial')
     basis = name_basis(leaf_element, restriction)
-    assignee = Variable(name)
 
-    if isinstance(sub_element, VectorElement):
-        lfs = lfs_child(lfs, idims[0])
-        assignee = Subscript(assignee, Variable(idims[0]))
+    if isinstance(sub_element, (VectorElement, TensorElement)):
+        # TODO we need a concept of symmetry here
+        lfs = lfs_child(lfs, idims, shape=shape_as_pymbolic(shape))
 
     from dune.perftool.pdelab.argument import pymbolic_coefficient
     coeff = pymbolic_coefficient(lfs, index, restriction)
+    assignee = Subscript(Variable(name), tuple(Variable(i) for i in idims))
+
     reduction_expr = Product((coeff, Subscript(Variable(basis), Variable(index))))
     instruction(expression=Reduction("sum", index, reduction_expr, allow_simultaneous=True),
                 assignee=assignee,
@@ -350,18 +373,15 @@ def evaluate_trialfunction_gradient(element, name, restriction, component):
     from ufl.functionview import select_subelement
     sub_element = select_subelement(element, component)
     rank = len(sub_element.value_shape()) + 1
-    assert rank in (1, 2)
 
     # We do then set some variables accordingly
-    shape = (name_dimension(),) * rank
-    if rank == 1:
-        shape_impl = ('fv',)
-    else:
-        shape_impl = ('fm',)
+    shape = sub_element.value_shape() + (element.cell().geometric_dimension(),)
+    shape_impl = ('arr',) * rank
+
     idims = tuple(dimension_iname(count=i) for i in range(rank))
     leaf_element = sub_element
-    from ufl import VectorElement
-    if isinstance(sub_element, VectorElement):
+    from ufl import VectorElement, TensorElement
+    if isinstance(sub_element, (VectorElement, TensorElement)):
         leaf_element = sub_element.sub_elements()[0]
 
     # and proceed to call the necessary generator functions
@@ -370,8 +390,9 @@ def evaluate_trialfunction_gradient(element, name, restriction, component):
     index = lfs_iname(leaf_element, restriction, context='trialgrad')
     basis = name_basis_gradient(leaf_element, restriction)
 
-    if isinstance(sub_element, VectorElement):
-        lfs = lfs_child(lfs, idims[0])
+    if isinstance(sub_element, (VectorElement, TensorElement)):
+        # TODO we need a concept of symmetry here
+        lfs = lfs_child(lfs, idims[:-1], shape=shape_as_pymbolic(shape[:-1]))
 
     from dune.perftool.pdelab.argument import pymbolic_coefficient
     coeff = pymbolic_coefficient(lfs, index, restriction)

python/dune/perftool/pdelab/driver.py

@@ -73,11 +73,17 @@ def isQk(fem):
 def isDG(fem):
     return fem._short_name is 'DG'
 
+def flatten_tensor(expr):
+    from ufl import TensorElement
+    assert isinstance(expr, TensorElement)
+
 
 def FEM_name_mangling(fem):
-    from ufl import MixedElement, VectorElement, FiniteElement
+    from ufl import MixedElement, VectorElement, FiniteElement, TensorElement
     if isinstance(fem, VectorElement):
         return FEM_name_mangling(fem.sub_elements()[0]) + "_" + str(fem.num_sub_elements())
+    if isinstance(fem, TensorElement):
+        return FEM_name_mangling(fem.sub_elements()[0]) + "_" + "_".join(str(i) for i in fem.value_shape())
     if isinstance(fem, MixedElement):
         name = ""
         for elem in fem.sub_elements():
@@ -413,8 +419,8 @@ def define_compositegfs_constraints(name, *children):
 
 @symbol
 def name_bctype_function(expr):
-    from ufl import MixedElement, VectorElement
-    if isinstance(expr, VectorElement):
+    from ufl import MixedElement, VectorElement, TensorElement
+    if isinstance(expr, (VectorElement, TensorElement)):
         element, (dirichlet, _) = get_constraints_metadata(expr)
         # get the correct name from the corresponding UFL file
         from dune.perftool.generation import get_global_context_value
@@ -461,16 +467,17 @@ def name_assembled_constraints(expr):
 @preamble
 def typedef_gfs(element, dirichlet, name):
     vb = type_vectorbackend()
-    from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement
+    from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement, TensorElement
     if isinstance(element, FiniteElement):
         gv = type_leafview()
         fem = type_fem(element)
         cass = type_constraintsassembler(dirichlet)
         return "typedef Dune::PDELab::GridFunctionSpace<{}, {}, {}, {}> {};".format(gv, fem, cass, vb, name)
-    if isinstance(element, VectorElement):
+    if isinstance(element, (VectorElement, TensorElement)):
         include_file("dune/pdelab/gridfunctionspace/vectorgridfunctionspace.hh", filetag="driver")
         gv = type_leafview()
         fem = type_fem(element.sub_elements()[0])
+        # TODO implement symmetry here
         dim = element.num_sub_elements()
         cass = type_constraintsassembler(dirichlet)
         return "typedef Dune::PDELab::VectorGridFunctionSpace<{}, {}, {}, {}, {}, {}> {};".format(gv, fem, dim, vb, vb, cass, name)
@@ -495,12 +502,12 @@ def type_gfs(expr):
 @preamble
 def define_gfs(element, dirichlet, name):
     gfstype = type_gfs(element)
-    from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement
+    from ufl import FiniteElement, MixedElement, VectorElement, EnrichedElement, RestrictedElement, TensorElement
     if isinstance(element, FiniteElement):
         gv = name_leafview()
         fem = name_fem(element)
         return "{} {}({}, {});".format(gfstype, name, gv, fem)
-    if isinstance(element, VectorElement):
+    if isinstance(element, (VectorElement, TensorElement)):
         gv = name_leafview()
         fem = name_fem(element.sub_elements()[0])
         return "{} {}({}, {});".format(gfstype, name, gv, fem)
@@ -705,8 +712,8 @@ def define_compositegfs_parameterfunction(name, *children):
 
 @symbol
 def name_boundary_function(expr):
-    from ufl import MixedElement, VectorElement
-    if isinstance(expr, VectorElement):
+    from ufl import MixedElement, VectorElement, TensorElement
+    if isinstance(expr, (VectorElement, TensorElement)):
         element, (_, boundary) = get_constraints_metadata(expr)
         from dune.perftool.generation import get_global_context_value
         name = get_global_context_value("namedata").get(id(boundary), "zero")
@@ -881,8 +888,8 @@ def print_matrix():
 
 @preamble
 def define_gfs_name(element, prefix=()):
-    from ufl import MixedElement, VectorElement
-    if isinstance(element, VectorElement):
+    from ufl import MixedElement, VectorElement, TensorElement
+    if isinstance(element, (VectorElement, TensorElement)):
         gfs = name_gfs(element)
         return "{}.name(\"data_{}\");".format(gfs, "_".join(str(p) for p in prefix))
     if isinstance(element, MixedElement):

python/dune/perftool/pdelab/localoperator.py

@@ -161,8 +161,8 @@ def generate_kernel(integrals):
         subdomain_id = integral.subdomain_id()
         subdomain_data = integral.subdomain_data()
 
-        from dune.perftool.ufl.dimensionindex import collect_dimension_index_aliases
-        dimension_index_aliases = collect_dimension_index_aliases(integrand)
+        from dune.perftool.ufl.dimensionindex import dimension_index_mapping
+        dimension_indices = dimension_index_mapping(integrand)
 
         # Now split the given integrand into accumulation expressions
         from dune.perftool.ufl.transformations.extract_accumulation_terms import split_into_accumulation_terms
@@ -170,7 +170,7 @@ def generate_kernel(integrals):
 
         # Get a transformer instance for this kernel
         from dune.perftool.loopy.transformer import UFL2LoopyVisitor
-        visitor = UFL2LoopyVisitor(measure, subdomain_id, dimension_index_aliases)
+        visitor = UFL2LoopyVisitor(measure, subdomain_id, dimension_indices)
 
         # Iterate over the terms and generate a kernel
         for term in accterms:

python/dune/perftool/ufl/dimensionindex.py

@@ -3,30 +3,32 @@
 from ufl.algorithms import MultiFunction
 
 
-class _CollectDimensionIndexAliases(MultiFunction):
+class _DimensionIndexMapping(MultiFunction):
     call = MultiFunction.__call__
 
     def __call__(self, o):
         self.shape = 0
-        return self.call(o)
+        self.dimension_index_dict = {}
+        self.call(o)
 
-    def expr(self, o):
-        return frozenset({}).union(*tuple(self.call(op) for op in o.ufl_operands))
+        return self.dimension_index_dict
 
-    def terminal(self, o):
-        return frozenset({})
+    def expr(self, o):
+        for op in o.ufl_operands:
+            self.call(op)
 
     def function_view(self, o):
-        self.shape = len(o.ufl_operands[1])
-        return frozenset({})
+        from ufl.functionview import select_subelement
+        subelement = select_subelement(o.ufl_operands[0].ufl_element(), o.ufl_operands[1])
+        self.shape = len(subelement.value_shape())
 
     def indexed(self, o):
-        ret = self.call(o.ufl_operands[0])
-        if self.shape:
-            ret = ret.union(frozenset({o.ufl_operands[1][:self.shape][0]}))
+        self.call(o.ufl_operands[0])
+        for i in range(self.shape):
+            from dune.perftool.pdelab.geometry import dimension_iname
+            self.dimension_index_dict[o.ufl_operands[1][i]] = dimension_iname(context='arg', count=i)
         self.shape = 0
-        return ret
 
 
-def collect_dimension_index_aliases(expr):
-    return _CollectDimensionIndexAliases()(expr)
+def dimension_index_mapping(expr):
+    return _DimensionIndexMapping()(expr)