def fiat_equivalent(self):
if self.is_mixed:
# EnrichedElement is actually a MixedElement
return FIAT.MixedElement(
[e.element.fiat_equivalent for e in self.elements],
ref_el=self.cell)
else:
return FIAT.EnrichedElement(*(e.fiat_equivalent
for e in self.elements))
def _create_vector_finiteelement(
element: ufl.VectorElement) -> FIAT.MixedElement:
fiat_element = FIAT.MixedElement(
map(_create_element, element.sub_elements()))
def reorder_for_vector_element(item, block_size):
"""Reorder the elements in item from XXYYZZ ordering to XYZXYZ."""
space_dim = len(item) // block_size
return [
item[i] for block in range(space_dim)
for i in range(block, len(item), space_dim)
]
def calculate_entity_dofs_of_vector_element(entity_dofs, block_size):
"""Get the entity DOFs of a VectorElement with XYZXYZ ordering."""
return {
dim: {
entity: [
block_size * i + j for i in e_dofs
for j in range(block_size)
]
for entity, e_dofs in dofs.items()
}
for dim, dofs in entity_dofs.items()
}
# Reorder from XXYYZZ to XYZXYZ
block_size = fiat_element.num_sub_elements()
fiat_element.mapping = types.MethodType(
lambda self:
[m for m in self._elements[0].mapping() for e in self._elements],
fiat_element)
fiat_element.dual.nodes = reorder_for_vector_element(
fiat_element.dual.nodes, block_size)
fiat_element.dual.entity_ids = calculate_entity_dofs_of_vector_element(
fiat_element.elements()[0].dual.entity_ids, block_size)
fiat_element.dual.entity_closure_ids = calculate_entity_dofs_of_vector_element(
fiat_element.elements()[0].dual.entity_closure_ids, block_size)
fiat_element.old_tabulate = fiat_element.tabulate
def tabulate(self, order, points, entity=None):
block_size = self.num_sub_elements()
scalar_dofs = len(self.dual.nodes) // block_size
return {
i: numpy.array([
item[j] for dim in range(scalar_dofs)
for j in range(dim, len(item), scalar_dofs)
])
for i, item in self.old_tabulate(order, points,
entity=entity).items()
}
fiat_element.tabulate = types.MethodType(tabulate, fiat_element)
return fiat_element
def fiat_equivalent(self):
# Avoid circular import dependency
from finat.mixed import MixedSubElement
if all(isinstance(e, MixedSubElement) for e in self.elements):
# EnrichedElement is actually a MixedElement
return FIAT.MixedElement([e.element.fiat_equivalent
for e in self.elements], ref_el=self.cell)
else:
return FIAT.EnrichedElement(*(e.fiat_equivalent
for e in self.elements))
def _create_mixed_finiteelement(
element: ufl.MixedElement) -> FIAT.MixedElement:
elements = []
def rextract(els):
for e in els:
if isinstance(e, ufl.MixedElement) \
and not isinstance(e, ufl.VectorElement) \
and not isinstance(e, ufl.TensorElement):
rextract(e.sub_elements())
else:
elements.append(e)
rextract(element.sub_elements())
return FIAT.MixedElement(map(_create_element, elements))
def _(element, vector_is_mixed):
# If we're just trying to get the scalar part of a vector element?
if not vector_is_mixed:
assert isinstance(element, (ufl.VectorElement, ufl.TensorElement))
return create_element(element.sub_elements()[0], vector_is_mixed)
elements = []
def rec(eles):
for ele in eles:
if isinstance(ele, ufl.MixedElement):
rec(ele.sub_elements())
else:
elements.append(ele)
rec(element.sub_elements())
fiat_elements = map(
partial(create_element, vector_is_mixed=vector_is_mixed), elements)
return FIAT.MixedElement(fiat_elements)
