Gem tabulation

FIAT/barycentric_interpolation.py

@@ -25,22 +25,24 @@ def barycentric_interpolation(nodes, wts, dmat, pts, order=0):
     https://doi.org/10.1137/S0036144502417715 Eq. (4.2) & (9.4)
     """
     if pts.dtype == object:
-        from sympy import simplify
-        sp_simplify = numpy.vectorize(simplify)
+        # Do not use barycentric interpolation at unknown points
+        phi = numpy.add.outer(-nodes, pts.flatten())
+        phis = [wi * numpy.prod(phi[:i], axis=0) * numpy.prod(phi[i+1:], axis=0) for i, wi in enumerate(wts)]
+        phi = numpy.asarray(phis)
     else:
-        sp_simplify = lambda x: x
-    phi = numpy.add.outer(-nodes, pts.flatten())
-    with numpy.errstate(divide='ignore', invalid='ignore'):
-        numpy.reciprocal(phi, out=phi)
-        numpy.multiply(phi, wts[:, None], out=phi)
-        numpy.multiply(1.0 / numpy.sum(phi, axis=0), phi, out=phi)
-    phi[phi != phi] = 1.0
-    phi = phi.reshape(-1, *pts.shape[:-1])
+        # Use the second barycentric interpolation formula
+        phi = numpy.add.outer(-nodes, pts.flatten())
+        with numpy.errstate(divide='ignore', invalid='ignore'):
+            numpy.reciprocal(phi, out=phi)
+            numpy.multiply(phi, wts[:, None], out=phi)
+            numpy.multiply(1.0 / numpy.sum(phi, axis=0), phi, out=phi)
+        # Replace nan with one
+        phi[phi != phi] = 1.0
 
-    phi = sp_simplify(phi)
+    phi = phi.reshape(-1, *pts.shape[:-1])
     results = {(0,): phi}
     for r in range(1, order+1):
-        phi = sp_simplify(numpy.dot(dmat, phi))
+        phi = numpy.dot(dmat, phi)
         results[(r,)] = phi
     return results
 

FIAT/expansions.py

@@ -86,19 +86,19 @@ def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
         scale = -scale
 
     num_members = math.comb(n + dim, dim)
-    results = tuple([None] * num_members for i in range(order+1))
-    phi, dphi, ddphi = results + (None,) * (2-order)
 
     outer = lambda x, y: x[:, None, ...] * y[None, ...]
     sym_outer = lambda x, y: outer(x, y) + outer(y, x)
 
     pad_dim = dim + 2
     dX = pad_jacobian(Jinv, pad_dim)
-    phi[0] = sum((ref_pts[i] - ref_pts[i] for i in range(dim)), scale)
-    if dphi is not None:
-        dphi[0] = (phi[0] - phi[0]) * dX[0]
-    if ddphi is not None:
-        ddphi[0] = outer(dphi[0], dX[0])
+
+    phi0 = numpy.array([sum((ref_pts[i] - ref_pts[i] for i in range(dim)), 0.0)])
+    results = [numpy.zeros((num_members,) + (dim,)*k + phi0.shape[1:], dtype=phi0.dtype)
+               for k in range(order+1)]
+
+    phi, dphi, ddphi = results + [None] * (2-order)
+    phi[0] += scale
     if dim == 0 or n == 0:
         return results
     if dim > 3 or dim < 0:
@@ -183,7 +183,7 @@ def C0_basis(dim, n, tabulations):
     # Recover facet bubbles
     for phi in tabulations:
         icur = 0
-        phi[icur] *= -1
+        phi[icur] *= -1.0
         for inext in range(1, dim+1):
             phi[icur] -= phi[inext]
         if dim == 2:
@@ -723,11 +723,15 @@ def compute_partition_of_unity(ref_el, pt, unique=True, tol=1E-12):
     :kwarg tol: the absolute tolerance.
     :returns: a list of (weighted) characteristic functions for each subcell.
     """
-    from sympy import Piecewise
+    import gem
     sd = ref_el.get_spatial_dimension()
     top = ref_el.get_topology()
     # assert singleton point
     pt = pt.reshape((sd,))
+    if isinstance(pt[0], gem.Node):
+        import gem as backend
+    else:
+        import sympy as backend
 
     # The distance to the nearest cell is equal to the distance to the parent cell
     best = ref_el.get_parent().distance_to_point_l1(pt, rescale=True)
@@ -739,7 +743,7 @@ def compute_partition_of_unity(ref_el, pt, unique=True, tol=1E-12):
     for cell in sorted(top[sd]):
         # Bin points based on l1 distance
         pt_near_cell = ref_el.distance_to_point_l1(pt, entity=(sd, cell), rescale=True) < tol
-        masks.append(Piecewise(*otherwise, (1.0, pt_near_cell), (0.0, True)))
+        masks.append(backend.Piecewise(*otherwise, (1.0, pt_near_cell), (0.0, True)))
         if unique:
             otherwise.append((0.0, pt_near_cell))
     # If the point is on a facet, divide the characteristic function by the facet multiplicity

finat/argyris.py

@@ -124,7 +124,7 @@ def _edge_transform(V, vorder, eorder, fiat_cell, coordinate_mapping, avg=False)
             V[s, v1id] = P1 * Bnt
             V[s, v0id] = P0 * Bnt
             if k > 0:
-                V[s, s + eorder] = -1 * Bnt
+                V[s, s + eorder] = -Bnt
 
 
 class Argyris(PhysicallyMappedElement, ScalarFiatElement):
@@ -167,7 +167,7 @@ def basis_transformation(self, coordinate_mapping):
 
                 # vertex points
                 V[s, v1id] = 15/8 * Bnt
-                V[s, v0id] = -1 * V[s, v1id]
+                V[s, v0id] = -V[s, v1id]
 
                 # vertex derivatives
                 for i in range(sd):
@@ -178,7 +178,7 @@ def basis_transformation(self, coordinate_mapping):
                 tau = [Jt[0]*Jt[0], 2*Jt[0]*Jt[1], Jt[1]*Jt[1]]
                 for i in range(len(tau)):
                     V[s, v1id+3+i] = 1/32 * Bnt * tau[i]
-                    V[s, v0id+3+i] = -1 * V[s, v1id+3+i]
+                    V[s, v0id+3+i] = -V[s, v1id+3+i]
 
         # Patch up conditioning
         h = coordinate_mapping.cell_size()

finat/bell.py

@@ -44,7 +44,7 @@ def basis_transformation(self, coordinate_mapping):
 
             # vertex points
             V[s, v1id] = 1/21 * Bnt
-            V[s, v0id] = -1 * V[s, v1id]
+            V[s, v0id] = -V[s, v1id]
 
             # vertex derivatives
             for i in range(sd):
@@ -55,7 +55,7 @@ def basis_transformation(self, coordinate_mapping):
             tau = [Jt[0]*Jt[0], 2*Jt[0]*Jt[1], Jt[1]*Jt[1]]
             for i in range(len(tau)):
                 V[s, v1id+3+i] = 1/252 * Bnt * tau[i]
-                V[s, v0id+3+i] = -1 * V[s, v1id+3+i]
+                V[s, v0id+3+i] = -V[s, v1id+3+i]
 
         # Patch up conditioning
         h = coordinate_mapping.cell_size()

finat/fiat_elements.py

@@ -1,12 +1,10 @@
 import FIAT
 import gem
 import numpy as np
-import sympy as sp
 from gem.utils import cached_property
 
 from finat.finiteelementbase import FiniteElementBase
-from finat.point_set import PointSet
-from finat.sympy2gem import sympy2gem
+from finat.point_set import PointSet, PointSingleton
 
 
 class FiatElement(FiniteElementBase):
@@ -67,62 +65,61 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         :param ps: the point set.
         :param entity: the cell entity on which to tabulate.
         '''
-        space_dimension = self._element.space_dimension()
-        value_size = np.prod(self._element.value_shape(), dtype=int)
-        fiat_result = self._element.tabulate(order, ps.points, entity)
-        result = {}
+        fiat_element = self._element
+        fiat_result = fiat_element.tabulate(order, ps.points, entity)
         # In almost all cases, we have
         # self.space_dimension() == self._element.space_dimension()
         # But for Bell, FIAT reports 21 basis functions,
         # but FInAT only 18 (because there are actually 18
         # basis functions, and the additional 3 are for
         # dealing with transformations between physical
         # and reference space).
-        index_shape = (self._element.space_dimension(),)
+        value_shape = self.value_shape
+        space_dimension = fiat_element.space_dimension()
+        if self.space_dimension() == space_dimension:
+            beta = self.get_indices()
+            index_shape = tuple(index.extent for index in beta)
+        else:
+            index_shape = (space_dimension,)
+            beta = tuple(gem.Index(extent=i) for i in index_shape)
+            assert len(beta) == len(self.get_indices())
+
+        zeta = self.get_value_indices()
+        basis_indices = beta + zeta
+
+        result = {}
         for alpha, fiat_table in fiat_result.items():
             if isinstance(fiat_table, Exception):
-                result[alpha] = gem.Failure(self.index_shape + self.value_shape, fiat_table)
+                result[alpha] = gem.Failure(index_shape + value_shape, fiat_table)
                 continue
 
+            point_indices = ()
+            replace_indices = ()
             derivative = sum(alpha)
-            shp = (space_dimension, value_size, *ps.points.shape[:-1])
-            table_roll = np.moveaxis(fiat_table.reshape(shp), 0, -1)
-
-            exprs = []
-            for table in table_roll:
-                if derivative == self.degree and not self.complex.is_macrocell():
-                    # Make sure numerics satisfies theory
-                    exprs.append(gem.Literal(table[0]))
-                elif derivative > self.degree:
-                    # Make sure numerics satisfies theory
-                    assert np.allclose(table, 0.0)
-                    exprs.append(gem.Literal(np.zeros(self.index_shape)))
+            if derivative == self.degree and self.complex.is_simplex():
+                # Ensure a cellwise constant tabulation
+                if fiat_table.dtype == object:
+                    replace_indices = tuple((i, 0) for i in ps.expression.free_indices)
                 else:
-                    point_indices = ps.indices
-                    point_shape = tuple(index.extent for index in point_indices)
-
-                    exprs.append(gem.partial_indexed(
-                        gem.Literal(table.reshape(point_shape + index_shape)),
-                        point_indices
-                    ))
-            if self.value_shape:
-                # As above, this extent may be different from that
-                # advertised by the finat element.
-                beta = tuple(gem.Index(extent=i) for i in index_shape)
-                assert len(beta) == len(self.get_indices())
-
-                zeta = self.get_value_indices()
-                result[alpha] = gem.ComponentTensor(
-                    gem.Indexed(
-                        gem.ListTensor(np.array(
-                            [gem.Indexed(expr, beta) for expr in exprs]
-                        ).reshape(self.value_shape)),
-                        zeta),
-                    beta + zeta
-                )
+                    fiat_table = fiat_table.reshape(*index_shape, *value_shape, -1)
+                    assert np.allclose(fiat_table, fiat_table[..., 0, None])
+                    fiat_table = fiat_table[..., 0]
+            elif derivative > self.degree:
+                # Ensure a zero tabulation
+                if fiat_table.dtype != object:
+                    assert np.allclose(fiat_table, 0.0)
+                fiat_table = np.zeros(index_shape + value_shape)
             else:
-                expr, = exprs
-                result[alpha] = expr
+                point_indices = ps.indices
+
+            point_shape = tuple(i.extent for i in point_indices)
+            fiat_table = fiat_table.reshape(index_shape + value_shape + point_shape)
+            gem_table = gem.as_gem(fiat_table)
+            expr = gem.Indexed(gem_table, basis_indices + point_indices)
+            expr = gem.ComponentTensor(expr, basis_indices)
+            if replace_indices:
+                expr, = gem.optimise.remove_componenttensors((expr,), subst=replace_indices)
+            result[alpha] = expr
         return result
 
     def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=None):
@@ -147,7 +144,20 @@ def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=Non
         esd = self.cell.construct_subelement(entity_dim).get_spatial_dimension()
         assert isinstance(refcoords, gem.Node) and refcoords.shape == (esd,)
 
-        return point_evaluation(self._element, order, refcoords, (entity_dim, entity_i))
+        # Coordinates on the reference entity (GEM)
+        Xi = tuple(gem.Indexed(refcoords, i) for i in np.ndindex(refcoords.shape))
+        ps = PointSingleton(Xi)
+        result = self.basis_evaluation(order, ps, entity=entity,
+                                       coordinate_mapping=coordinate_mapping)
+
+        # Apply symbolic simplification
+        vals = result.values()
+        vals = map(gem.optimise.ffc_rounding, vals, [1E-13]*len(vals))
+        vals = gem.optimise.constant_fold_zero(vals)
+        vals = map(gem.optimise.aggressive_unroll, vals)
+        vals = gem.optimise.remove_componenttensors(vals)
+        result = dict(zip(result.keys(), vals))
+        return result
 
     @cached_property
     def _dual_basis(self):
@@ -260,49 +270,6 @@ def mapping(self):
             return result
 
 
-def point_evaluation(fiat_element, order, refcoords, entity):
-    # Coordinates on the reference entity (SymPy)
-    esd, = refcoords.shape
-    Xi = sp.symbols('X Y Z')[:esd]
-
-    space_dimension = fiat_element.space_dimension()
-    value_size = np.prod(fiat_element.value_shape(), dtype=int)
-    fiat_result = fiat_element.tabulate(order, [Xi], entity)
-    result = {}
-    for alpha, fiat_table in fiat_result.items():
-        if isinstance(fiat_table, Exception):
-            result[alpha] = gem.Failure((space_dimension,) + fiat_element.value_shape(), fiat_table)
-            continue
-
-        # Convert SymPy expression to GEM
-        mapper = gem.node.Memoizer(sympy2gem)
-        mapper.bindings = {s: gem.Indexed(refcoords, (i,))
-                           for i, s in enumerate(Xi)}
-        gem_table = np.vectorize(mapper)(fiat_table)
-
-        table_roll = gem_table.reshape(space_dimension, value_size).transpose()
-
-        exprs = []
-        for table in table_roll:
-            exprs.append(gem.ListTensor(table.reshape(space_dimension)))
-        if fiat_element.value_shape():
-            beta = (gem.Index(extent=space_dimension),)
-            zeta = tuple(gem.Index(extent=d)
-                         for d in fiat_element.value_shape())
-            result[alpha] = gem.ComponentTensor(
-                gem.Indexed(
-                    gem.ListTensor(np.array(
-                        [gem.Indexed(expr, beta) for expr in exprs]
-                    ).reshape(fiat_element.value_shape())),
-                    zeta),
-                beta + zeta
-            )
-        else:
-            expr, = exprs
-            result[alpha] = expr
-    return result
-
-
 class Regge(FiatElement):  # symmetric matrix valued
     def __init__(self, cell, degree, **kwargs):
         super().__init__(FIAT.Regge(cell, degree, **kwargs))

finat/hct.py

@@ -71,7 +71,7 @@ def basis_transformation(self, coordinate_mapping):
 
             # vertex points
             V[s, v0id] = 1/5 * Bnt
-            V[s, v1id] = -1 * V[s, v0id]
+            V[s, v1id] = -V[s, v0id]
 
             # vertex derivatives
             for i in range(sd):

finat/physically_mapped.py

@@ -89,10 +89,6 @@ def basis_evaluation(self, order, ps, entity=None, coordinate_mapping=None):
         result = super().basis_evaluation(order, ps, entity=entity)
         return self.map_tabulation(result, coordinate_mapping)
 
-    def point_evaluation(self, order, refcoords, entity=None, coordinate_mapping=None):
-        result = super().point_evaluation(order, refcoords, entity=entity)
-        return self.map_tabulation(result, coordinate_mapping)
-
     def dual_transformation(self, Q, coordinate_mapping=None):
         M = self.basis_transformation(coordinate_mapping)
 

finat/point_set.py

@@ -76,7 +76,7 @@ def points(self):
 
     @cached_property
     def expression(self):
-        return gem.Literal(self.point)
+        return gem.as_gem(self.point)
 
 
 class UnknownPointsArray():

finat/sympy2gem.py

@@ -1,6 +1,5 @@
 from functools import singledispatch, reduce
 
-import numpy
 import sympy
 try:
     import symengine
@@ -130,18 +129,7 @@ def sympy2gem_le(node, self):
 @sympy2gem.register(sympy.Piecewise)
 @sympy2gem.register(symengine.Piecewise)
 def sympy2gem_conditional(node, self):
-    expr = None
-    pieces = []
-    for v, c in node.args:
-        if isinstance(c, (bool, numpy.bool, sympy.logic.boolalg.BooleanTrue)) and c:
-            expr = self(v)
-            break
-        pieces.append((v, c))
-    if expr is None:
-        expr = gem.Literal(float("nan"))
-    for v, c in reversed(pieces):
-        expr = gem.Conditional(self(c), self(v), expr)
-    return expr
+    return gem.Piecewise(*[(self(v), self(c)) for v, c in node.args])
 
 
 @sympy2gem.register(sympy.ITE)