Automate zany basis transformations

AGENTS.md

@@ -0,0 +1,157 @@
+# AGENTS.md for FIAT
+
+This document outlines the guidelines and architectural context for AI agents assisting with the FIAT codebase, functioning as a core component within the broader Firedrake ecosystem. 
+
+---
+
+## AI Contribution Policy
+
+When assisting with contributions to FIAT and the Firedrake project, AI agents and their human counterparts must adhere to the following strict policies:
+
+* The use of AI tools must be explicitly declared alongside the specific tool used.
+* Reviewer questions must be answered directly by the human, rather than acting as a relay to the AI.
+
+---
+
+## FIAT's Role in the Architecture
+
+FIAT operates within Firedrake's automated system for solving partial differential equations via the finite element method. Its specific architectural responsibilities include:
+
+* FIAT provides compile-time pre-tabulated basis functions.
+* These basis functions are utilized when the Two-Stage Form Compiler (TSFC) lowers Unified Form Language (UFL) into the GEM tensor language.
+* The resulting GEM expressions represent mathematical operations over quadrature points.
+* Mesh-topology bookkeeping routines within the ecosystem rely on FIAT ordering, such as the `create_cell_closure()` loop which builds a FIAT-ordered closure map necessary for subsequent code generation.
+
+---
+
+## Repository Layout
+
+Despite being called "fiat", this repository (Python package `firedrake_fiat`) contains **three** packages plus their tests, all in one tree:
+
+* `FIAT/` — reference-element definitions: cells, polynomial sets, dual bases (`FIAT/functional.py` holds the taxonomy of degree-of-freedom functionals), and element families.
+* `finat/` — symbolic layer on top of FIAT. Physically mapped elements (`finat/physically_mapped.py`, `finat/argyris.py`, `finat/morley.py`, `finat/hct.py`, `finat/piola_mapped.py`, …) construct GEM expressions for basis transformations. See `finat/AGENTS.md` for package-specific implementation conventions.
+* `gem/` — the tensor-algebra intermediate language used to express transformations symbolically (`gem/gem.py` for nodes, `gem/interpreter.py` to evaluate expressions numerically in tests). See `gem/AGENTS.md` for package-specific implementation conventions.
+* `test/` — note the singular name: tests live at `test/finat/test_zany_mapping.py`, `test/FIAT/...`, `test/gem/...` (not `tests/`).
+
+Read the subdirectory `AGENTS.md` for any package you are about to edit — they hold
+implementation-level conventions (GEM operator overloading rules, FInAT Jacobian sign
+conventions, etc.) that this top-level file does not repeat.
+
+## Environment and Setup
+
+Bugs can exist within Firedrake or any of its component packages, explicitly including FIAT. To effectively develop and debug FIAT:
+
+* Developers should use editable installs for subpackages like FIAT so that source code edits take effect without requiring a full reinstallation.
+* The active branch or commit of each component must be verified before assuming a bug originates in the top-level Firedrake package.
+
+---
+
+## Core Coding Rules
+
+Agents modifying FIAT code must follow these fundamental development principles:
+
+* Bug fixes must target the underlying mathematical or architectural root cause.
+* Developers must avoid merely patching specific failing test cases or edge cases.
+* Code complexity should be minimized by favoring the mathematical generality of finite elements over complicated special-case logic.
+* Memorized API shapes must not be trusted.
+* APIs across the ecosystem evolve, meaning properties can become methods, arguments can be renamed, and signatures can be deprecated.
+* Agents must verify current API signatures by reading the installed source code instead of relying on outdated training data.
+* Code documentation and comments must explain the present, correct code.
+* Comments must not detail what a removed or incorrect approach previously did.
+
+## Style and Conventions
+
+When writing Python code for FIAT, maintain the ecosystem's structural and stylistic integrity:
+
+* Class attributes must be declared in one visible location.
+* Attributes must be initialized in the `__init__` constructor or declared as a `functools.cached_property` if they are expensive to compute.
+* Ad hoc lazy initialization discovering attributes via `hasattr`, `setattr`, or `getattr` scattered across methods is strictly prohibited.
+* Boolean attributes must be used to record initialization intent and state instead of probing for the presence of state built by an initialization function.
+* New code must include type hints on all function and method signatures.
+* Public-facing APIs must include properly formatted `numpydoc`-style docstrings.
+
+## Pull Request Expectations
+
+* All changes are expected to arrive through GitHub Pull Requests.
+* Keep diffs reviewable and focused.
+* Before concluding work verify that the relevant subset of the pytest test suite
+  succeeds locally.
+* CI (`.github/workflows/test.yml`) checks `flake8` and `pydocstyle .` (both
+  configured in `setup.cfg`); run them locally and fix all findings before pushing.
+* Watch for `pydocstyle` D413 (blank line after the last numpydoc section) and D417
+  (every argument, including keyword arguments, needs a description) in new docstrings.
+
+---
+
+## Pattern Matching and Mathematical Reasoning
+
+This section used to be a rough draft of aspirations. It is now grounded in the concrete
+experience of automating the Kirby (2017) / Aznaran-Kirby-Farrell (2022) / Brubeck & Kirby
+(2025) transformation theory (`finat/zany.py`, `finat/functional.py`) for Morley, Hermite,
+Argyris, Bell, Mardal-Tai-Winther, Johnson-Mercier, and Guzman-Neilan. The lessons below are
+about *how to design and debug this kind of code*, not just about this one project.
+
+### Design strategies that generalize
+
+* **Design by duality, never by direct construction.** Do not try to write physical basis
+  functions directly (that requires inverting a Vandermonde system). Write the physical *node*
+  (a linear functional — point evaluation or moment, always computable), push it forward, and
+  solve for its expansion in the reference nodes. This is the master strategy behind the whole
+  framework and generalizes to any element whose nodes, not whose basis functions, are the
+  simple objects.
+* **Recover mathematical type from data, not from a class hierarchy.** `from_fiat` never asks
+  "is this an `IntegralMomentOfNormalDerivative`?"; it reads `pt_dict`/`deriv_dict` and derives
+  (order, direction, rank) numerically. This makes the framework forward-compatible with FIAT
+  functional classes that do not exist yet, as long as they reduce to the same numeric shape —
+  which essentially all "linear functional built from point/derivative data" dofs do. Prefer
+  this kind of structural recovery over adding another `isinstance`/type-tag branch whenever a
+  new functional needs support.
+* **Turn a mathematical case-split into a template method, not a conditional.** When a theorem
+  reads "for pullback type X do A, for pullback type Y do B" against an otherwise identical
+  algorithmic skeleton, encode the skeleton once (`ZanyPhysicallyMappedElement
+  .basis_transformation`) and the case-split as small, named hook methods on sibling mixins
+  (`ScalarPhysicallyMappedElement`, `PiolaPhysicallyMappedElement`). Keep the underlying linear
+  algebra as free functions taking plain arrays/GEM expressions with no `self`
+  (`FacetFrame`, `_scalar_facet_rows`, `_piola_facet_rows`, ...): the class layer should only
+  ever be a thin adapter from "element structure" to "which pure-math function to call," never
+  a place where new mathematics is derived. This also keeps the generic infrastructure
+  (`PhysicallyMappedElement` in `finat/physically_mapped.py`, used by hand-coded elements like
+  Arnold-Winther and HCT) free of any awareness of the theory-specific case-split.
+* **Validate one element at a time against an independently computed ground truth, not just
+  against the closed-form answer you derived.** `test/finat/test_zany_mapping.py
+  ::check_zany_mapping` tabulates the FIAT element on an actual physical cell and least-squares
+  fits the physical values against the pulled-back reference tabulation — a computation
+  completely independent of the symbolic derivation. Its assertion message pretty-prints the
+  (row, column) location of the relative error; read those indices first, they usually
+  localize the bug to one entity/dof-type combination before you write any new code.
+* **Generalize from one worked example, verify to machine precision, then extend — never
+  generalize speculatively ahead of a second concrete example.** The actual sequence that
+  worked: derive the mechanism on Morley alone, check it reproduces the deleted hand-coded
+  matrix bit-for-bit, *then* extend to Hermite, then Argyris/Bell, then the Piola family. Each
+  extension needed a genuinely new piece (interior invariance for Piola interior moments,
+  mixed-order jets for Bell/Argyris vertices, the reciprocal-basis correction for 3D) that a
+  premature generalization from Morley alone would not have anticipated.
+
+### Working with the human collaborator
+
+* **The literature is necessary but not sufficient — treat a stuck derivation as a cue to ask,
+  not to keep re-deriving.** Kirby (2017) and Aznaran-Kirby-Farrell (2022) work out 2D (or a
+  simplification that hides a 3D-only effect); none of the papers mention the reciprocal-basis
+  transformation law 3D Piola elements need. That fix came directly from the user's domain
+  expertise ("the papers do not consider the reciprocal basis that FIAT implements... Consider
+  this"), not from re-reading the papers harder. When a derivation matches every 2D case but
+  fails exactly where a paper's simplifying assumption would hide something, that is the moment
+  to ask what convention the *implementation* (not the paper) actually uses.
+* **Expect, and design for, architectural correction mid-project.** This project's real
+  trajectory was: prototype on Morley → generalize into a `Functional`/free-function design →
+  "don't extend further, make the computation generic first, get rid of the very specific
+  helpers" → extend to Argyris/Bell → extend to the Piola family → "refactor into
+  `Scalar`/`PiolaPhysicallyMappedElement` mixins so the `if piola` disappears from the loop" →
+  "move the method body into the Mixin" (keep the generic `PhysicallyMappedElement` free of
+  zany-specific knowledge). None of these were bug fixes; each was the human reshaping the
+  design once its true shape became visible. Treat working code as a draft the architecture
+  will still move under, and re-run the full verification (test suite, `flake8`, `pydocstyle`)
+  after every such reshaping, not only after functional changes.
+* **Keep this file's *why*, not just the *what*.** Git history has the diffs; this file exists
+  so the next session does not have to rediscover the reciprocal-basis subtlety, the
+  sign-invariance requirement, or the identity-row idiom by re-reading a diff from scratch.

FIAT/expansions.py

@@ -62,6 +62,55 @@ def jacobi_factors(x, y, z, dx, dy, dz):
     return fa, fb, fc, dfa, dfb, dfc
 
 
+def _product_derivative_term(factor: object,
+                             operand: numpy.ndarray,
+                             factor_axes: tuple[int, ...],
+                             rank: int) -> numpy.ndarray:
+    """Return one Leibniz term in an ordered derivative tensor."""
+    factor = numpy.asarray(factor)
+    factor_rank = len(factor_axes)
+    operand_rank = rank - factor_rank
+    factor_index = (slice(None),) * factor_rank + (None,) * operand_rank + (Ellipsis,)
+    operand_index = (None,) * factor_rank + (Ellipsis,)
+    product = factor[factor_index] * operand[operand_index]
+
+    if rank:
+        dim = factor.shape[0] if operand_rank == 0 else operand.shape[0]
+        factor_source = {axis: source for source, axis in enumerate(factor_axes)}
+        operand_axes = tuple(axis for axis in range(rank) if axis not in factor_axes)
+        operand_source = {axis: factor_rank + source
+                          for source, axis in enumerate(operand_axes)}
+        permutation = tuple((factor_source | operand_source)[axis]
+                            for axis in range(rank))
+        product = product.transpose(permutation + tuple(range(rank, product.ndim)))
+        product = product.reshape((dim,) * rank + operand.shape[operand_rank:])
+
+    return product
+
+
+def _product_derivative(factor: object,
+                        dfactor: object | None,
+                        ddfactor: object | None,
+                        operands: list[numpy.ndarray],
+                        rank: int) -> numpy.ndarray:
+    """Differentiate a recurrence factor times a basis derivative tensor."""
+    # The Dubiner recurrence factors are linear or quadratic, so higher
+    # factor derivatives vanish and the Leibniz sum stops at order two.
+    result = _product_derivative_term(factor, operands[rank], (), rank)
+
+    if dfactor is not None and rank >= 1:
+        for axis in range(rank):
+            result += _product_derivative_term(dfactor, operands[rank-1], (axis,), rank)
+
+    if ddfactor is not None and rank >= 2:
+        for axis0 in range(rank):
+            for axis1 in range(axis0 + 1, rank):
+                result += _product_derivative_term(ddfactor, operands[rank-2],
+                                                   (axis0, axis1), rank)
+
+    return result
+
+
 def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
     """Tabulate a Dubiner expansion set using the recurrence from (Kirby 2010).
 
@@ -77,8 +126,6 @@ def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
 
     :returns: A tuple with tabulations of the expansion set and its derivatives.
     """
-    if order > 2:
-        raise ValueError("Higher order derivatives not supported")
     if variant not in [None, "bubble", "dual"]:
         raise ValueError(f"Invalid variant {variant}")
     if variant == "bubble":
@@ -95,7 +142,7 @@ def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
     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 = results[0]
     phi[0] = scale
     if dim == 0 or n == 0:
         return results
@@ -127,14 +174,12 @@ def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
 
             fcur = a * fa - b * fb
             phi[inext] = fcur * phi[icur]
-            if dphi is not None:
+            if order:
                 dfcur = a * dfa - b * dfb
-                dphi[inext] = phi[icur] * dfcur
-                dphi[inext] += fcur * dphi[icur]
-                if ddphi is not None:
-                    ddphi[inext] = outer(dphi[icur], dfcur)
-                    ddphi[inext] += outer(dfcur, dphi[icur])
-                    ddphi[inext] += fcur * ddphi[icur]
+                cur = [result[icur] for result in results]
+                for rank in range(1, order+1):
+                    results[rank][inext] = _product_derivative(fcur, dfcur, None,
+                                                               cur, rank)
 
             # general i by recurrence
             for i in range(1, n - sum(sub_index)):
@@ -145,26 +190,17 @@ def dubiner_recurrence(dim, n, order, ref_pts, Jinv, scale, variant=None):
                 fprev = -c * fc
                 phi[inext] = fcur * phi[icur]
                 phi[inext] += fprev * phi[iprev]
-                if dphi is None:
-                    continue
 
                 dfcur = a * dfa - b * dfb
                 dfprev = -c * dfc
-                dphi[inext] = phi[icur] * dfcur
-                dphi[inext] += phi[iprev] * dfprev
-                dphi[inext] += fcur * dphi[icur]
-                dphi[inext] += fprev * dphi[iprev]
-                if ddphi is None:
-                    continue
-
                 ddfprev = -c * ddfc
-                ddphi[inext] = phi[iprev] * ddfprev
-                ddphi[inext] += outer(dphi[icur], dfcur)
-                ddphi[inext] += outer(dfcur, dphi[icur])
-                ddphi[inext] += outer(dphi[iprev], dfprev)
-                ddphi[inext] += outer(dfprev, dphi[iprev])
-                ddphi[inext] += fcur * ddphi[icur]
-                ddphi[inext] += fprev * ddphi[iprev]
+                cur = [result[icur] for result in results]
+                prev = [result[iprev] for result in results]
+                for rank in range(1, order+1):
+                    results[rank][inext] = _product_derivative(fcur, dfcur, None,
+                                                               cur, rank)
+                    results[rank][inext] += _product_derivative(fprev, dfprev, ddfprev,
+                                                                prev, rank)
 
         # normalize
         d = codim + 1
@@ -291,7 +327,7 @@ def __init__(self, ref_el, scale=None, variant=None):
         self.scale = scale
         self.variant = variant
         self.continuity = "C0" if variant == "bubble" else None
-        self.recurrence_order = 2
+        self.recurrence_order = math.inf
         self._dmats_cache = {}
         self._cell_node_map_cache = {}
 
@@ -353,7 +389,7 @@ def _tabulate_on_cell(self, n, pts, order=0, cell=0, direction=None):
         def distance(alpha, beta):
             return sum(ai != bi for ai, bi in zip(alpha, beta))
 
-        # Only use dmats if tabulate failed
+        # Use dmats only for derivatives above the configured recurrence order.
         for i in range(len(phi), order + 1):
             dmats = self.get_dmats(n, cell=cell)
             for alpha in mis(sd, i):

FIAT/macro.py

@@ -484,20 +484,31 @@ def __init__(self, ref_el, degree, order=1, vorder=None, shape=(), **kwargs):
         # Impose C^vorder super-smoothness at interior vertices
         # C^forder automatically gives C^{forder+dim-1} at the interior vertex
         verts = numpy.asarray(ref_el.get_vertices())
+        has_vertex_constraints = False
         for vorder in set(order[0].values()):
             vids = [i for i in order[0] if order[0][i] == vorder]
             facets = chain.from_iterable(ref_el.connectivity[(0, sd-1)][v] for v in vids)
             forder = min(order[sd-1][f] for f in facets)
             sorder = forder + sd - 1
             if vorder > sorder:
+                has_vertex_constraints = True
                 jumps = expansion_set.tabulate_jumps(degree, verts[vids], order=vorder)
                 rows.extend(numpy.vstack(jumps[r].T) for r in range(sorder+1, vorder+1))
 
         if len(rows) > 0:
             for row in rows:
                 row *= 1 / max(numpy.max(abs(row)), 1)
-            dual_mat = numpy.vstack(rows)
-            coeffs = polynomial_set.spanning_basis(dual_mat, nullspace=True)
+            if has_vertex_constraints:
+                # Project each constraint block onto the current nullspace so
+                # high-order vertex constraints are not buried in one large SVD.
+                coeffs = numpy.eye(expansion_set.get_num_members(degree))
+                for row in rows:
+                    restricted_row = numpy.dot(row, coeffs.T)
+                    nsp = polynomial_set.spanning_basis(restricted_row, nullspace=True, rtol=1e-12)
+                    coeffs = numpy.dot(nsp, coeffs)
+            else:
+                dual_mat = numpy.vstack(rows)
+                coeffs = polynomial_set.spanning_basis(dual_mat, nullspace=True)
         else:
             coeffs = numpy.eye(expansion_set.get_num_members(degree))