reworks neumann bc to use predicates. Only works in 2d

python/dune/codegen/blockstructured/accumulation.py

@@ -2,11 +2,11 @@ from dune.codegen.blockstructured.tools import sub_element_inames
 from dune.codegen.generation import accumulation_mixin, instruction
 from dune.codegen.loopy.target import dtype_floatingpoint
 from dune.codegen.options import get_form_option
-from dune.codegen.pdelab.geometry import world_dimension
+from dune.codegen.pdelab.geometry import world_dimension, name_intersection_geometry_wrapper
 from dune.codegen.pdelab.localoperator import determine_accumulation_space, GenericAccumulationMixin
 from dune.codegen.pdelab.argument import name_accumulation_variable
 from dune.codegen.pdelab.localoperator import boundary_predicates
-from dune.codegen.generation.loopy import function_mangler, globalarg
+from dune.codegen.generation.loopy import function_mangler, globalarg, temporary_variable
 import loopy as lp
 import pymbolic.primitives as prim
 
@@ -17,9 +17,9 @@ from loopy.match import Writes
 class BlockStructuredAccumulationMixin(GenericAccumulationMixin):
     def generate_accumulation_instruction(self, expr):
         if get_form_option('vectorization_blockstructured'):
-            return generate_accumulation_instruction(expr, self)
+            return generate_accumulation_instruction_vectorized(expr, self)
         else:
-            return GenericAccumulationMixin.generate_accumulation_instruction(self, expr)
+            return generate_accumulation_instruction(expr, self)
 
 
 def name_accumulation_alias(container, accumspace):
@@ -50,17 +50,77 @@ def residual_weight_mangler(knl, func, arg_dtypes):
         return lp.CallMangleInfo(func, (lp.types.NumpyType(dtype_floatingpoint()),), ())
 
 
+def blockstructured_boundary_predicated(measure, subdomain_id):
+    predicates = []
+
+    if subdomain_id not in ['everywhere', 'otherwise']:
+        subelem_inames = sub_element_inames()
+
+        face_id = "face_id"
+        temporary_variable(face_id, shape=())
+        instruction(code="{} = {}.indexInInside();".format(face_id, name_intersection_geometry_wrapper()),
+                    assignees=(face_id,))
+
+        def face_id_equals(id):
+            return prim.Comparison(prim.Variable(face_id), "==", id)
+
+        k = get_form_option("number_of_blocks")
+
+        if(world_dimension() == 2):
+            predicates.append(prim.If(face_id_equals(0), prim.Comparison(prim.Variable(subelem_inames[0]), "==", 0), True))
+            predicates.append(prim.If(face_id_equals(1), prim.Comparison(prim.Variable(subelem_inames[0]), "==", k - 1), True))
+            predicates.append(prim.If(face_id_equals(2), prim.Comparison(prim.Variable(subelem_inames[1]), "==", 0), True))
+            predicates.append(prim.If(face_id_equals(3), prim.Comparison(prim.Variable(subelem_inames[1]), "==", k - 1), True))
+        else:
+            raise NotImplementedError()
+    return frozenset(predicates)
+
+
 def generate_accumulation_instruction(expr, visitor):
     # Collect the lfs and lfs indices for the accumulate call
     test_lfs = determine_accumulation_space(visitor.test_info, 0)
     # In the jacobian case, also determine the space for the ansatz space
     ansatz_lfs = determine_accumulation_space(visitor.trial_info, 1)
 
+    # Collect the lfs and lfs indices for the accumulate call
+    accumvar = name_accumulation_variable(test_lfs.get_restriction() + ansatz_lfs.get_restriction())
+
+    predicates = boundary_predicates(visitor.measure, visitor.subdomain_id)
+    predicates = predicates.union(blockstructured_boundary_predicated(visitor.measure, visitor.subdomain_id))
+
+    rank = 1 if ansatz_lfs.lfs is None else 2
+
+    from dune.codegen.pdelab.argument import PDELabAccumulationFunction
+    from pymbolic.primitives import Call
+    accexpr = Call(PDELabAccumulationFunction(accumvar, rank),
+                   (test_lfs.get_args() + ansatz_lfs.get_args() + (expr,))
+                   )
+
+    from dune.codegen.generation import instruction
+    quad_inames = visitor.quadrature_inames()
+    lfs_inames = frozenset(visitor.test_info.inames)
+    if visitor.trial_info:
+        lfs_inames = lfs_inames.union(visitor.trial_info.inames)
+
+    instruction(assignees=(),
+                expression=accexpr,
+                forced_iname_deps=lfs_inames.union(frozenset(quad_inames)),
+                forced_iname_deps_is_final=True,
+                predicates=predicates
+                )
+
+def generate_accumulation_instruction_vectorized(expr, visitor):
+    # Collect the lfs and lfs indices for the accumulate call
+    test_lfs = determine_accumulation_space(visitor.test_info, 0)
+    # In the jacobian case, also determine the space for the ansatz space
+    ansatz_lfs = determine_accumulation_space(visitor.trial_info, 1)
+
     # Collect the lfs and lfs indices for the accumulate call
     accumvar = name_accumulation_variable(test_lfs.get_restriction() + ansatz_lfs.get_restriction())
     accumvar_alias = name_accumulation_alias(accumvar, test_lfs)
 
     predicates = boundary_predicates(visitor.measure, visitor.subdomain_id)
+    predicates.append(blockstructured_boundary_predicated(visitor.measure, visitor.subdomain_id))
 
     quad_inames = visitor.quadrature_inames()
     lfs_inames = visitor.test_info.inames

python/dune/codegen/blockstructured/geometry.py

@@ -15,9 +15,12 @@ from dune.codegen.pdelab.geometry import (AxiparallelGeometryMixin,
                                           component_iname,
                                           enforce_boundary_restriction,
                                           restricted_name,
-                                          name_geometry,
+                                          name_cell_geometry
                                           )
-from dune.codegen.blockstructured.tools import sub_element_inames
+from dune.codegen.blockstructured.tools import (sub_element_inames,
+                                                name_point_in_macro,
+                                                )
+from dune.codegen.ufl.modified_terminals import Restriction
 import pymbolic.primitives as prim
 from loopy.match import Writes
 
@@ -25,22 +28,27 @@ 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())
+        if self.measure == 'cell':
+            return self.to_global(self.quadrature_position_in_macro())
+        else:
+            assert self.measure == 'exterior_facet' or self.measure == 'interior_facet'
+            micro_qp = self.to_cell(self.quadrature_position_in_micro())
+
+            macro_qp = prim.Variable(name_point_in_macro(micro_qp, self), )
+
+            return self.to_global(macro_qp)
 
     def to_global(self, local):
         assert isinstance(local, prim.Expression)
         name = get_pymbolic_basename(local) + "_global"
 
-        if self.measure == 'cell':
-            compute_multilinear_to_global_transformation(name, local, self)
-            return prim.Variable(name)
-        elif self.measure == 'exterior_facet' or self.measure == 'interior_facet':
-            return GenericPDELabGeometryMixin.to_global(self, local)
-        else:
-            raise NotImplementedError
+        # TODO need to assert somehow that local is of codim 0
+        compute_multilinear_to_global_transformation(name, local, self)
+        return prim.Variable(name)
 
     def facet_jacobian_determinant(self, o):
-        raise NotImplementedError("I think this case was never really implemented")
+        return prim.Quotient(GenericPDELabGeometryMixin.facet_jacobian_determinant(self, o),
+                             prim.Power(get_form_option("number_of_blocks"), local_dimension()))
 
     def jacobian_determinant(self, o):
         name = 'detjac'
@@ -87,16 +95,9 @@ class AxiparallelBlockStructuredGeometryMixin(AxiparallelGeometryMixin, BlockStr
         assert isinstance(local, prim.Expression)
         name = get_pymbolic_basename(local) + "_global"
 
-        if self.measure == 'cell':
-            compute_axiparallel_to_global_transformation(name, local, self)
-            return prim.Variable(name)
-        elif self.measure == 'exterior_facet' or self.measure == 'interior_facet':
-            return AxiparallelGeometryMixin.to_global(self, local)
-        else:
-            raise NotImplementedError
-
-    def facet_jacobian_determinant(self, o):
-        raise NotImplementedError("I think this case was never really implemented")
+        # TODO need to assert somehow that local is of codim 0
+        compute_axiparallel_to_global_transformation(name, local, self)
+        return prim.Variable(name)
 
     def jacobian_determinant(self, o):
         return prim.Quotient(AxiparallelGeometryMixin.jacobian_determinant(self, o),
@@ -328,59 +329,6 @@ def name_jacobian_inverse_transposed(restriction, visitor):
     define_jacobian_inverse_transposed(name, visitor)
     return name
 
-#
-# # scale Jacobian according to the order of the blockstructure
-# def pymbolic_jacobian_inverse(i, j, restriction):
-#     if get_form_option("constant_transformation_matrix"):
-#         from dune.codegen.pdelab.geometry import name_constant_jacobian_inverse_transposed
-#         name_jit = name_constant_jacobian_inverse_transposed(restriction)
-#     else:
-#         name_jit = name_jacobian_inverse_transposed(restriction)
-#     return prim.Product((get_form_option("number_of_blocks"),
-#                          prim.Subscript(prim.Variable(name_jit), (j, i))))
-#
-#
-# # scale determinant according to the order of the blockstructure
-# def pymbolic_facet_jacobian_determinant():
-#     from dune.codegen.pdelab.geometry import name_facet_jacobian_determinant
-#     return prim.Quotient(prim.Variable(name_facet_jacobian_determinant()),
-#                          prim.Power(get_form_option("number_of_blocks"), local_dimension()))
-
-
-# translate a point in the micro element into macro coordinates
-def define_point_in_macro(name, point_in_micro, visitor):
-    dim = local_dimension()
-    if get_form_option('vectorization_blockstructured'):
-        temporary_variable(name, shape=(dim,), managed=True)
-    else:
-        temporary_variable(name, shape=(dim,), shape_impl=('fv',))
-
-    # point_macro = (point_micro + index_micro) / number_of_blocks
-    # where index_micro = tensor index of the micro element
-    subelem_inames = sub_element_inames()
-    for i in range(dim):
-        if isinstance(point_in_micro, prim.Subscript):
-            expr = prim.Subscript(point_in_micro.aggregate, point_in_micro.index + (i,))
-        else:
-            expr = prim.Subscript(point_in_micro, (i,))
-        expr = prim.Sum((expr, prim.Variable(subelem_inames[i]),))
-        expr = prim.Quotient(expr, get_form_option('number_of_blocks'))
-        # TODO relax within inames
-        instruction(assignee=prim.Subscript(prim.Variable(name), (i,)),
-                    expression=expr,
-                    within_inames=frozenset(subelem_inames + visitor.quadrature_inames()),
-                    tags=frozenset({subelem_inames[i]})
-                    )
-
-
-# TODO add subelem inames if this function gets called
-# TODO change input parameter to string
-def name_point_in_macro(point_in_micro, visitor):
-    assert isinstance(point_in_micro, prim.Expression)
-    name = get_pymbolic_basename(point_in_micro) + "_macro"
-    define_point_in_macro(name, point_in_micro, visitor)
-    return name
-
 
 def compute_multilinear_to_global_transformation(name, local, visitor):
     dim = world_dimension()
@@ -445,12 +393,20 @@ def compute_axiparallel_to_global_transformation(name, local, visitor):
 def define_element_corners(name):
     from dune.codegen.pdelab.driver import get_form
     n_corners = get_form().ufl_cell().num_vertices()
-    temporary_variable(name, shape_impl=('fv', 'fv'), shape=(n_corners, local_dimension()))
+    temporary_variable(name, shape_impl=('fv', 'fv'), shape=(n_corners, world_dimension()))
 
     iname = "i_corner"
     domain(iname, n_corners)
-    from dune.codegen.generation import instruction
-    instruction(code="{0}[{1}] = {2}.corner({3});".format(name, iname, name_geometry(), iname),
+
+    it = get_global_context_value("integral_type")
+    if it == 'cell':
+        restriction = Restriction.NONE
+    elif it == 'exterior_facet':
+        restriction = Restriction.POSITIVE
+    else:
+        raise NotImplementedError()
+
+    instruction(code="{0}[{1}] = {2}.corner({3});".format(name, iname, name_cell_geometry(restriction), iname),
                 assignees=frozenset({name}),
                 within_inames=frozenset({iname}), within_inames_is_final=True)
 

python/dune/codegen/blockstructured/tools.py

@@ -2,14 +2,10 @@ from dune.codegen.ufl.modified_terminals import Restriction
 from dune.codegen.tools import get_pymbolic_basename
 from dune.codegen.generation import (iname,
                                      domain,
-                                     get_global_context_value,
                                      temporary_variable,
                                      instruction)
-from dune.codegen.pdelab.geometry import (local_dimension,
-                                          world_dimension,
-                                          name_localcenter,)
-
-from dune.codegen.generation.counter import get_counted_variable
+from dune.codegen.pdelab.geometry import (world_dimension,
+                                          )
 from dune.codegen.options import get_form_option
 import pymbolic.primitives as prim
 
@@ -19,7 +15,7 @@ import pymbolic.primitives as prim
 @iname
 def sub_element_inames():
     name = "subel"
-    dim = local_dimension()
+    dim = world_dimension()
     dim_names = ["x", "y", "z"] + [str(i) for i in range(4, dim + 1)]
     inames = tuple()
     for i in range(dim):
@@ -28,72 +24,6 @@ def sub_element_inames():
     return inames
 
 
-# define inames for boundary integration
-# In the case of boundary integrationsince we have only d-1 loops,
-# but we need always d inames to compute the macro index.
-# Therefore we must find out which iname is constant.
-# Since loo.py cannot handle if-else blocks very well, this whole
-# computation seems a bit cumbersome.
-def sub_facet_inames():
-    subelem_inames = sub_element_inames()
-
-    center = pymbolic_in_cell_coordinates(prim.Variable(name_localcenter()), Restriction.POSITIVE)
-
-    # check if iname[index] must be constant or not
-    def predicate(index):
-        return prim.Comparison(prim.Subscript(center, index), "==", 0.5)
-
-    def conditional_instruction(index):
-        # instruction for not constant iname
-        # special case for the third iname
-        instruction(assignee=prim.Variable(inames[index]),
-                    expression=prim.Variable(subelem_inames[1 if index == 2 else 0]),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
-                    predicates=frozenset([predicate(index)])
-                    )
-        # instruction for constant iname
-        instruction(assignee=prim.Variable(inames[index]),
-                    expression=prim.Product(((k - 1), prim.Subscript(center, (index,)))),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
-                    predicates=frozenset([prim.LogicalNot(predicate(index))])
-                    )
-
-    k = get_form_option("number_of_blocks")
-
-    inames = ("x",)
-    temporary_variable(inames[0])
-    conditional_instruction(0)
-    if world_dimension() < 3:
-        inames = inames + ("y",)
-        temporary_variable(inames[1])
-        conditional_instruction(1)
-    else:
-        inames = inames + ("y",)
-        temporary_variable(inames[1])
-        # one additional condition is needed in 3d for the second iname
-        instruction(assignee=prim.Variable(inames[1]),
-                    expression=prim.Variable(subelem_inames[0]),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
-                    predicates=frozenset([prim.LogicalAnd((predicate(1), predicate(2)))])
-                    )
-        instruction(assignee=prim.Variable(inames[1]),
-                    expression=prim.Variable(subelem_inames[1]),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
-                    predicates=frozenset([prim.LogicalAnd((predicate(1), predicate(0)))])
-                    )
-        instruction(assignee=prim.Variable(inames[1]),
-                    expression=prim.Product(((k - 1), prim.Subscript(center, (1,)))),
-                    within_inames=frozenset(subelem_inames).union(frozenset(quadrature_inames())),
-                    predicates=frozenset([prim.LogicalNot(predicate(1))])
-                    )
-
-        inames = inames + ("z",)
-        temporary_variable(inames[2])
-        conditional_instruction(2)
-
-    return inames
-
-
 # compute sequential index for given tensor index, the same as index in base-k to base-10
 def tensor_index_to_sequential_index(indices, k):
     return prim.Sum(tuple(prim.Variable(index) * k ** i for i, index in enumerate(indices)))
@@ -106,11 +36,7 @@ def sequential_index_to_tensor_index(iname, k):
 
 # compute index for higher order FEM for a given Q1 index
 def micro_index_to_macro_index(element, inames):
-    it = get_global_context_value("integral_type")
-    if it == "cell":
-        subelem_inames = sub_element_inames()
-    elif it == "exterior_facet" or it == "interior_facet":
-        subelem_inames = sub_facet_inames()
+    subelem_inames = sub_element_inames()
 
     k = get_form_option("number_of_blocks")
     p = element.degree()
@@ -120,7 +46,8 @@ def micro_index_to_macro_index(element, inames):
 
 # translate a point in the micro element into macro coordinates
 def define_point_in_macro(name, point_in_micro, visitor):
-    dim = local_dimension()
+    # TODO this won't work for 2d mannifolds
+    dim = world_dimension()
     if get_form_option('vectorization_blockstructured'):
         temporary_variable(name, shape=(dim,), managed=True)
     else: