diff --git a/python/dune/codegen/pdelab/argument.py b/python/dune/codegen/pdelab/argument.py
index dc1acd660c137c42be6fb65bb687bafe03fbc730..f7135618bd32d21045860a89bc9666f020652392 100644
--- a/python/dune/codegen/pdelab/argument.py
+++ b/python/dune/codegen/pdelab/argument.py
@@ -14,9 +14,7 @@ from dune.codegen.generation import (domain,
)
from dune.codegen.loopy.target import dtype_floatingpoint
from dune.codegen.pdelab.index import name_index
-from dune.codegen.pdelab.spaces import (lfs_iname,
- name_lfs_bound,
- )
+from dune.codegen.pdelab.spaces import lfs_iname
from dune.codegen.pdelab.restriction import restricted_name
from dune.codegen.ufl.modified_terminals import Restriction
from dune.codegen.options import get_form_option
diff --git a/python/dune/codegen/pdelab/basis.py b/python/dune/codegen/pdelab/basis.py
index 885194af5817bd30d6cc339603ab84f00bc7b990..575eac61385e8f0b2e8651ad442f483a9fe44ab8 100644
--- a/python/dune/codegen/pdelab/basis.py
+++ b/python/dune/codegen/pdelab/basis.py
@@ -73,7 +73,7 @@ class BasisMixinBase(object):
def implement_apply_function_gradient(self, element, restriction, index):
raise NotImplementedError("Basis Mixins should implement linearization point gradient evaluation")
- def implement_grid_function(self, coeff, restriction, grad):
+ def implement_grid_function(self, coeff, restriction, grad, index):
raise NotImplementedError("Basis Mixins should implement grid function evaluation")
@@ -241,15 +241,18 @@ class GenericBasisMixin(BasisMixinBase):
forced_iname_deps_is_final=True,
)
- def implement_gridfunction(self, coeff, restriction, grad):
- name = "coeff{}{}".format(coeff.count(), "_grad" if grad else "")
- self.define_grid_function(name, coeff, restriction, grad)
+ def implement_grid_function(self, coeff, restriction, grad, index):
+ rawname = "coeff{}{}".format(coeff.count(), "_grad" if grad else "")
+ if index is not None:
+ rawname = rawname + '_index{}'.format(index)
+ name = restricted_name(rawname, restriction)
+ self.define_grid_function(name, coeff, restriction, grad, index)
return prim.Subscript(prim.Variable(name), (0,))
- def define_grid_function(self, name, coeff, restriction, grad):
+ def define_grid_function(self, name, coeff, restriction, grad, index):
diffOrder = 1 if grad else 0
- gridfunction = name_gridfunction_member(coeff, restriction, diffOrder)
+ gridfunction = name_gridfunction_member(coeff, restriction, index, diffOrder)
bind_gridfunction_to_element(gridfunction, restriction)
temporary_variable(name,
@@ -258,11 +261,12 @@ class GenericBasisMixin(BasisMixinBase):
managed=False,
)
- quadpos = self.to_cell(self.quadrature_position())
- instruction(code="{} = {}({});".format(name, gridfunction, quadpos),
+ qp = self.to_cell(self.quadrature_position())
+ instruction(code="{} = {}({});".format(name, gridfunction, qp),
assignees=frozenset({name}),
within_inames=frozenset(self.quadrature_inames()),
within_inames_is_final=True,
+ read_variables=frozenset({get_pymbolic_basename(qp)}),
)
diff --git a/python/dune/codegen/pdelab/localoperator.py b/python/dune/codegen/pdelab/localoperator.py
index 835fe44e24df9f0849391e1f50422512df43c044..2ac71789809d3531b79cd2eb817afe7a2088eff7 100644
--- a/python/dune/codegen/pdelab/localoperator.py
+++ b/python/dune/codegen/pdelab/localoperator.py
@@ -167,29 +167,33 @@ def localoperator_basename(form_ident):
return get_form_option("classname", form_ident)
-def name_gridfunction_member(coeff, restriction, diffOrder=0):
+def name_gridfunction_member(coeff, restriction, index, diffOrder=0):
# We reuse the grid function for volume integrals in skeleton integrals
if restriction == Restriction.POSITIVE:
restriction = Restriction.NONE
restr = "_n" if restriction == Restriction.NEGATIVE else ""
name = "local_gridfunction_coeff{}_diff{}{}".format(coeff.count(), diffOrder, restr)
- define_gridfunction_member(name, coeff, restriction, diffOrder)
+ if index is not None:
+ name = name + '_index{}'.format(index)
+ define_gridfunction_member(name, coeff, restriction, index, diffOrder)
return name
-def name_gridfunction_constructor_argument(coeff):
- _type = type_gridfunction_template_parameter(coeff)
+def name_gridfunction_constructor_argument(coeff, index):
+ _type = type_gridfunction_template_parameter(coeff, index)
name = "gridfunction_coeff{}_".format(coeff.count())
+ if index is not None:
+ name = name + '_index{}'.format(index)
constructor_parameter("const {}&".format(_type), name, classtag="operator")
return name
@class_member(classtag="operator")
-def define_gridfunction_member(name, coeff, restriction, diffOrder):
- _type = type_gridfunction_template_parameter(coeff)
- param = name_gridfunction_constructor_argument(coeff)
+def define_gridfunction_member(name, coeff, restriction, index, diffOrder):
+ _type = type_gridfunction_template_parameter(coeff, index)
+ param = name_gridfunction_constructor_argument(coeff, index)
if diffOrder > 0:
- other = name_gridfunction_member(coeff, restriction, diffOrder - 1)
+ other = name_gridfunction_member(coeff, restriction, index, diffOrder - 1)
init = "derivative({})".format(other)
initializer_list(name, [init], classtag="operator")
return "mutable decltype({}) {};".format(init, name)
@@ -200,8 +204,11 @@ def define_gridfunction_member(name, coeff, restriction, diffOrder):
@template_parameter(classtag="operator")
-def type_gridfunction_template_parameter(coeff):
- return "GRIDFUNCTION_COEFF{}".format(coeff.count())
+def type_gridfunction_template_parameter(coeff, index):
+ name = "GRIDFUNCTION_COEFF{}".format(coeff.count())
+ if index is not None:
+ name = name + '_INDEX{}'.format(index)
+ return name
def class_type_from_cache(classtag):
@@ -261,7 +268,7 @@ def determine_accumulation_space(info, number):
subel = element.extract_component(info.element_index)[1]
# And generate a local function space for it!
- from dune.codegen.pdelab.spaces import name_lfs, name_lfs_bound, lfs_iname
+ from dune.codegen.pdelab.spaces import name_lfs, lfs_iname
lfs = name_lfs(element, info.restriction, info.element_index)
from dune.codegen.generation import valuearg
from loopy.types import NumpyType
@@ -874,7 +881,16 @@ def local_operator_default_settings(operator, form):
# Iterate over the needed grid functions in correct order
for c in sorted(filter(lambda c: c.count() > 2, form.coefficients()), key=lambda c: c.count()):
- name_gridfunction_constructor_argument(c)
+ from ufl import MixedElement, VectorElement
+ if isinstance(c.ufl_element(), MixedElement):
+ # For non vector elements we need to do something more evolved. We
+ # probably would need to extract the correct index (tuple?) and
+ # pass that to name_gridfunction_constructor_argument
+ assert isinstance(c.ufl_element(), VectorElement)
+ for i in range(c.ufl_element().num_sub_elements()):
+ name_gridfunction_constructor_argument(c, i)
+ else:
+ name_gridfunction_constructor_argument(c, None)
# Add right base classes for stationary/instationary operators
base_class('Dune::PDELab::LocalOperatorDefaultFlags', classtag="operator")
diff --git a/python/dune/codegen/pdelab/spaces.py b/python/dune/codegen/pdelab/spaces.py
index 19de9bc7e7d0fa8bf14327eea39ff559d3450da7..d22cbe26b1c01cd8d9187e36836171595b441a70 100644
--- a/python/dune/codegen/pdelab/spaces.py
+++ b/python/dune/codegen/pdelab/spaces.py
@@ -72,10 +72,14 @@ def define_lfs(name, father, child):
return "auto {} = child({}, _{});".format(name, father, child)
-@preamble
def define_lfs_size(lfs, element, restriction):
name = name_lfs_bound(lfs)
valuearg(name, dtype=numpy.int32)
+ _define_lfs_size(name, lfs)
+
+
+@preamble
+def _define_lfs_size(name, lfs):
return "auto {} = {}.size();".format(name, lfs)
diff --git a/python/dune/codegen/sumfact/basis.py b/python/dune/codegen/sumfact/basis.py
index 83e5206b5289719f3c4e85f4ca523976cb6feb79..354e096a5ea7343ae4dfaecc6f63b7c489df196a 100644
--- a/python/dune/codegen/sumfact/basis.py
+++ b/python/dune/codegen/sumfact/basis.py
@@ -40,7 +40,7 @@ from dune.codegen.sumfact.symbolic import SumfactKernel, SumfactKernelInterfaceB
from dune.codegen.options import get_form_option
from dune.codegen.pdelab.driver import FEM_name_mangling
from dune.codegen.pdelab.restriction import restricted_name
-from dune.codegen.pdelab.spaces import name_lfs, name_lfs_bound, name_leaf_lfs
+from dune.codegen.pdelab.spaces import name_lfs, name_leaf_lfs
from dune.codegen.tools import maybe_wrap_subscript, ImmutableCuttingRecord
from dune.codegen.pdelab.basis import shape_as_pymbolic
from dune.codegen.sumfact.accumulation import sumfact_iname
diff --git a/python/dune/codegen/ufl/visitor.py b/python/dune/codegen/ufl/visitor.py
index e774e03d90e8bc61a1108767867ab046bfa56c81..cbc5b1c76a03d85419e87cf67b268516dffac2c4 100644
--- a/python/dune/codegen/ufl/visitor.py
+++ b/python/dune/codegen/ufl/visitor.py
@@ -172,10 +172,18 @@ class UFL2LoopyVisitor(ModifiedTerminalTracker):
# and exports it through a getter method 'getTime'
return prim.Call(prim.Variable("getTime"), ())
else:
+ index = None
+ if isinstance(o.ufl_element(), MixedElement):
+ index = self.indices[0]
+ assert isinstance(index, int)
+ self.indices = self.indices[1:]
+ if len(self.indices) == 0:
+ self.indices = None
+
if self.reference_grad:
raise CodegenUFLError("Coefficient gradients should not be transformed to reference element")
- return self.implement_gridfunction(o, restriction, self.grad)
+ return self.implement_grid_function(o, restriction, self.grad, index)
def variable(self, o):
# Right now only scalar varibables are supported