Sum factorisation on simplices

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
   # each Firedrake release to a specific UFL minor version (e.g. 2025.3.x)
   "fenics-ufl @ git+https://github.com/FEniCS/ufl.git@main",
   # TODO RELEASE
-  "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@main",
+  # DROP BEFORE MERGE: pinned to the paired FIAT branch for CI; revert to @main
+  "firedrake-fiat @ git+https://github.com/firedrakeproject/fiat.git@pbrubeck/simplex-sum-factor",
   "h5py>3.12.1",
   "firedrake-rtree",
   "immutabledict",

tests/firedrake/regression/test_helmholtz_bernstein.py

@@ -27,7 +27,7 @@ def mesh(request):
 def test_bernstein(mesh, degree):
     # Solve with Bernstein polynomials
     B = FunctionSpace(mesh, "Bernstein", degree)
-    xb = helmholtz(B)
+    xb = helmholtz(B, dx(scheme="collapsed"))
 
     # Solve with Lagrange polynomials
     L = FunctionSpace(mesh, "Lagrange", degree)
@@ -39,15 +39,15 @@ def test_bernstein(mesh, degree):
     assert np.allclose(xl.dat.data, xp.dat.data)
 
 
-def helmholtz(V):
+def helmholtz(V, measure=dx):
     # Define variational problem
     u = TrialFunction(V)
     v = TestFunction(V)
     f = Function(V)
     x = SpatialCoordinate(V.mesh())
     f.project(np.prod([cos(2*pi*xi) for xi in x]))
-    a = (inner(grad(u), grad(v)) + inner(u, v)) * dx
-    L = inner(f, v) * dx
+    a = (inner(grad(u), grad(v)) + inner(u, v)) * measure
+    L = inner(f, v) * measure
 
     # Compute solution
     x = Function(V)

tests/firedrake/regression/test_quadrature.py

@@ -52,3 +52,50 @@ def test_quadrature_element(mesh, family, mat_type, diagonal):
         a = inner(u, v) * dx
 
     assemble(a, mat_type=mat_type, diagonal=diagonal)
+
+
+@pytest.mark.parametrize("family", ["DG", "CG"])
+@pytest.mark.parametrize("cell", ["triangle", "tetrahedron"])
+@pytest.mark.parametrize("degree", [1, 3])
+def test_collapsed_quadrature_sum_factorisation(cell, degree, family):
+    """``dx(scheme="collapsed")`` on a simplicial "DG"/variant="integral"
+    (i.e. `finat.spectral.Legendre`) or "CG"/variant="integral" (i.e.
+    `finat.spectral.IntegratedLegendre`, exercising the
+    `FIAT.expansions.C0_basis` recombination) space must produce the same
+    assembled residual and matrix as ``dx(scheme="canonical")``, even
+    though it takes the sum-factorized (Duffy/lattice) tabulation path
+    in ``tsfc.fem`` rather than the standard dense one.
+
+    ``"canonical"`` is the same collapsed Gauss-Jacobi rule as
+    ``"collapsed"`` (same points, same weights,
+    `finat.quadrature.collapsed_gauss_jacobi_quadrature`), but as a plain
+    `finat.point_set.PointSet` rather than a `CollapsedTensorProductPointSet`
+    -- so `finat.duffy.DuffyElement._duffy_applies` is `False` and it always
+    takes the dense FIAT tabulation path. Comparing against it (rather than
+    the default scheme, a different quadrature rule entirely for these
+    degrees) isolates any discrepancy to the sum-factorized tabulation
+    itself, not to a difference in quadrature choice.
+    """
+    mesh = {"triangle": UnitSquareMesh(2, 2),
+            "tetrahedron": UnitCubeMesh(1, 1, 1)}[cell]
+    V = FunctionSpace(mesh, family, degree, variant="integral")
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    w = Function(V)
+    w.dat.data[:] = np.random.default_rng(0).random(w.dat.data.shape)
+
+    # translate_coefficient path (forward transform): residual with a
+    # derivative, mixing both the coefficient and argument sum-factorized
+    # tabulations.
+    L = inner(grad(w), grad(v)) * dx(scheme="canonical")
+    L_collapsed = inner(grad(w), grad(v)) * dx(scheme="collapsed")
+    b = assemble(L)
+    b_collapsed = assemble(L_collapsed)
+    assert np.allclose(b.dat.data, b_collapsed.dat.data, rtol=1e-10, atol=1e-10)
+
+    # translate_argument path (backward transform): mass matrix.
+    a = inner(u, v) * dx(scheme="canonical")
+    a_collapsed = inner(u, v) * dx(scheme="collapsed")
+    M = assemble(a).M.values
+    M_collapsed = assemble(a_collapsed).M.values
+    assert np.allclose(M, M_collapsed, rtol=1e-10, atol=1e-10)

tests/tsfc/test_codegen.py

@@ -1,6 +1,7 @@
+import numpy
 import pytest
 
-from gem import impero_utils
+from gem import gem, impero_utils
 from gem.gem import Index, Indexed, IndexSum, Product, Variable
 
 
@@ -24,6 +25,131 @@ def gencode(expr):
     assert len(gencode(e1).children) == len(gencode(e2).children)
 
 
+def test_jagged_index_codegen(monkeypatch):
+    import islpy as isl
+    import loopy as lp
+    import tsfc.loopy
+
+    # Execute the generated code so we check the numbers, not just the loop bounds
+    monkeypatch.setattr(tsfc.loopy, "target", lp.ExecutableCTarget())
+
+    n = 4
+    extent = n + 1
+    npts = 3
+    ndof = (n + 1) * (n + 2) // 2
+
+    rng = numpy.random.default_rng(7)
+    # Table zero-padded outside the simplex lattice p + q > n, and a
+    # clamped Morton index table for the coefficient gather
+    B = rng.random((extent, extent, npts))
+    morton = numpy.zeros((extent, extent), dtype=gem.uint_type)
+    for p_, q_ in numpy.ndindex(morton.shape):
+        if p_ + q_ > n:
+            B[p_, q_] = 0.0
+        else:
+            morton[p_, q_] = (p_ + q_) * (p_ + q_ + 1) // 2 + q_
+    c = rng.random(ndof)
+
+    i = Index(name="i", extent=npts)
+    p = Index(name="p", extent=extent)
+    q = gem.JaggedIndex(name="q", extent=extent, parents=(p,))
+
+    dof = gem.VariableIndex(Indexed(gem.Literal(morton, dtype=gem.uint_type), (p, q)))
+    integrand = Product(Indexed(Variable("c", (ndof,)), (dof,)),
+                        Indexed(gem.Literal(B), (p, q, i)))
+    expr = IndexSum(integrand, (p, q))
+
+    u = Variable("u", (npts,))
+    impero_c = impero_utils.compile_gem([(Indexed(u, (i,)), expr)], (i, p, q))
+    args = [lp.GlobalArg("u", dtype=numpy.float64, shape=(npts,)),
+            lp.GlobalArg("c", dtype=numpy.float64, shape=(ndof,))]
+    knl, _ = tsfc.loopy.generate(impero_c, args, numpy.float64)
+
+    # The jagged loop must have a domain parametrized by its parent iname
+    assert any(dom.get_var_names(isl.dim_type.param)
+               for dom in knl.default_entrypoint.domains)
+
+    u_out = numpy.zeros(npts)
+    knl(c=c, u=u_out)
+    u_ref = numpy.tensordot(c[morton], B, axes=((0, 1), (0, 1)))
+    assert numpy.allclose(u_out, u_ref, rtol=1e-14)
+
+
+@pytest.mark.parametrize("element_name", ["Legendre", "IntegratedLegendre"])
+@pytest.mark.parametrize("cellname,degree", [("triangle", 3), ("tetrahedron", 2)])
+def test_duffy_scatter_and_contract(cellname, degree, element_name):
+    """Route B of the simplex sum-factorization milestone 2 design:
+    `finat.duffy.DuffyElement.basis_evaluation` must reproduce the standard
+    dense FIAT tabulation, via the dof numbering FIAT already uses, from
+    `duffy_evaluation`'s lattice-indexed, sum-factorized tabulation, and a
+    generic contraction of that same flat-dof-indexed result against a
+    coefficient vector (as `tsfc.fem.translate_coefficient` performs
+    uniformly for every FInAT element, `DuffyElement` included) must
+    reproduce the dense tensordot. `Legendre` (continuity=None) reads
+    exactly one lattice point per dof (lattice-lexicographic order);
+    `IntegratedLegendre` (continuity="C0") additionally exercises the
+    `FIAT.expansions.C0_basis` recombination, where each dof combines a
+    handful of lattice points.
+
+    Verified via `gem.interpreter.evaluate` rather than a compiled loopy
+    kernel: the GEM expressions built here (in particular
+    `duffy_evaluation`'s `gem.VariableIndex`-based scatter index arithmetic)
+    schedule fine once merged into a real PyOP2 wrapper kernel (as confirmed
+    via `firedrake.assemble` on real forms), but are not guaranteed
+    schedulable by loopy in isolation -- scheduling an isolated,
+    unwrapped kernel is not a configuration real Firedrake usage ever
+    exercises. The GEM interpreter checks the same numerical correctness
+    without depending on loopy scheduling at all.
+    """
+    from FIAT.reference_element import UFCTetrahedron, UFCTriangle
+    from finat.quadrature import make_quadrature
+    from finat.spectral import Legendre, IntegratedLegendre
+    from gem.gem import Index, Indexed, IndexSum, Product, Variable
+    from gem.interpreter import evaluate
+    from gem.optimise import remove_componenttensors
+
+    cell = {"triangle": UFCTriangle, "tetrahedron": UFCTetrahedron}[cellname]()
+    element_cls = {"Legendre": Legendre, "IntegratedLegendre": IntegratedLegendre}[element_name]
+    element = element_cls(cell, degree)
+    ndof = element.space_dimension()
+
+    quad_rule = make_quadrature(cell, 2 * degree, scheme="collapsed")
+    point_set = quad_rule.point_set
+    point_indices = point_set.indices
+    point_shape = tuple(index.extent for index in point_indices)
+
+    entity = (cell.get_dimension(), 0)
+    dense_dict = element._element.tabulate(1, point_set.points)
+
+    rng = numpy.random.default_rng(1)
+    coefficients = rng.random(ndof)
+
+    # translate_argument path: dispatched transparently through basis_evaluation
+    scattered_dict = element.basis_evaluation(1, point_set, entity)
+    for alpha, dense in dense_dict.items():
+        dense = dense.reshape((ndof,) + point_shape)
+
+        r = Index(extent=ndof)
+        table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))])
+        u_out, = evaluate([table])
+        assert u_out.fids == (r,) + point_indices
+        assert numpy.allclose(u_out.arr, dense, rtol=1e-12, atol=1e-12)
+
+    # translate_coefficient path: generic contraction of the same
+    # flat-dof-indexed basis_evaluation result against a coefficient vector
+    # (no element-specific dispatch, see tsfc.fem.translate_coefficient)
+    c = Variable("c", (ndof,))
+    r = Index(extent=ndof)
+    for alpha, dense in dense_dict.items():
+        dense = dense.reshape((ndof,) + point_shape)
+        table, = remove_componenttensors([Indexed(scattered_dict[alpha], (r,))])
+        value = IndexSum(Product(Indexed(c, (r,)), table), (r,))
+        v_out, = evaluate([value], bindings={c: coefficients})
+        assert v_out.fids == point_indices
+        v_ref = numpy.tensordot(coefficients, dense, axes=(0, 0))
+        assert numpy.allclose(v_out.arr, v_ref, rtol=1e-12, atol=1e-12)
+
+
 if __name__ == "__main__":
     import os
     import sys

tests/tsfc/test_pickle_gem.py

@@ -17,6 +17,20 @@ def test_pickle_gem(protocol):
     assert repr(expr) == repr(unpickled)
 
 
+@pytest.mark.parametrize('protocol', range(3))
+def test_pickle_jagged_index(protocol):
+    p = gem.Index(name='p', extent=4)
+    q = gem.JaggedIndex(name='q', extent=4, parents=(p,))
+    expr = gem.IndexSum(gem.Indexed(gem.Variable('A', (4, 4)), (p, q)), (p, q))
+
+    unpickled = pickle.loads(pickle.dumps(expr, protocol))
+    assert repr(expr) == repr(unpickled)
+    up, uq = unpickled.multiindex
+    assert isinstance(uq, gem.JaggedIndex)
+    assert uq.extent == 4
+    assert uq.parents == (up,)
+
+
 @pytest.mark.parametrize('protocol', range(3))
 def test_listtensor(protocol):
     expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()])

tests/tsfc/test_sum_factorisation.py

@@ -3,7 +3,7 @@
 
 from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction,
                  TensorProductCell, dx, action, interval, triangle,
-                 quadrilateral, curl, dot, div, grad)
+                 tetrahedron, quadrilateral, curl, dot, div, grad, inner)
 from finat.ufl import (FiniteElement, VectorElement, EnrichedElement,
                        TensorProductElement, HCurlElement, HDivElement)
 
@@ -168,7 +168,80 @@ def test_vector_laplace_action(cell, order):
     assert (rates < order).all()
 
 
-if __name__ == "__main__":
-    import os
-    import sys
-    pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:])
+def simplex_mass(cell, family, degree):
+    m = Mesh(VectorElement('CG', cell, 1))
+    variant = None if family == "Bernstein" else "integral"
+    V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant))
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    return inner(u, v) * dx(scheme='collapsed')
+
+
+def simplex_laplacian(cell, family, degree):
+    m = Mesh(VectorElement('CG', cell, 1))
+    variant = None if family == "Bernstein" else "integral"
+    V = FunctionSpace(m, FiniteElement(family, cell, degree, variant=variant))
+    u = TrialFunction(V)
+    v = TestFunction(V)
+    return inner(grad(u), grad(v)) * dx(scheme='collapsed')
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4)])
+def test_simplex_mass_action(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(action(simplex_mass(cell, family, degree)))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 3), (tetrahedron, 4.4)])
+def test_simplex_laplacian_action(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(action(simplex_laplacian(cell, family, degree)))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+def test_simplex_laplacian_action_compact_codegen(family):
+    form = action(simplex_laplacian(triangle, family, 3))
+    kernel, = compile_form(form, parameters=dict(mode='spectral'))
+    temporaries = kernel.ast.default_entrypoint.temporary_variables
+    assert len(temporaries) < 70
+
+
+def test_bernstein_laplacian_action_compact_literals():
+    degree = 5
+    form = action(simplex_laplacian(tetrahedron, "Bernstein", degree))
+    kernel, = compile_form(form, parameters=dict(mode='spectral'))
+    temporaries = kernel.ast.default_entrypoint.temporary_variables
+    literals = [numpy.asarray(temporary.initializer)
+                for temporary in temporaries.values()
+                if temporary.initializer is not None]
+    lattice_size = (degree + 1) ** 3
+    assert max(literal.size for literal in literals) <= lattice_size
+    assert sum(literal.size for literal in literals) < 10 * lattice_size
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)])
+def test_simplex_mass_bilinear(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(simplex_mass(cell, family, degree))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()
+
+
+@pytest.mark.parametrize('family', ["DG", "CG", "Bernstein"])
+@pytest.mark.parametrize(('cell', 'order'), [(triangle, 5), (tetrahedron, 7)])
+def test_simplex_laplacian_bilinear(cell, family, order):
+    degrees = list(range(3, 9)) if cell is triangle else list(range(3, 8))
+    flops = [count_flops(simplex_laplacian(cell, family, degree))
+             for degree in degrees]
+    rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees))
+    assert (rates < order).all()

tsfc/kernel_interface/common.py

@@ -14,9 +14,10 @@
 import numpy
 from FIAT.reference_element import TensorProductCell
 from finat.cell_tools import max_complex
+from finat.duffy import DuffyElement
 from finat.quadrature import AbstractQuadratureRule
 from gem.node import traversal
-from gem.optimise import constant_fold_zero
+from gem.optimise import constant_fold_zero, unflatten_returns
 from gem.optimise import remove_componenttensors as prune
 from numpy import asarray
 from tsfc import fem
@@ -210,6 +211,8 @@ def compile_gem(self, ctx):
             assignments.extend(mode.flatten(var_reps.items(), ctx['index_cache']))
 
         if assignments:
+            # Rewrite flat FlattenedTensor scatters as jagged lattice loops
+            assignments = unflatten_returns(assignments)
             return_variables, expressions = zip(*assignments)
         else:
             return_variables = []
@@ -344,6 +347,13 @@ def set_quad_rule(params, cell, integral_type, functions):
         scheme = quad_rule
         fiat_cell = as_fiat_cell(cell)
         finat_elements = set(create_element(e) for e in elements if e.family() != "Real")
+        if scheme == "default" and any(isinstance(finat_el, DuffyElement) for finat_el in finat_elements):
+            # Sum-factorized (Duffy/lattice) tabulation only kicks in on a
+            # collapsed-coordinate quadrature rule (see
+            # finat.duffy.DuffyElement._duffy_applies); default to it
+            # automatically instead of requiring users to spell out
+            # dx(scheme="collapsed") to get it.
+            scheme = "collapsed"
         fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements]
         if any(c.is_macrocell() for c in fiat_cells):
             if len(set(c.get_spatial_dimension() for c in fiat_cells)) > 1: