diff --git a/tsfc/fem.py b/tsfc/fem.py index de8f5ee6..27439f1c 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -5,33 +5,30 @@ import itertools from functools import singledispatch +import gem import numpy - import ufl -from ufl.corealg.map_dag import map_expr_dag, map_expr_dags -from ufl.corealg.multifunction import MultiFunction -from ufl.classes import ( - Argument, CellCoordinate, CellEdgeVectors, CellFacetJacobian, - CellOrientation, CellOrigin, CellVertices, CellVolume, Coefficient, - FacetArea, FacetCoordinate, GeometricQuantity, Jacobian, - JacobianDeterminant, NegativeRestricted, QuadratureWeight, - PositiveRestricted, ReferenceCellVolume, ReferenceCellEdgeVectors, - ReferenceFacetVolume, ReferenceNormal, SpatialCoordinate -) -from ufl.domain import extract_unique_domain - -from FIAT.reference_element import make_affine_mapping -from FIAT.reference_element import UFCSimplex - -import gem +from FIAT.reference_element import UFCSimplex, make_affine_mapping +from finat.physically_mapped import (NeedsCoordinateMappingElement, + PhysicalGeometry) +from finat.point_set import PointSet, PointSingleton +from finat.quadrature import make_quadrature from gem.node import traversal -from gem.optimise import ffc_rounding, constant_fold_zero +from gem.optimise import constant_fold_zero, ffc_rounding from gem.unconcatenate import unconcatenate from gem.utils import cached_property - -from finat.physically_mapped import PhysicalGeometry, NeedsCoordinateMappingElement -from finat.point_set import PointSet, PointSingleton -from finat.quadrature import make_quadrature +from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors, + CellFacetJacobian, CellOrientation, CellOrigin, + CellVertices, CellVolume, Coefficient, FacetArea, + FacetCoordinate, GeometricQuantity, Jacobian, + JacobianDeterminant, NegativeRestricted, + PositiveRestricted, QuadratureWeight, + ReferenceCellEdgeVectors, ReferenceCellVolume, + ReferenceFacetVolume, ReferenceNormal, + SpatialCoordinate) +from ufl.corealg.map_dag import map_expr_dag, map_expr_dags +from ufl.corealg.multifunction import MultiFunction +from ufl.domain import extract_unique_domain from tsfc import ufl2gem from tsfc.finatinterface import as_fiat_cell, create_element @@ -40,8 +37,8 @@ construct_modified_terminal) from tsfc.parameters import is_complex from tsfc.ufl_utils import (ModifiedTerminalMixin, PickRestriction, - entity_avg, one_times, simplify_abs, - preprocess_expression, TSFCConstantMixin) + TSFCConstantMixin, entity_avg, one_times, + preprocess_expression, simplify_abs) class ContextBase(ProxyKernelInterface): @@ -175,31 +172,35 @@ def detJ_at(self, point): return map_expr_dag(context.translator, expr) def reference_normals(self): - if not (isinstance(self.interface.fiat_cell, UFCSimplex) and - self.interface.fiat_cell.get_spatial_dimension() == 2): - raise NotImplementedError("Only works for triangles for now") - return gem.Literal(numpy.asarray([self.interface.fiat_cell.compute_normal(i) for i in range(3)])) + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_faces = len(cell.get_topology()[sd-1]) + + return gem.Literal(numpy.asarray([cell.compute_normal(i) for i in range(num_faces)])) def reference_edge_tangents(self): - return gem.Literal(numpy.asarray([self.interface.fiat_cell.compute_edge_tangent(i) for i in range(3)])) + cell = self.interface.fiat_cell + num_edges = len(cell.get_topology()[1]) + return gem.Literal(numpy.asarray([cell.compute_edge_tangent(i) for i in range(num_edges)])) def physical_tangents(self): - if not (isinstance(self.interface.fiat_cell, UFCSimplex) and - self.interface.fiat_cell.get_spatial_dimension() == 2): - raise NotImplementedError("Only works for triangles for now") - - rts = [self.interface.fiat_cell.compute_tangents(1, f)[0] for f in range(3)] - jac = self.jacobian_at([1/3, 1/3]) - + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_edges = len(cell.get_topology()[1]) els = self.physical_edge_lengths() + rts = gem.ListTensor([cell.compute_tangents(1, i)[0] / els[i] for i in range(num_edges)]) + jac = self.jacobian_at(cell.make_points(sd, 0, sd+1)[0]) - return gem.ListTensor([[(jac[0, 0]*rts[i][0] + jac[0, 1]*rts[i][1]) / els[i], - (jac[1, 0]*rts[i][0] + jac[1, 1]*rts[i][1]) / els[i]] - for i in range(3)]) + return rts @ jac.T def physical_normals(self): + cell = self.interface.fiat_cell + if not (isinstance(cell, UFCSimplex) and cell.get_dimension() == 2): + raise NotImplementedError("Can't do physical normals on that cell yet") + + num_edges = len(cell.get_topology()[1]) pts = self.physical_tangents() - return gem.ListTensor([[pts[i, 1], -1*pts[i, 0]] for i in range(3)]) + return gem.ListTensor([[pts[i, 1], -1*pts[i, 0]] for i in range(num_edges)]) def physical_edge_lengths(self): expr = ufl.classes.CellEdgeVectors(extract_unique_domain(self.mt.terminal)) @@ -208,8 +209,11 @@ def physical_edge_lengths(self): elif self.mt.restriction == '-': expr = NegativeRestricted(expr) - expr = ufl.as_vector([ufl.sqrt(ufl.dot(expr[i, :], expr[i, :])) for i in range(3)]) - config = {"point_set": PointSingleton([1/3, 1/3])} + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_edges = len(cell.get_topology()[1]) + expr = ufl.as_vector([ufl.sqrt(ufl.dot(expr[i, :], expr[i, :])) for i in range(num_edges)]) + config = {"point_set": PointSingleton(cell.make_points(sd, 0, sd+1)[0])} config.update(self.config) context = PointSetContext(**config) expr = self.preprocess(expr, context) @@ -443,7 +447,8 @@ def callback(facet_i): @translate.register(ReferenceCellEdgeVectors) def translate_reference_cell_edge_vectors(terminal, mt, ctx): - from FIAT.reference_element import TensorProductCell as fiat_TensorProductCell + from FIAT.reference_element import \ + TensorProductCell as fiat_TensorProductCell fiat_cell = ctx.fiat_cell if isinstance(fiat_cell, fiat_TensorProductCell): raise NotImplementedError("ReferenceCellEdgeVectors not implemented on TensorProductElements yet") diff --git a/tsfc/finatinterface.py b/tsfc/finatinterface.py index 320b3e32..659a69dc 100644 --- a/tsfc/finatinterface.py +++ b/tsfc/finatinterface.py @@ -46,6 +46,7 @@ "Gauss-Lobatto-Legendre": finat.GaussLobattoLegendre, "HDiv Trace": finat.HDivTrace, "Hellan-Herrmann-Johnson": finat.HellanHerrmannJohnson, + "Johnson-Mercier": finat.JohnsonMercier, "Nonconforming Arnold-Winther": finat.ArnoldWintherNC, "Conforming Arnold-Winther": finat.ArnoldWinther, "Hermite": finat.Hermite,