propagate dtype through finat element construction

FIAT/expansions.py

@@ -459,7 +459,7 @@ def _tabulate(self, n, pts, order=0):
 
         if pts.dtype == object:
             # If binning is undefined, scale by the characteristic function of each subcell
-            tol = calibrate_tolerance(1E-12, numpy.array(self.ref_el.vertices).dtype)
+            tol = calibrate_tolerance(1E-12, self.ref_el.target_dtype)
             Xi = compute_partition_of_unity(self.ref_el, pts, unique=unique, tol=tol)
             for cell, phi in phis.items():
                 for alpha in phi:

FIAT/macro.py

@@ -151,6 +151,7 @@ def __init__(self, parent, vertices, topology):
         self._interior_facets = interior_facets
 
         super().__init__(parent.shape, vertices, topology)
+        self.target_dtype = parent.target_dtype
 
     def get_child_to_parent(self):
         """Maps split complex facet tuple to its parent entity tuple."""

FIAT/polynomial_set.py

@@ -16,6 +16,7 @@
 # an entire set of polynomials)
 
 import numpy
+
 from itertools import chain
 from FIAT import expansions
 

FIAT/reference_element.py

@@ -143,6 +143,7 @@ def __init__(self, shape, vertices, topology):
         comprising the facet."""
         self.shape = shape
         self.vertices = vertices
+        self.target_dtype = numpy.asarray(vertices).dtype
         self.topology = topology
 
         # Given the topology, work out for each entity in the cell,
@@ -409,7 +410,8 @@ def compute_normal(self, facet_i, cell=None):
         if cell is None:
             cell = next(k for k, facets in enumerate(self.connectivity[(sd, sd-1)])
                         if facet_i in facets)
-        verts = numpy.asarray(self.get_vertices_of_subcomplex(t[sd][cell]))
+        # Always compute in double precision for a robust SVD below.
+        verts = numpy.asarray(self.get_vertices_of_subcomplex(t[sd][cell]), dtype=float)
         # Interval case
         if self.get_shape() == LINE:
             v_i = t[1][cell].index(t[0][facet_i][0])
@@ -429,7 +431,7 @@ def compute_normal(self, facet_i, cell=None):
             self.get_vertices_of_subcomplex(t[sd-1][facet_i])
 
         # now I find everything normal to the facet.
-        vcf = numpy.asarray(vert_coords_of_facet)
+        vcf = numpy.asarray(vert_coords_of_facet, dtype=float)
         facet_span = vcf[1:, :] - vcf[:1, :]
         (_, sf, vft) = numpy.linalg.svd(facet_span)
 
@@ -1658,9 +1660,7 @@ def ufc_hypercube(spatial_dim, dtype=None):
     """Factory function that maps spatial dimension to an instance of
     the UFC reference hypercube of that dimension.
 
-    :arg dtype: optional numpy dtype to cast the vertex coordinates to.
-        Defaults to the plain Python `float` used by the hardcoded
-        vertex coordinates.
+    :arg dtype: optional working dtype for tabulation.
     """
     if spatial_dim == 0:
         cell = Point()
@@ -1673,7 +1673,7 @@ def ufc_hypercube(spatial_dim, dtype=None):
     else:
         raise RuntimeError(f"Can't create UFC hypercube of dimension {spatial_dim}.")
     if dtype is not None:
-        cell.vertices = cast_vertices(cell.vertices, dtype)
+        cell.target_dtype = numpy.dtype(dtype)
     return cell
 
 
@@ -1696,9 +1696,7 @@ def ufc_simplex(spatial_dim, dtype=None):
     """Factory function that maps spatial dimension to an instance of
     the UFC reference simplex of that dimension.
 
-    :arg dtype: optional numpy dtype to cast the vertex coordinates to.
-        Defaults to the plain Python `float` used by the hardcoded
-        vertex coordinates.
+    :arg dtype: optional working dtype for tabulation.
     """
     if spatial_dim == 0:
         cell = Point()
@@ -1711,20 +1709,22 @@ def ufc_simplex(spatial_dim, dtype=None):
     else:
         raise RuntimeError(f"Can't create UFC simplex of dimension {spatial_dim}.")
     if dtype is not None:
-        cell.vertices = cast_vertices(cell.vertices, dtype)
+        cell.target_dtype = numpy.dtype(dtype)
     return cell
 
 
 def symmetric_simplex(spatial_dim, dtype=None):
     A = numpy.array([[2, 1, 1],
                      [0, numpy.sqrt(3), numpy.sqrt(3)/3],
-                     [0, 0, numpy.sqrt(6)*(2/3)]], dtype=dtype)
+                     [0, 0, numpy.sqrt(6)*(2/3)]], dtype=float)
     A = A[:spatial_dim, :][:, :spatial_dim]
     b = A.sum(axis=1) * (-1 / (1 + spatial_dim))
     Ref1 = ufc_simplex(spatial_dim, dtype=dtype)
     v = numpy.dot(Ref1.get_vertices(), A.T) + b[None, :]
     vertices = tuple(map(tuple, v))
-    return SymmetricSimplex(Ref1.get_shape(), vertices, Ref1.get_topology())
+    cell = SymmetricSimplex(Ref1.get_shape(), vertices, Ref1.get_topology())
+    cell.target_dtype = numpy.dtype(float if dtype is None else dtype)
+    return cell
 
 
 def ufc_cell(cell, dtype=None):
@@ -1738,7 +1738,9 @@ def ufc_cell(cell, dtype=None):
 
     if " * " in celltype:
         # Tensor product cell
-        return TensorProductCell(*(ufc_cell(c, dtype=dtype) for c in celltype.split(" * ")))
+        ref_el = TensorProductCell(*(ufc_cell(c, dtype=dtype) for c in celltype.split(" * ")))
+        ref_el.target_dtype = numpy.dtype(float if dtype is None else dtype)
+        return ref_el
     elif celltype == "quadrilateral":
         return ufc_hypercube(2, dtype=dtype)
     elif celltype == "hexahedron":

finat/element_factory.py

@@ -111,13 +111,15 @@
 
 
 @cache
-def as_fiat_cell(cell):
+def as_fiat_cell(cell, dtype=None):
     """Convert a ufl cell to a FIAT cell.
 
-    :arg cell: the :class:`ufl.Cell` to convert."""
+    :arg cell: the :class:`ufl.Cell` to convert.
+    :arg dtype: the dtype to build the reference cell's coordinates with.
+        Defaults to FIAT's own default (float64) if not given."""
     if not isinstance(cell, ufl.AbstractCell):
         raise ValueError("Expecting a UFL Cell")
-    return ufc_cell(cell)
+    return ufc_cell(cell, dtype=dtype)
 
 
 @singledispatch
@@ -153,15 +155,15 @@ def convert(element, **kwargs):
 # Base finite elements first
 @convert.register(finat.ufl.FiniteElement)
 def convert_finiteelement(element, **kwargs):
-    cell = as_fiat_cell(element.cell)
+    cell = as_fiat_cell(element.cell, dtype=kwargs.get("dtype"))
     if element.family() in {"Quadrature", "Boundary Quadrature"}:
         degree = element.degree()
         scheme = element.quadrature_scheme() or "default"
         if degree is None or scheme is None:
             raise ValueError("Quadrature scheme and degree must be specified!")
 
         codim = 1 if element.family() == "Boundary Quadrature" else 0
-        return finat.make_quadrature_element(cell, degree, scheme, codim), set()
+        return finat.make_quadrature_element(cell, degree, scheme, codim), {"dtype"}
 
     make_finat_element = supported_elements[element.family()]
 
@@ -186,7 +188,7 @@ def convert_finiteelement(element, **kwargs):
         finat_elem, deps = _create_element(element, **kwargs)
         return finat.FlattenedDimensions(finat_elem), deps
 
-    deps = set()
+    deps = {"dtype"}
     finat_kwargs = {}
     kind = element.variant()
     if kind is None:
@@ -205,7 +207,7 @@ def convert_finiteelement(element, **kwargs):
             finat_kwargs["variant"] = kind
             finat_kwargs["shift_axes"] = kwargs["shift_axes"]
             finat_kwargs["restriction"] = kwargs["restriction"]
-            deps = {"shift_axes", "restriction"}
+            deps |= {"shift_axes", "restriction"}
         else:
             # Let FIAT handle the general case
             make_finat_element = finat.Lagrange
@@ -227,7 +229,7 @@ def convert_finiteelement(element, **kwargs):
             finat_kwargs["shift_axes"] = kwargs["shift_axes"]
             finat_kwargs["restriction"] = kwargs["restriction"]
             finat_kwargs["continuous"] = False
-            deps = {"shift_axes", "restriction"}
+            deps |= {"shift_axes", "restriction"}
         else:
             # Let FIAT handle the general case
             make_finat_element = finat.DiscontinuousLagrange
@@ -330,18 +332,21 @@ def convert_restrictedelement(element, **kwargs):
 _cache = weakref.WeakKeyDictionary()
 
 
-def create_element(ufl_element, shape_innermost=True, shift_axes=0, restriction=None):
+def create_element(ufl_element, shape_innermost=True, shift_axes=0, restriction=None, dtype=None):
     """Create a FInAT element (suitable for tabulating with) given a UFL element.
 
     :arg ufl_element: The UFL element to create a FInAT element from.
     :arg shape_innermost: Vector/tensor indices come after basis function indices
     :arg restriction: cell restriction in interior facet integrals
                       (only for runtime tabulated elements)
+    :arg dtype: the dtype to build the element's reference cell with.
+               Defaults to FIAT's own default (float64) if not given.
     """
     finat_element, deps = _create_element(ufl_element,
                                           shape_innermost=shape_innermost,
                                           shift_axes=shift_axes,
-                                          restriction=restriction)
+                                          restriction=restriction,
+                                          dtype=dtype)
     return finat_element
 
 

test/FIAT/unit/test_hct.py

@@ -7,9 +7,9 @@
 from FIAT.macro import CkPolynomialSet
 
 
-@pytest.fixture
-def cell():
-    K = ufc_simplex(2)
+@pytest.fixture(params=(numpy.float64, numpy.float32), ids=("float64", "float32"))
+def cell(request):
+    K = ufc_simplex(2, dtype=request.param)
     K.vertices = ((0.0, 0.1), (1.17, -0.09), (0.15, 1.84))
     return K
 
@@ -76,7 +76,9 @@ def test_full_polynomials(cell, reduced):
     assert span_greater_equal(tab, C1_tab)
 
 
+@pytest.mark.parametrize("dtype", (numpy.float64, numpy.float32),
+                         ids=("float64", "float32"))
 @pytest.mark.parametrize("degree, space_dimension", [(13, 127), (14, 144)])
-def test_hct_high_order_degree(degree: int, space_dimension: int) -> None:
-    fe = HCT(ufc_simplex(2), degree)
+def test_hct_high_order_degree(dtype, degree: int, space_dimension: int) -> None:
+    fe = HCT(ufc_simplex(2, dtype=dtype), degree)
     assert fe.space_dimension() == space_dimension

test/FIAT/unit/test_macro.py

@@ -375,11 +375,15 @@ def test_macro_expansion(cell, split, variant, degree):
 @pytest.mark.parametrize("order", (0, 1))
 @pytest.mark.parametrize("variant", (None, "bubble"))
 @pytest.mark.parametrize("degree", (1, 4))
-def test_Ck_basis(cell, order, degree, variant):
+@pytest.mark.parametrize("split", (AlfeldSplit, IsoSplit))
+@pytest.mark.parametrize("dtype", (numpy.float64, numpy.float32),
+                         ids=("float64", "float32"))
+def test_Ck_basis(cell, order, degree, variant, split, dtype):
     # Test that we can correctly tabulate on points on facets.
     # This breaks if we were binning points into more than one cell without a partition of unity.
     # It suffices to tabulate on the vertices of the simplicial complex.
-    A = AlfeldSplit(cell)
+    cell = ufc_simplex(cell.get_spatial_dimension(), dtype=dtype)
+    A = split(cell)
     Ck = CkPolynomialSet(A, degree, order=order, variant=variant)
     U = Ck.get_expansion_set()
     cell_node_map = U.get_cell_node_map(degree)
@@ -397,11 +401,12 @@ def test_Ck_basis(cell, order, degree, variant):
         assert numpy.allclose(local_phis, phis[:, ipts])
 
 
-def test_C2_double_alfeld():
+@pytest.mark.parametrize("dtype", (numpy.float64, numpy.float32))
+def test_C2_double_alfeld(dtype):
     from FIAT.c2_elements import AlfeldC2Space
     # Construct the quintic C2 spline on the double Alfeld split
     # See Section 7.5 of Lai & Schumacher
-    K = ufc_simplex(2)
+    K = ufc_simplex(2, dtype=dtype)
     degree = 5
     P = AlfeldC2Space(K, degree)
     assert P.get_num_members() == 27

test/FIAT/unit/test_powell_sabin.py

@@ -6,15 +6,19 @@
 from FIAT.reference_element import make_lattice, ufc_simplex
 
 
-@pytest.fixture
-def cell():
-    return ufc_simplex(2)
+@pytest.fixture(params=(numpy.float64, numpy.float32),
+                ids=("float64", "float32"))
+def cell(request):
+    return ufc_simplex(2, dtype=request.param)
 
 
 @pytest.mark.parametrize("el", (PS6, PS12))
 def test_powell_sabin_constant(cell, el):
     # Test that bfs associated with point evaluation sum up to 1
     fe = el(cell)
+    assert numpy.asarray(fe.get_reference_complex().vertices).dtype == numpy.float64
+    assert fe.get_reference_complex().target_dtype == cell.target_dtype
+    assert fe.get_coeffs().dtype == numpy.float64
 
     pts = make_lattice(cell.get_vertices(), 3)
     tab = fe.tabulate(2, pts)

test/FIAT/unit/test_precision.py

@@ -46,7 +46,9 @@ def test_dtype_propagates_into_symbolic_tabulation():
     fe64 = macro_element(dtype=numpy.float64)
     fe32 = macro_element(dtype=numpy.float32)
     assert numpy.array(fe64.ref_complex.vertices).dtype == numpy.float64
-    assert numpy.array(fe32.ref_complex.vertices).dtype == numpy.float32
+    assert numpy.array(fe32.ref_complex.vertices).dtype == numpy.float64
+    assert fe64.ref_complex.target_dtype == numpy.float64
+    assert fe32.ref_complex.target_dtype == numpy.float32
     x0 = gem.Variable("x0", ())
 
     tab64 = fe64.tabulate(0, (x0,))[(0,)][0]

test/FIAT/unit/test_stokes_complex.py

@@ -143,8 +143,10 @@ def check_stokes_complex(spaces, degree):
 
 @pytest.mark.parametrize("reduced", (False, True), ids=("full", "reduced"))
 @pytest.mark.parametrize("sobolev", ("H1", "H1div"))
-@pytest.mark.parametrize("cell", (T,))
-def test_hct_stokes_complex(cell, sobolev, reduced):
+@pytest.mark.parametrize("dtype", (numpy.float64, numpy.float32),
+                         ids=("float64", "float32"))
+def test_hct_stokes_complex(dtype, sobolev, reduced):
+    cell = ufc_simplex(2, dtype=dtype)
     if sobolev == "H1":
         if reduced:
             spaces = [rHCT(cell), rAQ(cell), DG(cell, 0)]
@@ -246,11 +248,18 @@ def test_gn_trace(sd):
 
 
 @pytest.mark.parametrize("cell", (T, S))
-@pytest.mark.parametrize("family", ("AQ", "CH", "GN", "GN2"))
-def test_minimal_stokes_space(cell, family):
+@pytest.mark.parametrize("family,dtype", (
+    ("AQ", numpy.float64),
+    ("CH", numpy.float64),
+    ("CH", numpy.float32),
+    ("GN", numpy.float64),
+    ("GN2", numpy.float64),
+))
+def test_minimal_stokes_space(cell, family, dtype):
     # Test that the C0 Stokes space is spanned by a C0 basis
     # Also test that its divergence is constant
     sd = cell.get_spatial_dimension()
+    cell = ufc_simplex(sd, dtype=dtype)
     if family == "GN":
         degree = 1
         space = GuzmanNeilanSpace

test/finat/test_create_finat_element.py

@@ -1,3 +1,4 @@
+import numpy
 import pytest
 
 import ufl
@@ -166,6 +167,30 @@ def test_cache_hit_vector(ufl_vector_element):
     assert A is B
 
 
+def test_dtype_reaches_reference_cell(ufl_element):
+    """Check that target dtype is independent of construction dtype."""
+    default = create_element(ufl_element)
+    single = create_element(ufl_element, dtype=numpy.float32)
+    double = create_element(ufl_element, dtype=numpy.float64)
+
+    assert numpy.array(default.cell.vertices).dtype == numpy.float64
+    assert numpy.array(single.cell.vertices).dtype == numpy.float64
+    assert numpy.array(double.cell.vertices).dtype == numpy.float64
+    assert default.cell.target_dtype == numpy.float64
+    assert single.cell.target_dtype == numpy.float32
+    assert double.cell.target_dtype == numpy.float64
+
+
+def test_dtype_cache_distinguishes(ufl_element):
+    """Different dtypes for the same UFL element must not share a cache entry."""
+    single_a = create_element(ufl_element, dtype=numpy.float32)
+    single_b = create_element(ufl_element, dtype=numpy.float32)
+    double = create_element(ufl_element, dtype=numpy.float64)
+
+    assert single_a is single_b
+    assert single_a is not double
+
+
 if __name__ == "__main__":
     import os
     import sys