Adds computation of macro index on facets in 3d

python/dune/perftool/blockstructured/tools.py

@@ -27,9 +27,63 @@ def sub_element_inames():
     return inames
 
 
-def tensor_index_to_sequential_index(inames, k):
-    index = prim.Sum(tuple(prim.Variable(name)*k**i for i,name in enumerate(inames)))
-    return index
+def sub_facet_inames():
+    subelem_inames = sub_element_inames()
+    
+    center = pymbolic_in_cell_coordinates(prim.Variable(name_localcenter()), Restriction.NEGATIVE)
+
+    def predicate(index):
+        return prim.Comparison(prim.Subscript(center, index), "==", 0.5)
+
+    def conditional_instruction(index):
+        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(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_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])
+        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
+
+
+def tensor_index_to_sequential_index(indices, k):
+    return prim.Sum(tuple(index * k ** i for i, index in enumerate(indices)))
 
 
 # TODO better name
@@ -39,56 +93,15 @@ def to_tensor_index(iname):
 
 # TODO 3d
 def micro_index_to_macro_index(iname):
-
-    subelem_inames_orig = sub_element_inames()
-
-    k = get_option("number_of_blocks")
-
     it = get_global_context_value("integral_type")
     if it == "cell":
-        sublem_inames = subelem_inames_orig
+        subelem_inames = sub_element_inames()
     elif it == "exterior_facet" or it == "interior_facet":
-        center = pymbolic_in_cell_coordinates(prim.Variable(name_localcenter()), Restriction.NEGATIVE)
-
-        def predicate(x):
-            return frozenset([prim.Comparison(prim.Subscript(center, 0), x, 0.5)])
+        subelem_inames = sub_facet_inames()
 
-        sublem_inames = ("x",)
-        temporary_variable(sublem_inames[0])
-        instruction(assignee=prim.Variable(sublem_inames[0]),
-                    expression=prim.Variable(subelem_inames_orig[0]),
-                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
-                    predicates=predicate("==")
-                    )
-        instruction(assignee=prim.Variable(sublem_inames[0]),
-                    expression=prim.Product(((k-1),prim.Subscript(center, (0,)))),
-                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
-                    predicates=predicate("!=")
-                    )
-        sublem_inames = sublem_inames + ("y",)
-        temporary_variable(sublem_inames[1])
-        instruction(assignee=prim.Variable(sublem_inames[1]),
-                    expression=prim.Product(((k-1),prim.Subscript(center, (1,)))),
-                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
-                    predicates=predicate("==")
-                    )
-        instruction(assignee=prim.Variable(sublem_inames[1]),
-                    expression=prim.Variable(subelem_inames_orig[0]),
-                    within_inames=frozenset(subelem_inames_orig).union(frozenset(quadrature_inames())),
-                    predicates=predicate("!=")
-                    )
-    #
-    # modified_index = prim.Sum((prim.Variable(x), index))
-    # if local_dimension() > 1:
-    #     modified_index = prim.Sum((modified_index, prim.Product(((k+1), prim.Variable(y)))))
-    #     index_div_2 = prim.FloorDiv(index, 2)
-    #     index_div_2 = prim.Product((index_div_2, k-1))
-    #     modified_index = prim.Sum((modified_index, index_div_2))
-    #     if local_dimension() > 2:
-    #         modified_index = prim.Sum((modified_index, prim.Variable(z)))
-
-    modified_index = prim.Sum(tuple(to_tensor_index(iname)[n] * (k + 1) ** n for n in range(world_dimension()))
-                              + (tensor_index_to_sequential_index(sublem_inames, get_option("number_of_blocks") + 1), ))
+    k = get_option("number_of_blocks")
+    modified_index = prim.Sum((tensor_index_to_sequential_index(to_tensor_index(iname), k+1),
+                               tensor_index_to_sequential_index(tuple(prim.Variable(iname) for iname in subelem_inames), k + 1)))
 
     return modified_index