more fixes

python/dune/codegen/blockstructured/basis.py

@@ -21,11 +21,11 @@ from dune.codegen.pdelab.basis import (GenericBasisMixin,
 from dune.codegen.pdelab.driver import (isPk,
                                         isQk,
                                         )
-from dune.codegen.pdelab.geometry import world_dimension
+from dune.codegen.pdelab.geometry import world_dimension, component_iname
 from dune.codegen.pdelab.spaces import type_leaf_gfs, name_lfs
 from dune.codegen.pdelab.restriction import restricted_name
 from dune.codegen.blockstructured.spaces import lfs_inames
-from dune.codegen.blockstructured.tools import tensor_index_to_sequential_index
+from dune.codegen.blockstructured.tools import tensor_index_to_sequential_index, sub_element_inames
 
 from ufl import MixedElement
 
@@ -81,12 +81,35 @@ class BlockStructuredBasisMixin(GenericBasisMixin):
                     read_variables=frozenset({get_pymbolic_basename(qp)}),
                     )
 
-    def implement_trialfunction_gradient(self, element, restriction, index):
-        rawname = "gradu_{}".format(index)
-        name = restricted_name(rawname, restriction)
-        container = name_coefficientcontainer(restriction)
-        self.evaluate_coefficient_gradient(element, name, container, restriction, index)
-        return prim.Variable(name)
+    @kernel_cached
+    def evaluate_coefficient_gradient(self, element, name, container, restriction, index):
+        sub_element = element
+        if isinstance(element, MixedElement):
+            sub_element = element.extract_component(index)[1]
+        from ufl import FiniteElement
+        assert isinstance(sub_element, FiniteElement)
+
+        temporary_variable(name, shape=(element.cell().geometric_dimension(),), managed=True)
+
+        dimindex = component_iname(count=0)
+
+        lfs = name_lfs(element, restriction, index)
+        basis = self.implement_reference_gradient(sub_element, restriction, 0, context='trialgrad')
+        basisindex = get_pymbolic_indices(basis)[:-1]
+        from dune.codegen.tools import maybe_wrap_subscript
+        basis = maybe_wrap_subscript(basis, prim.Variable(dimindex))
+
+        from dune.codegen.blockstructured.argument import pymbolic_coefficient
+        coeff = pymbolic_coefficient(container, lfs, sub_element, basisindex)
+
+        assignee = prim.Subscript(prim.Variable(name), (prim.Variable(dimindex),))
+        reduction_expr = prim.Product((coeff, basis))
+
+        instruction(expression=Reduction("sum", basisindex, reduction_expr, allow_simultaneous=True),
+                    assignee=assignee,
+                    forced_iname_deps=frozenset(self.quadrature_inames() + (dimindex,) + sub_element_inames()),
+                    forced_iname_deps_is_final=True,
+                    )
 
     @kernel_cached
     def evaluate_coefficient(self, element, name, container, restriction, index):
@@ -103,7 +126,6 @@ class BlockStructuredBasisMixin(GenericBasisMixin):
         basis = self.implement_basis(sub_element, restriction, 0, context='trial')
         basisindex = get_pymbolic_indices(basis)[:-1]
 
-        # TODO get rid ot this!
         from dune.codegen.blockstructured.argument import pymbolic_coefficient
         coeff = pymbolic_coefficient(container, lfs, sub_element, basisindex)
 
@@ -112,7 +134,7 @@ class BlockStructuredBasisMixin(GenericBasisMixin):
 
         instruction(expression=Reduction("sum", basisindex, reduction_expr, allow_simultaneous=True),
                     assignee=assignee,
-                    forced_iname_deps=frozenset(self.quadrature_inames()),
+                    forced_iname_deps=frozenset(self.quadrature_inames() + sub_element_inames()),
                     forced_iname_deps_is_final=True,
                     )
 

python/dune/codegen/blockstructured/geometry.py

@@ -21,6 +21,9 @@ from loopy.match import Writes
 
 @geometry_mixin("blockstructured_multilinear")
 class BlockStructuredGeometryMixin(GenericPDELabGeometryMixin):
+    def spatial_coordinate(self, o):
+        return self.to_global(self.quadrature_position_in_macro())
+
     def facet_jacobian_determinant(self, o):
         raise NotImplementedError("I think this case was never really implemented")
 
@@ -35,9 +38,7 @@ class AxiparallelBlockStructuredGeometryMixin(AxiparallelGeometryMixin, BlockStr
                              prim.Power(get_form_option("number_of_blocks"), local_dimension()))
 
     def jacobian_determinant(self, o):
-        name = 'detjac'
-        self.define_jacobian_determinant(name)
-        return prim.Quotient(prim.Variable(name),
+        return prim.Quotient(AxiparallelGeometryMixin.jacobian_determinant(self, o),
                              prim.Power(get_form_option("number_of_blocks"), local_dimension()))
 
     def jacobian_inverse(self, o):

python/dune/codegen/blockstructured/quadrature.py

+from dune.codegen.error import CodegenError
+
 from dune.codegen.generation import (backend,
                                      quadrature_mixin,
                                      )
@@ -8,16 +10,19 @@ import pymbolic.primitives as prim
 
 @quadrature_mixin("blockstructured")
 class BlockstructuredQuadratureMixin(GenericQuadratureMixin):
-    quadrature_position_in_micro = GenericQuadratureMixin.quadrature_position
+    def quadrature_position_in_micro(self, index=None):
+        return GenericQuadratureMixin.quadrature_position(self, index)
 
     def quadrature_position_in_macro(self, index=None):
-        original = quadrature_position_in_micro(self, index)
+        original = self.quadrature_position_in_micro(index)
         qp =  prim.Variable(name_point_in_macro(original, self), )
         if index is not None:
             return prim.Subscript(qp, (index,))
         else:
             return qp
 
+    def quadrature_position(self, index=None):
+        raise CodegenError('Decide if the quadrature point is in the macro or micro element')
 #
 # @backend(interface="quad_pos", name='blockstructured')
 # def pymbolic_quadrature_position():