diff --git a/python/dune/codegen/blockstructured/basis.py b/python/dune/codegen/blockstructured/basis.py index 69097e24dc9763f9b9b7d5cb285bfcda792f1bbc..1f9436f5f61a16b674cd089871f4975d4560befa 100644 --- a/python/dune/codegen/blockstructured/basis.py +++ b/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, ) diff --git a/python/dune/codegen/blockstructured/geometry.py b/python/dune/codegen/blockstructured/geometry.py index 5bb48aebc4c692617fbef78af29e6baab6291d03..dba4940e1b03236db6bd78062b21acedd5a867e8 100644 --- a/python/dune/codegen/blockstructured/geometry.py +++ b/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): diff --git a/python/dune/codegen/blockstructured/quadrature.py b/python/dune/codegen/blockstructured/quadrature.py index ac45eda11a113c9b9d5642d2e09915334c9d9b78..7effba474d44545b230255179e159ff66a4127a9 100644 --- a/python/dune/codegen/blockstructured/quadrature.py +++ b/python/dune/codegen/blockstructured/quadrature.py @@ -1,3 +1,5 @@ +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():