From e7b44968ba2c1f3a768181f9173ef21de9c79040 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 16 Jul 2026 09:43:28 +0100 Subject: [PATCH 01/10] TSFC: support same-mesh ufl.Interpolate in Form kernels --- AGENTS.md | 82 ++++++++++- firedrake/assemble.py | 13 +- .../regression/test_adjoint_operators.py | 14 ++ .../firedrake/regression/test_interp_dual.py | 79 +++++++++++ tsfc/driver.py | 94 +------------ tsfc/fem.py | 132 +++++++++++++++++- tsfc/modified_terminals.py | 8 +- tsfc/ufl_utils.py | 33 ++++- 8 files changed, 351 insertions(+), 104 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index ffc2ca2295..c5568c4e5e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,8 +1,10 @@ # Firedrake Firedrake is an automated system for the portable solution of partial differential equations using -the finite element method (FEM). The codebase is primarily Python, relying heavily on code generation -and high-performance C backends to achieve scalability and speed. +the finite element method. The codebase is primarily Python, relying heavily on code generation +and high-performance C backends to achieve scalability and speed. Firedrake is highly composable and +fully differentiable, enabling the automatic generation of tangent linear and adjoint models for +PDE-constrained optimization. Firedrake's full contribution process is documented at [Contributing to Firedrake](https://firedrakeproject.org/contribute.html). In short, for AI-assisted @@ -30,6 +32,12 @@ toolchain: 2. **Lowering to Loopy:** The GEM expressions are then lowered into **loopy** kernels. * **PyOP2:** Finally, the generated loopy kernels are wrapped and executed by PyOP2, which handles the parallel execution of loops over mesh cells and facets. +* **pyadjoint:** Firedrake integrates with `pyadjoint` for algorithmic differentiation. Annotated + operations (assembly, interpolation, variational solves, boundary condition application, ...) are + recorded on a `Tape` as a DAG of `Block`s; a `ReducedFunctional` composed with one or more `Control`s + can then be evaluated, differentiated (adjoint or tangent-linear), and checked with a `taylor_test`. + Taping happens at the level of Firedrake's own operations, not the underlying numerics, so any new + feature assembled purely from already-annotated building blocks is differentiable automatically. ## Core Working Rules @@ -56,6 +64,13 @@ toolchain: prose explaining what the removed, incorrect approach used to do or why it was wrong. Keep comments and documentation focused on the current, correct code; a reader should never need the history of what used to be there to understand why the present code is right. +* **Composability And Differentiability:** New features are expected to compose with existing ones + without special-casing, and, via `pyadjoint`, to remain differentiable when built from already-taped + operations. Prefer the annotated, top-level API (e.g. `firedrake.assemble`) over a lower-level + equivalent that bypasses it (e.g. calling an `Interpolator`'s own `.assemble()` directly) even when + both give the same forward numbers — the lower-level call silently drops out of the tape, and no test + will notice unless it specifically exercises pyadjoint. When a change could plausibly sit on the tape, + verify differentiability explicitly with a `taylor_test`, not just a forward-value check. ## Coding Style And Conventions @@ -90,6 +105,29 @@ toolchain: excluded from that rule, so its snippet renders in the docs but never runs. Prefer `::` — reach for the directive only for an illustrative fragment naming things the demo never defines. +## Design And Debugging Method + +How to spend effort on any feature or bug: the first three shape a design before writing code, +the last two localize a failure before reading code. + +* **Design by nearest working neighbor.** Some existing feature already solves a structurally + identical problem. Grep for the invariant yours must satisfy (the type handled, the hook fired, + the kwarg accepted), read how the neighbor earns it, and implement only the delta. +* **Write the state contract before the code.** For anything flowing through a cached, replayed, + or lazily-refreshed system, answer up front: who owns it, when is it refreshed, what happens on + reuse or in-place mutation, how does a replay recover it. Stale-state bugs fail far from their + cause and only under reuse. +* **Classify the mathematical structure before choosing machinery.** Affine and linear + dependencies have closed-form contributions (identities, fixed operators, exact zeros for higher + derivatives); recognizing an exact zero lets you skip a code path instead of fixing it. +* **Attribute before you analyze.** Shrink *where* with one-delta experiments — each ingredient + toggled in isolation against a known-good baseline, consecutive CI failure sets diffed, a single + hunk reverted, the merge base rerun — then read code for *why*. +* **Census the consumers before changing a lifecycle.** Changing *when* or *how often* something + is computed (rather than its value) is safe only after grepping every access site: some consumer + calls it in a context you did not design for (per nonlinear iteration, inside a PETSc callback, + on every attribute read). + ## Testing Requirements * **Pull Requests:** All PRs must include comprehensive tests demonstrating that the new feature works @@ -117,6 +155,12 @@ toolchain: in the install docs to get a component installed in editable mode so source edits take effect without reinstalling, and check which branch/commit of each component is actually active before assuming a fix belongs in Firedrake itself. +* **Branch pairing across the stack:** Firedrake's `main` and `release` branches go hand in hand with + the `main` and `release` branches of its components (FIAT, UFL, ...). A CI failure may be + unreproducible locally simply because a component checkout in the venv sits on some other branch — + check `git -C $VIRTUAL_ENV/src/ branch`, switch to the branch matching the Firedrake branch + under test, run `firedrake-clean`, and reproduce again before hunting for the bug in Firedrake + itself. * **`petsc4py`/PETSc version skew:** `petsc4py` is a compiled extension built against one specific PETSc checkout. If you switch the PETSc branch/commit underneath an existing venv (e.g. to bisect a PETSc-side issue) without rebuilding `petsc4py` against it, `import firedrake` fails with a confusing @@ -153,12 +197,25 @@ toolchain: conclude a parallel code path is untested just because a plain, unmarked `pytest` run was green. * **Splitting for CI:** `firedrake-run-split-tests` shards the suite by process count for CI; look at it (and `.github/workflows/pr.yml`/`core.yml`) if a failure only reproduces in CI and not locally. +* **CI triage:** `gh pr checks ` lists job statuses. When `gh run view --log` returns nothing + (it does for very large logs), download the log with + `gh api repos///actions/jobs//logs` and grep for `FAILED`. Before debugging + anything, fetch the failure list of the *previous* run of the same PR: the difference between the + two failure sets attributes each failure to the commits pushed in between. * **Narrow reproduction first:** Run the single failing test node (`pytest path::test_name -k ...`) before the full module; the suite is large and full-module reruns are slow to iterate against. +* **Test mathematical correctness, not just that it runs or looks structurally right.** Neither "no + exception was raised" nor `==` agreement between two expressions proves the result is + correct — two independently-built expressions can match structurally while sharing the same wrong + derivative or simplification rule. Verify the actual mathematical claim: evaluate numerically and + compare against a hand-computed or finite-difference value, or use a Taylor test for anything + claiming to be a derivative. +* **Taylor-test-everything is the immune system:** Taylor-test a `ReducedFunctional`, ensuring that any + new feature built from existing, annotated Firedrake operations is automatically differentiable. ### Debugging -* **Generated kernels (niche, rarely needed):** By default, generated C is compiled optimized and +* **Generated kernels:** By default, generated C is compiled optimized and without debug symbols, so a debugger attached to the Python process cannot meaningfully step through it. Set `PYOP2_DEBUG=1` to compile with `-O0 -g` instead, which is the prerequisite for using `gdb`/`cgdb` on the compiled kernel at all. @@ -169,7 +226,7 @@ toolchain: computed differently per rank and fed into code generation (e.g. a rank-local decision that should be a collective/global one) — make that decision the same on every rank, rather than patching the generated source or the difference itself. -* **Parallel deadlocks (niche, rarely needed):** `PYOP2_SPMD_STRICT=1` adds barriers around calls +* **Parallel deadlocks:** `PYOP2_SPMD_STRICT=1` adds barriers around calls marked `@collective` and around cache access, trading overhead for a much narrower failure point when ranks disagree about control flow. * **Logging:** `firedrake.logging.set_log_level()` (or the `PYOP2_LOG_LEVEL` environment variable) @@ -178,6 +235,23 @@ toolchain: standard PETSc options (`-ksp_view`, `-snes_view`, `-ksp_monitor`, `-log_view`, `-start_in_debugger`) can be passed through Firedrake's `solver_parameters` or the command line exactly as in a plain PETSc application. +* **Errors inside PETSc callbacks do not surface as their own traceback:** under pytest they often + appear as a bare `Segmentation fault` with no Firedrake frames; standalone they appear as + `petsc4py.PETSc.Error: error code 101` whose *first* chained traceback (e.g. a `TypeError` about + the callback context in `petscsnes.pxi`) describes the corrupted callback state, not the cause. + Rerun the failing test as a standalone script to expose the chained tracebacks, and read the PETSc + call stack inside the error (`PCSetUp_MG` → `SNESComputeFunction`, ...) to identify *which* + callback was executing. +* **Construction-time code re-runs inside solver callbacks:** geometric multigrid coarsens the + entire problem lazily inside `PCSetUp`, through the `coarsen` singledispatch in + `firedrake/mg/ufl_utils.py` — whatever a feature does at construction time (e.g. `DirichletBC` + interpolating or projecting its boundary value into the space) is re-executed per level inside + that PETSc callback. Objects that carry solvers or attach DM hooks (a `Projector`, a variational + solver) must be built once and cached, never rebuilt on each call of an accessor that may fire + there: repeatedly constructing and garbage-collecting a solver stack inside `PCSetUp_MG` corrupts + the DM callback state and segfaults far from the allocation site. Extruded (hexahedral) + hierarchies are the stress test: tensor-product elements (NCE/NCF) have no dual-basis + interpolation, so paths that interpolate on simplices take the projection fallback there. ### Reproducible Environments diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6dd7c4f03f..67a32be63f 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -170,7 +170,18 @@ def get_assembler(form, *args, **kwargs): # Preprocess the DAG and restructure the DAG # Only pre-process `form` once beforehand to avoid pre-processing for each assembly call form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) - if isinstance(form, (ufl.form.Form, slate.TensorBase)) and not BaseFormAssembler.base_form_operands(form): + base_form_operands = BaseFormAssembler.base_form_operands(form) + can_compile = not base_form_operands + if isinstance(form, ufl.form.Form) and len(form.ufl_domains()) == 1: + domain, = form.ufl_domains() + can_compile = all( + isinstance(op, ufl.Interpolate) + and op.ufl_function_space().ufl_domain() == domain + and set(extract_domains(op)) <= {domain} + for op in base_form_operands + ) + + if isinstance(form, (ufl.form.Form, slate.TensorBase)) and can_compile: diagonal = kwargs.pop('diagonal', False) if len(form.arguments()) == 0: return ZeroFormAssembler(form, form_compiler_parameters=fc_params) diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index 57faf80477..76331e9893 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -83,6 +83,20 @@ def test_interpolate_with_arguments(rg): assert taylor_test(rf, f, h) > 1.9 +@pytest.mark.skipcomplex +def test_interpolate_in_form(rg): + mesh = UnitSquareMesh(3, 3) + V = FunctionSpace(mesh, "CG", 1) + W = FunctionSpace(mesh, "DG", 0) + x, y = SpatialCoordinate(mesh) + f = Function(V).interpolate(x + 2 * y) + + J = assemble(interpolate(f, W) ** 2 * dx) + rf = ReducedFunctional(J, Control(f)) + + assert taylor_test(rf, f, rg.uniform(V)) > 1.9 + + @pytest.mark.skipcomplex # Taping for complex-valued 0-forms not yet done def test_interpolate_scalar_valued(rg): mesh = IntervalMesh(10, 0, 1) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index b0ce971a95..d9bd4c08d4 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -395,3 +395,82 @@ def test_assemble_action_adjoint(V1, V2): assert isinstance(res4, Cofunction) assert res4.function_space() == V1.dual() assert np.allclose(res.dat.data, res4.dat.data) + + +@pytest.mark.parallel(2) +def test_interpolate_in_form_compiled_reuse(): + from firedrake.assemble import OneFormAssembler, get_assembler + + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 2) + W = FunctionSpace(mesh, "DG", 1) + x, y = SpatialCoordinate(mesh) + u = Function(V) + v = TestFunction(W) + interpolation = interpolate(u, W) + form = inner(interpolation, v) * dx + assembler = get_assembler(form) + + assert isinstance(assembler, OneFormAssembler) + for scale in (1, 3): + u.interpolate(scale * (x + y)) + actual = assembler.assemble() + expected = assemble(inner(assemble(interpolation), v) * dx) + assert np.allclose(actual.dat.data, expected.dat.data) + + +def test_interpolate_in_form_mapped_derivative(): + mesh = UnitSquareMesh(2, 2) + V = VectorFunctionSpace(mesh, "CG", 2) + W = FunctionSpace(mesh, "RT", 1) + x, y = SpatialCoordinate(mesh) + u = Function(V).interpolate(as_vector((x**2 + y, x - y**2))) + interpolation = interpolate(u, W) + + actual = assemble(inner(grad(interpolation), grad(interpolation)) * dx) + interpolated = assemble(interpolation) + expected = assemble(inner(grad(interpolated), grad(interpolated)) * dx) + assert np.isclose(actual, expected) + + +def test_interpolate_in_bilinear_form(): + mesh = UnitIntervalMesh(3) + V = FunctionSpace(mesh, "CG", 1) + W = FunctionSpace(mesh, "DG", 0) + u = TrialFunction(V) + v = TestFunction(W) + operator = assemble(inner(interpolate(u, W), v) * dx) + + x, = SpatialCoordinate(mesh) + f = Function(V).interpolate(x + 1) + actual = assemble(action(operator, f)) + expected = assemble(inner(assemble(interpolate(f, W)), v) * dx) + assert np.allclose(actual.dat.data, expected.dat.data) + + +def test_interpolate_in_interior_facet_form(): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 2) + W = FunctionSpace(mesh, "DG", 1) + x, y = SpatialCoordinate(mesh) + u = Function(V).interpolate(x**2 + y) + v = TestFunction(W) + interpolation = interpolate(u, W) + + actual = assemble(jump(interpolation) * jump(v) * dS) + interpolated = assemble(interpolation) + expected = assemble(jump(interpolated) * jump(v) * dS) + assert np.allclose(actual.dat.data, expected.dat.data) + + +def test_cross_mesh_interpolate_in_form_uses_base_form_assembler(): + from firedrake.assemble import BaseFormAssembler, get_assembler + + source_mesh = UnitSquareMesh(1, 1) + target_mesh = UnitSquareMesh(1, 1) + V = FunctionSpace(source_mesh, "CG", 1) + W = FunctionSpace(target_mesh, "CG", 1) + v = TestFunction(W) + form = inner(interpolate(Function(V), W), v) * dx(domain=target_mesh) + + assert isinstance(get_assembler(form), BaseFormAssembler) diff --git a/tsfc/driver.py b/tsfc/driver.py index 2c480f9c55..71efba33d0 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -3,7 +3,6 @@ import sys from itertools import chain import numpy -from finat.physically_mapped import NeedsCoordinateMappingElement import ufl from ufl.algorithms import extract_coefficients @@ -20,7 +19,6 @@ from tsfc import fem, ufl_utils from tsfc.logging import logger -from tsfc.modified_terminals import analyse_modified_terminal from tsfc.parameters import default_parameters, is_complex from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants, simplify_abs import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy @@ -81,6 +79,7 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di assert isinstance(form, Form) GREEN = "\033[1;37;32m%s\033[0m" + form = ufl_utils.lower_form_interpolations(form) # Determine whether in complex mode: complex_mode = parameters and is_complex(parameters.get("scalar_type")) @@ -332,33 +331,9 @@ def compile_expression_dual_evaluation(expression, ufl_element, *, if isinstance(to_element, finat.QuadratureElement): kernel_cfg["quadrature_rule"] = to_element._rule - dual_arg, operand = expression.argument_slots() - - # Create callable for translation of UFL expression to gem - fn = DualEvaluationCallable(operand, kernel_cfg) - - # Get the gem expression for dual evaluation and corresponding basis - # indices needed for compilation of the expression - if isinstance(to_element, NeedsCoordinateMappingElement): - ctx = fem.PointSetContext(**kernel_cfg) - mt = analyse_modified_terminal(ufl.Coefficient(dual_arg.ufl_function_space().dual())) - coordinate_mapping = fem.CoordinateMapping(mt, ctx) - else: - coordinate_mapping = None - evaluation, point_indices, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) - quadrature_multiindex = tuple(point_indices) - - # Compute the action against the dual argument - if isinstance(dual_arg, ufl.Cofunction): - gem_dual = builder.coefficient_map[dual_arg] - if complex_mode: - evaluation = gem.MathFunction('conj', evaluation) - # The dual argument contracts over the nodes, so the basis indices are - # reduction indices like the points, not indices of the return value. - evaluation = evaluation * gem_dual[basis_indices] - quadrature_multiindex += tuple(basis_indices) - basis_indices = () - else: + evaluation, quadrature_multiindex, basis_indices = fem.dual_evaluate(expression, to_element, kernel_cfg) + dual_arg, _ = expression.argument_slots() + if not isinstance(dual_arg, ufl.Cofunction): argument_multiindices[dual_arg.number()] = basis_indices argument_multiindices = dict(sorted(argument_multiindices.items())) @@ -387,64 +362,3 @@ def compile_expression_dual_evaluation(expression, ufl_element, *, builder.set_output(return_var) # Build kernel tuple return builder.construct_kernel(impero_c, index_names, needs_external_coords, parameters["add_petsc_events"], name=name) - - -class DualEvaluationCallable(object): - """ - Callable representing a function to dual evaluate. - - When called, this takes in a - :class:`finat.point_set.AbstractPointSet` and returns a GEM - expression for evaluation of the function at those points. - - :param expression: UFL expression for the function to dual evaluate. - :param kernel_cfg: A kernel configuration for creation of a - :class:`GemPointContext` or a :class:`PointSetContext` - - Not intended for use outside of - :func:`compile_expression_dual_evaluation`. - """ - def __init__(self, expression, kernel_cfg): - self.expression = expression - self.kernel_cfg = kernel_cfg - - def __call__(self, ps): - """The function to dual evaluate. - - :param ps: The :class:`finat.point_set.AbstractPointSet` for - evaluating at - :returns: a gem expression representing the evaluation of the - input UFL expression at the given point set ``ps``. - For point set points with some shape ``(*value_shape)`` - (i.e. ``()`` for scalar points ``(x)`` for vector points - ``(x, y)`` for tensor points etc) then the gem expression - has shape ``(*value_shape)`` and free indices corresponding - to the input :class:`finat.point_set.AbstractPointSet`'s - free indices alongside any input UFL expression free - indices. - """ - - if not isinstance(ps, finat.point_set.AbstractPointSet): - raise ValueError("Callable argument not a point set!") - - # Avoid modifying saved kernel config - kernel_cfg = self.kernel_cfg.copy() - - if isinstance(ps, finat.point_set.UnknownPointSet): - # Run time known points - kernel_cfg.update(point_indices=ps.indices, point_expr=ps.expression) - # GemPointContext's aren't allowed to have quadrature rules - kernel_cfg.pop("quadrature_rule", None) - translation_context = fem.GemPointContext(**kernel_cfg) - else: - # Compile time known points - kernel_cfg.update(point_set=ps) - translation_context = fem.PointSetContext(**kernel_cfg) - - gem_expr, = fem.compile_ufl(self.expression, translation_context, point_sum=False) - # In some cases ps.indices may be dropped from expr, but nothing - # new should now appear - argument_multiindices = kernel_cfg["argument_multiindices"].values() - assert set(gem_expr.free_indices) <= set(chain(ps.indices, *argument_multiindices)) - - return gem_expr diff --git a/tsfc/fem.py b/tsfc/fem.py index 943089052e..3657885a4b 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -3,8 +3,10 @@ import collections import itertools +from itertools import chain from functools import cached_property, singledispatch +import finat import gem import numpy import ufl @@ -13,7 +15,9 @@ from FIAT.reference_element import TensorProductCell from finat.physically_mapped import (NeedsCoordinateMappingElement, PhysicalGeometry) +from finat.finiteelementbase import FiniteElementBase from finat.point_set import PointSet, PointSingleton +from finat.point_set import AbstractPointSet, UnknownPointSet from finat.quadrature import make_quadrature from finat.element_factory import as_fiat_cell, create_element from gem.node import traversal @@ -37,6 +41,7 @@ from tsfc.kernel_interface import ProxyKernelInterface from tsfc.kernel_interface.common import lower_integral_type from tsfc.modified_terminals import (analyse_modified_terminal, + ModifiedTerminal, construct_modified_terminal) from tsfc.parameters import is_complex from tsfc.ufl_utils import (ModifiedTerminalMixin, PickRestriction, @@ -327,6 +332,80 @@ def needs_coordinate_mapping(element): return isinstance(create_element(element), NeedsCoordinateMappingElement) +def dual_evaluate(expression: ufl.Interpolate, to_element: FiniteElementBase, kernel_cfg: dict) -> tuple: + """Translate an interpolation operand and evaluate its target dual basis. + + Parameters + ---------- + expression : ufl.Interpolate + Preprocessed interpolation expression in the target reference frame. + to_element : finat.FiniteElementBase + Target FInAT element. + kernel_cfg : dict + Configuration for the point-evaluation translation context. + + Returns + ------- + tuple + The GEM expression for the local interpolated values, the multiindex + to contract it on, and the basis indices of the return value. + """ + dual_arg, operand = expression.argument_slots() + fn = DualEvaluationCallable(operand, kernel_cfg) + + if isinstance(to_element, NeedsCoordinateMappingElement): + ctx = PointSetContext(**kernel_cfg) + coefficient = ufl.Coefficient(dual_arg.ufl_function_space().dual()) + coordinate_mapping = CoordinateMapping(analyse_modified_terminal(coefficient), ctx) + else: + coordinate_mapping = None + evaluation, point_indices, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) + quadrature_multiindex = tuple(point_indices) + + if isinstance(dual_arg, ufl.Cofunction): + gem_dual = kernel_cfg["interface"].coefficient_map[dual_arg] + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + # The dual argument contracts over the nodes, so the basis indices are + # reduction indices like the points, not indices of the return value. + evaluation = evaluation * gem_dual[basis_indices] + quadrature_multiindex += tuple(basis_indices) + basis_indices = () + + return evaluation, quadrature_multiindex, basis_indices + + +class DualEvaluationCallable: + """Translate an expression at points requested by a FInAT dual basis.""" + + def __init__(self, expression: ufl.core.expr.Expr, kernel_cfg: dict) -> None: + self.expression = expression + self.kernel_cfg = kernel_cfg + + def __call__(self, point_set: AbstractPointSet) -> gem.Node: + if not isinstance(point_set, AbstractPointSet): + raise ValueError("Callable argument not a point set!") + + kernel_cfg = self.kernel_cfg.copy() + if isinstance(point_set, UnknownPointSet): + kernel_cfg.update(point_indices=point_set.indices, + point_expr=point_set.expression) + kernel_cfg.pop("quadrature_rule", None) + translation_context = GemPointContext(**kernel_cfg) + else: + kernel_cfg.update(point_set=point_set) + translation_context = PointSetContext(**kernel_cfg) + + gem_expr, = compile_ufl(self.expression, translation_context, point_sum=False) + argument_multiindices = kernel_cfg["argument_multiindices"] + if hasattr(argument_multiindices, "values"): + argument_multiindices = argument_multiindices.values() + assert set(gem_expr.free_indices) <= set( + chain(point_set.indices, *argument_multiindices) + ) + return gem_expr + + @serial_cache(hashkey=lambda *args: args) def get_quadrature_rule(fiat_cell, integration_dim, quadrature_degree, scheme): integration_cell = fiat_cell.construct_subcomplex(integration_dim) @@ -739,11 +818,41 @@ def translate_constant_value(terminal, mt, ctx): return ctx.constant(terminal) +@translate.register(ufl.Interpolate) +def translate_interpolate(terminal: ufl.Interpolate, mt: ModifiedTerminal, ctx: ContextBase) -> gem.Node: + domain = extract_unique_domain(terminal) + element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) + kernel_cfg = { + "interface": CellVolumeKernelInterface( + ctx, domain, mt.restriction), + "ufl_cell": domain.ufl_cell(), + "integration_dim": as_fiat_cell(domain.ufl_cell()).get_dimension(), + "argument_multiindices": ctx.argument_multiindices, + "index_cache": ctx.index_cache, + "scalar_type": ctx.scalar_type, + } + if isinstance(element, finat.QuadratureElement): + kernel_cfg["quadrature_rule"] = element._rule + + evaluation, quadrature_multiindex, basis_indices = dual_evaluate(terminal, element, kernel_cfg) + # The interpolation points are internal to the local solve, so contract + # them here: only the form's own quadrature points stay free. + evaluation = gem.IndexSum(evaluation, quadrature_multiindex) + vec = gem.ComponentTensor(evaluation, basis_indices) + return translate_element(terminal, mt, ctx, vec, element, beta=element.get_indices()) + + @translate.register(Coefficient) def translate_coefficient(terminal, mt, ctx): - domain = extract_unique_domain(terminal) vec = ctx.coefficient(terminal, mt.restriction) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) + return translate_element(terminal, mt, ctx, vec, element) + + +def translate_element(terminal: ufl.core.expr.Expr, mt: ModifiedTerminal, ctx: ContextBase, + vec: gem.Node, element: FiniteElementBase, beta: tuple | None = None) -> gem.Node: + """Evaluate local finite element values at the current points.""" + domain = extract_unique_domain(terminal) # Collect FInAT tabulation for all entities per_derivative = collections.defaultdict(list) @@ -771,15 +880,26 @@ def take_singleton(xs): for alpha, tables in per_derivative.items()} # Coefficient evaluation - beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) + if beta is None: + beta = ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) zeta = element.get_value_indices() vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) value_dict = {} for alpha, table in per_derivative.items(): table_qi = gem.Indexed(table, beta + zeta) + if not hasattr(vec_beta, "index_ordering"): + value = gem.IndexSum(gem.Product(vec_beta, table_qi), beta) + value_dict[alpha] = gem.ComponentTensor(gem.optimise.contraction(value), zeta) + continue + summands = [] + argument_multiindices = ctx.argument_multiindices + if hasattr(argument_multiindices, "values"): + argument_multiindices = argument_multiindices.values() + unsummed_indices = set(chain(*argument_multiindices)) + unsummed_indices.update(ctx.unsummed_coefficient_indices) for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): - indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) + indices = tuple(i for i in var.index_ordering() if i not in unsummed_indices) value = gem.IndexSum(gem.Product(expr, var), indices) summands.append(gem.optimise.contraction(value)) optimised_value = gem.optimise.make_sum(summands) @@ -788,7 +908,11 @@ def take_singleton(xs): # Change from FIAT to UFL arrangement result = fiat_to_ufl(value_dict, mt.local_derivatives) assert result.shape == mt.expr.ufl_shape - assert set(result.free_indices) - ctx.unsummed_coefficient_indices <= set(ctx.point_indices) + argument_multiindices = ctx.argument_multiindices + if hasattr(argument_multiindices, "values"): + argument_multiindices = argument_multiindices.values() + allowed_indices = set(chain(ctx.point_indices, *argument_multiindices)) + assert set(result.free_indices) - ctx.unsummed_coefficient_indices <= allowed_indices # Detect Jacobian of affine cells if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2 diff --git a/tsfc/modified_terminals.py b/tsfc/modified_terminals.py index a26e5c2980..2b63022188 100644 --- a/tsfc/modified_terminals.py +++ b/tsfc/modified_terminals.py @@ -23,7 +23,7 @@ from ufl.classes import (ReferenceValue, ReferenceGrad, NegativeRestricted, PositiveRestricted, Restricted, ConstantValue, - Jacobian, SpatialCoordinate, Zero) + Interpolate, Jacobian, SpatialCoordinate, Zero) from ufl.checks import is_cellwise_constant from ufl.domain import extract_unique_domain @@ -82,7 +82,7 @@ def __str__(self): def is_modified_terminal(v): "Check if v is a terminal or a terminal wrapped in terminal modifier types." - while not v._ufl_is_terminal_: + while not (v._ufl_is_terminal_ or isinstance(v, Interpolate)): if v._ufl_is_terminal_modifier_: v = v.ufl_operands[0] else: @@ -92,7 +92,7 @@ def is_modified_terminal(v): def strip_modified_terminal(v): "Extract core Terminal from a modified terminal or return None." - while not v._ufl_is_terminal_: + while not (v._ufl_is_terminal_ or isinstance(v, Interpolate)): if v._ufl_is_terminal_modifier_: v = v.ufl_operands[0] else: @@ -115,7 +115,7 @@ def analyse_modified_terminal(expr): # Start with expr and strip away layers of modifiers t = expr - while not t._ufl_is_terminal_: + while not (t._ufl_is_terminal_ or isinstance(t, Interpolate)): if isinstance(t, ReferenceValue): assert reference_value is None, "Got twice pulled back terminal!" reference_value = True diff --git a/tsfc/ufl_utils.py b/tsfc/ufl_utils.py index ce25dc3087..e5f7af13ea 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -5,6 +5,7 @@ import numpy import ufl +from ufl.algorithms.map_integrands import map_integrand_dags from ufl import as_tensor, indices, replace from ufl.algorithms import compute_form_data as ufl_compute_form_data from ufl.algorithms import estimate_total_polynomial_degree @@ -23,7 +24,7 @@ from ufl.geometry import QuadratureWeight from ufl.geometry import Jacobian, JacobianDeterminant, JacobianInverse from ufl.classes import (Abs, Argument, CellOrientation, - Expr, FloatValue, Division, + Expr, FloatValue, Division, ReferenceValue, Product, ScalarValue, Sqrt, Zero, CellVolume, FacetArea) from ufl.utils.sorting import sorted_by_count @@ -37,6 +38,36 @@ preserve_geometry_types = (CellVolume, FacetArea) +class InterpolateMapper(MultiFunction): + """Represent interpolation in the target element's reference frame.""" + + expr = MultiFunction.reuse_if_untouched + + def interpolate(self, o: ufl.Interpolate, operand: Expr) -> Expr: + dual_arg, _ = o.argument_slots() + domain = extract_unique_domain(o) + element = o.ufl_function_space().ufl_element() + operand = apply_mapping(operand, element, domain) + expr = o._ufl_expr_reconstruct_(operand, v=dual_arg) + return element.pullback.apply(ReferenceValue(expr), domain) + + +def lower_form_interpolations(form: ufl.Form) -> ufl.Form: + """Represent interpolation nodes in a form in reference space. + + Parameters + ---------- + form : ufl.Form + Form containing interpolation nodes. + + Returns + ------- + ufl.Form + Form with target-element mappings made explicit. + """ + return map_integrand_dags(InterpolateMapper(), form) + + def compute_form_data(form, do_apply_function_pullbacks=True, do_apply_integral_scaling=True, From 241a209fb2f15305b06a54b06f43683f39580429 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 16 Jul 2026 15:16:45 +0100 Subject: [PATCH 02/10] FormAssembler: handle same-mesh interpolation --- firedrake/assemble.py | 200 +++++++++++++++++++---- firedrake/interpolation.py | 143 +++++++++++++++- pyop2/parloop.py | 12 +- tsfc/kernel_interface/firedrake_loopy.py | 10 +- tsfc/ufl_utils.py | 16 ++ 5 files changed, 333 insertions(+), 48 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 67a32be63f..0e32a3d0b5 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -19,7 +19,7 @@ from firedrake import (extrusion_utils as eutils, parameters, solving, tsfc_interface, utils) from firedrake.adjoint_utils import annotate_assemble -from firedrake.ufl_expr import extract_domains +from firedrake.ufl_expr import extract_domains as _extract_domains from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit from firedrake.matrix import MatrixBase, Matrix, ImplicitMatrix from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace @@ -43,6 +43,17 @@ """Entry used in form cache to try and reuse assemblers where possible.""" +def extract_domains(expression): + """Extract domains, including interpolation argument slots.""" + domains = list(_extract_domains(expression)) + if isinstance(expression, ufl.Interpolate): + dual_arg, _ = expression.argument_slots() + for domain in _extract_domains(dual_arg): + if domain not in domains: + domains.append(domain) + return tuple(domains) + + @PETSc.Log.EventDecorator() @annotate_assemble def assemble(expr, *args, **kwargs): @@ -904,8 +915,8 @@ def preprocess_base_form(expr, mat_type=None, form_compiler_parameters=None): expr = BaseFormAssembler.restructure_base_form_postorder(expr) # Preprocessing the form makes a new object -> current form caching mechanism # will populate `expr`'s cache which is now different than `original_expr`'s cache so we need - # to transmit the cache. All of this only holds when both are `ufl.Form` objects. - if isinstance(original_expr, ufl.form.Form) and isinstance(expr, ufl.form.Form): + # to transmit the cache. Both objects must support the assembler cache. + if isinstance(original_expr, (ufl.form.Form, ufl.Interpolate)) and isinstance(expr, (ufl.form.Form, ufl.Interpolate)): expr._cache = original_expr._cache return expr @@ -960,8 +971,8 @@ class FormAssembler(AbstractFormAssembler): def __new__(cls, *args, **kwargs): form = args[0] - if not isinstance(form, (ufl.Form, slate.TensorBase)): - raise TypeError(f"The first positional argument must be of ufl.Form or slate.TensorBase: got {type(form)} ({form})") + if not isinstance(form, (ufl.Form, ufl.Interpolate, slate.TensorBase)): + raise TypeError(f"The first positional argument must be of ufl.Form, ufl.Interpolate, or slate.TensorBase: got {type(form)} ({form})") # It is expensive to construct new assemblers because extracting the data # from the form is slow. Since all of the data structures in the assembler # are persistent apart from the output tensor, we stash the assembler on the @@ -982,9 +993,10 @@ def __new__(cls, *args, **kwargs): self = super().__new__(cls) self._initialised = False self.__init__(*args, **kwargs) - if _FORM_CACHE_KEY not in form._cache: - form._cache[_FORM_CACHE_KEY] = {} - form._cache[_FORM_CACHE_KEY][key] = self + if key is not None: + if _FORM_CACHE_KEY not in form._cache: + form._cache[_FORM_CACHE_KEY] = {} + form._cache[_FORM_CACHE_KEY][key] = self return self @classmethod @@ -1122,6 +1134,8 @@ def local_kernels(self): self._form, "form", diagonal=self.diagonal, parameters=self._form_compiler_params ) + elif isinstance(self._form, ufl.Interpolate): + kernels = self._compile_interpolate() elif isinstance(self._form, slate.TensorBase): kernels = slac.compile_expression( self._form, @@ -1133,6 +1147,91 @@ def local_kernels(self): (k, subdomain_id) for k in kernels for subdomain_id in k.kinfo.subdomain_id ) + def _compile_interpolate(self): + """Compile an interpolation as a form kernel.""" + from firedrake.interpolation import compile_expression + + all_meshes = extract_domains(self._form) + all_constants = extract_firedrake_constants(self._form) + parameters = {"scalar_type": ScalarType} + parameters.update(self._form_compiler_params) + + dual_arg, operand = self._form.argument_slots() + target_mesh = dual_arg.function_space().mesh().unique() + source_mesh = (ufl.domain.extract_unique_domain(operand) or target_mesh).unique() + target_element = dual_arg.function_space().dual().ufl_element() + expression_kernel = compile_expression( + target_mesh.comm, self._form, target_element, + domain=source_mesh, parameters=parameters, + ) + + arguments = list(expression_kernel.arguments) + active_domain = all_meshes.index(source_mesh) + coefficient_offset = ( + 1 + + expression_kernel.oriented + + expression_kernel.needs_cell_sizes + ) + if expression_kernel.needs_external_coords: + coordinate_arg = arguments[coefficient_offset] + arguments[coefficient_offset] = kernel_args.CoordinatesKernelArg( + coordinate_arg.loopy_arg + ) + + active_domain_numbers = tsfc_interface.ActiveDomainNumbers( + coordinates=( + (active_domain,) + if expression_kernel.needs_external_coords + else () + ), + cell_orientations=( + (active_domain,) if expression_kernel.oriented else () + ), + cell_sizes=( + (active_domain,) if expression_kernel.needs_cell_sizes else () + ), + exterior_facets=(), + interior_facets=(), + orientations_cell=(), + orientations_exterior_facet=(), + orientations_interior_facet=(), + ) + + coefficient_numbers = tuple( + (i, None) for i in expression_kernel.coefficient_numbers + ) + constants = extract_firedrake_constants(self._form) + constant_numbers = tuple( + all_constants.index(constant) for constant in constants + ) + access = self._form.options.access or op2.WRITE + kernel = tsfc_interface.as_pyop2_local_kernel( + expression_kernel.ast, + expression_kernel.name, + len(arguments), + access=access, + flop_count=expression_kernel.flop_count, + events=(expression_kernel.event,), + ) + kinfo = tsfc_interface.KernelInfo( + kernel=kernel, + integral_type="cell", + subdomain_id=("everywhere",), + domain_number=all_meshes.index(target_mesh), + active_domain_numbers=active_domain_numbers, + coefficient_numbers=coefficient_numbers, + constant_numbers=constant_numbers, + needs_cell_facets=False, + pass_layer_arg=False, + arguments=tuple(arguments), + events=(expression_kernel.event,), + ) + return ( + tsfc_interface.SplitKernel( + (None,) * len(self._form.arguments()), kinfo + ), + ) + @property @abc.abstractmethod def diagonal(self): @@ -1290,6 +1389,11 @@ def _as_pyop2_type(tensor, indices=None): return tensor.dat def execute_parloops(self, tensor): + if isinstance(self._form, ufl.Interpolate) and self._form.options.access is not op2.INC: + for parloop in self.parloops(tensor): + parloop() + return + # We are repeatedly incrementing into the same Dat so intermediate halo exchanges # can be skipped. with tensor.dat.frozen_halo(op2.INC): @@ -1305,7 +1409,7 @@ def result(self, tensor): def TwoFormAssembler(form, *args, **kwargs): - assert isinstance(form, (ufl.form.Form, slate.TensorBase)) + assert isinstance(form, (ufl.form.Form, ufl.Interpolate, slate.TensorBase)) mat_type = kwargs.pop('mat_type', None) sub_mat_type = kwargs.pop('sub_mat_type', None) mat_type, sub_mat_type = _get_mat_type(mat_type, sub_mat_type, form.arguments()) @@ -1494,6 +1598,8 @@ def _all_assemblers(self): def _apply_bc(self, tensor, bc, u=None): assert u is None + if isinstance(self._form, ufl.Interpolate): + return op2tensor = tensor.M spaces = tuple(a.function_space() for a in tensor.a.arguments()) V = bc.function_space() @@ -1630,22 +1736,27 @@ def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomai if isinstance(form, ufl.Form): sig = form.signature() + elif isinstance(form, ufl.Interpolate): + sig = hash(form) elif isinstance(form, slate.TensorBase): sig = form.expression_hash # The form signature does not store this information. This should be accessible from # the UFL so we don't need this nasty hack. subdomain_key = [] - for val in form.subdomain_data().values(): - for k, v in val.items(): - for i, vi in enumerate(v): - if vi is not None: - extruded = vi._extruded - constant_layers = extruded and vi.constant_layers - subset = isinstance(vi, op2.Subset) - subdomain_key.append((k, i, extruded, constant_layers, subset)) - else: - subdomain_key.append((k, i)) + if isinstance(form, ufl.Interpolate): + subdomain_key.append(("cell", 0, form.options.subset is not None)) + else: + for val in form.subdomain_data().values(): + for k, v in val.items(): + for i, vi in enumerate(v): + if vi is not None: + extruded = vi._extruded + constant_layers = extruded and vi.constant_layers + subset = isinstance(vi, op2.Subset) + subdomain_key.append((k, i, extruded, constant_layers, subset)) + else: + subdomain_key.append((k, i)) return (domain_ids + (sig, subdomain_id) @@ -1747,6 +1858,8 @@ def _mesh(self): @cached_property def _needs_subset(self): + if isinstance(self._form, ufl.Interpolate): + return self._form.options.subset is not None subdomain_data = self._form.subdomain_data()[self._mesh] if not all(sd is None for sd in subdomain_data.get(self._integral_type, [None])): return True @@ -1789,7 +1902,7 @@ def _make_mat_global_kernel_arg(self, Vrow, Vcol): if any(isinstance(e, finat.EnrichedElement) and e.is_mixed for e in {relem, celem}): subargs = tuple(self._make_mat_global_kernel_arg(Vrow_sub, Vcol_sub) for Vrow_sub, Vcol_sub in product(Vrow, Vcol)) - shape = len(relem.elements), len(celem.elements) + shape = len(Vrow), len(Vcol) return op2.MixedMatKernelArg(subargs, shape) else: rmap_arg, cmap_arg = (V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg for V in [Vrow, Vcol]) @@ -2052,18 +2165,28 @@ def get_indicess(self): def _filter_bcs(self, row, col): assert len(self._form.arguments()) == 2 and not self._diagonal - if len(self.test_function_space) > 1: - bcrow = tuple(bc for bc in self._bcs - if bc.function_space_index() == row) - else: - bcrow = self._bcs - if len(self.trial_function_space) > 1: - bccol = tuple(bc for bc in self._bcs - if bc.function_space_index() == col - and isinstance(bc, DirichletBC)) - else: - bccol = tuple(bc for bc in self._bcs if isinstance(bc, DirichletBC)) + def matches_space(bc, V): + V = V.dual() if isinstance(V, FiredrakeDualSpace) else V + bc_space = bc.function_space() + if V.parent is None: + return bc_space.parent is None and bc_space == V + return bc_space.parent == V.parent and bc_space.index == V.index + + test_space = self.test_function_space + if len(test_space) > 1: + test_space = test_space[row] + bcrow = tuple( + bc for bc in self._bcs if matches_space(bc, test_space) + ) + + trial_space = self.trial_function_space + if len(trial_space) > 1: + trial_space = trial_space[col] + bccol = tuple( + bc for bc in self._bcs + if isinstance(bc, DirichletBC) and matches_space(bc, trial_space) + ) return bcrow, bccol def needs_unrolling(self): @@ -2095,7 +2218,6 @@ def collect_lgmaps(self): if len(self._form.arguments()) == 2 and not self._diagonal: if not self._bcs: return None - if any(i is not None for i in self._local_knl.indices): i, j = self._local_knl.indices row_bcs, col_bcs = self._filter_bcs(i, j) @@ -2141,6 +2263,8 @@ def _mesh(self): @cached_property def _iterset(self): + if isinstance(self._form, ufl.Interpolate): + return self._form.options.subset or self._mesh.cell_set try: subdomain_data = self._form.subdomain_data()[self._mesh][self._integral_type] except KeyError: @@ -2342,10 +2466,16 @@ def iter_active_cell_sizes(form, kinfo): @staticmethod def iter_active_coefficients(form, kinfo): """Yield the form coefficients referenced in ``kinfo``.""" - all_coefficients = form.coefficients() + if isinstance(form, ufl.Interpolate): + all_coefficients = ufl.algorithms.extract_coefficients(form) + else: + all_coefficients = form.coefficients() for idx, subidxs in kinfo.coefficient_numbers: - for subidx in subidxs: - yield all_coefficients[idx].subfunctions[subidx] + if subidxs is None: + yield all_coefficients[idx] + else: + for subidx in subidxs: + yield all_coefficients[idx].subfunctions[subidx] @staticmethod def iter_constants(form, kinfo): diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 58823f5451..0b0b58f829 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -8,7 +8,7 @@ from dataclasses import asdict, dataclass from numbers import Number -from ufl.algorithms import extract_arguments, replace +from ufl.algorithms import extract_arguments, extract_coefficients, replace from ufl.domain import extract_unique_domain from ufl.classes import Expr from ufl.duals import is_dual @@ -163,8 +163,6 @@ def _interpolator(self): """ arguments = self.arguments() has_mixed_arguments = any(len(arg.function_space()) > 1 for arg in arguments) - if len(arguments) == 2 and has_mixed_arguments: - return MixedInterpolator(self) operand, = self.ufl_operands target_mesh = self.target_space.mesh() @@ -707,6 +705,68 @@ def __init__(self, expr, source_mesh, target_mesh): # Default access for forward 1-form or 2-form (forward and adjoint) self.access = op2.WRITE + @cached_property + def _form_interpolate(self): + options = asdict(self.ufl_interpolate.options) + options.update(subset=self.subset, access=self.access) + return self.ufl_interpolate._ufl_expr_reconstruct_( + self.operand, v=self.dual_arg, **options + ) + + @property + def _needs_adjoint_weighting(self): + return ( + isinstance(self.dual_arg, Cofunction) + and any(not V.finat_element.is_dg() for V in self.target_space) + ) + + @cached_property + def _weighted_dual_arg(self): + return Function(self.dual_arg.function_space()) + + @cached_property + def _adjoint_weight(self): + W = self.dual_arg.function_space() + weight = W.make_dat() + if len(W) > 1: + spaces_and_weights = zip(W, weight) + else: + spaces_and_weights = ((W, weight),) + + source_mesh = self.source_mesh.unique() + target_mesh = self.target_mesh.unique() + for i, (V, component_weight) in enumerate(spaces_and_weights): + node_map = get_interp_node_map(source_mesh, target_mesh, V) + size = V.finat_element.space_dimension() * V.block_size + kernel_code = f""" + void multiplicity_{i}(PetscScalar *restrict w) {{ + for (PetscInt i=0; i<{size}; i++) w[i] += 1; + }}""" + kernel = op2.Kernel(kernel_code, f"multiplicity_{i}") + op2.par_loop( + kernel, target_mesh.cell_set, + component_weight(op2.INC, node_map), + ) + with weight.vec as weight_vec: + weight_vec.reciprocal() + return weight + + @cached_property + def _assembler_form(self): + if self._needs_adjoint_weighting: + return self._form_interpolate._ufl_expr_reconstruct_( + self.operand, v=self._weighted_dual_arg + ) + return self._form_interpolate + + def _update_weighted_dual_arg(self): + self.dual_arg.dat.copy(self._weighted_dual_arg.dat) + with ( + self._adjoint_weight.vec_ro as weight, + self._weighted_dual_arg.dat.vec as dual, + ): + dual.pointwiseMult(dual, weight) + def _get_tensor(self, mat_type: Literal["aij", "baij"]) -> op2.Mat | Function | Cofunction: """Return a suitable tensor to interpolate into. @@ -768,7 +828,82 @@ def _get_monolithic_sparsity(self, mat_type: Literal["aij", "baij"]) -> op2.Spar block_sparse=(mat_type == "baij")) return sparsity + def _get_form_assembler(self, bcs=None, mat_type=None, sub_mat_type=None): + """Return the form assembler for this interpolation rank.""" + from firedrake.assemble import ( + OneFormAssembler, TwoFormAssembler, ZeroFormAssembler, + ) + + if self.rank == 0: + return ZeroFormAssembler(self._assembler_form) + elif self.rank == 1: + return OneFormAssembler( + self._assembler_form, + bcs=bcs, + needs_zeroing=self.access is op2.INC, + ) + elif self.rank == 2: + return TwoFormAssembler( + self._assembler_form, + bcs=bcs, + mat_type=mat_type, + sub_mat_type=sub_mat_type, + needs_zeroing=True, + allocation_integral_types=("cell",), + ) + else: + raise ValueError( + f"Cannot interpolate an expression with {self.rank} arguments" + ) + def _get_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None): + if self.source_mesh.unique() is not self.target_mesh.unique(): + return self._get_transmesh_callable( + tensor=tensor, + bcs=bcs, + mat_type=mat_type, + sub_mat_type=sub_mat_type, + ) + + assembler = self._get_form_assembler( + bcs=bcs, mat_type=mat_type, sub_mat_type=sub_mat_type, + ) + copyout = () + inputs = tuple( + coefficient for coefficient in extract_coefficients(self._assembler_form) + if isinstance(coefficient, Function | Cofunction) + ) + if isinstance(self.dual_arg, Cofunction): + inputs += (self.dual_arg,) + if ( + isinstance(tensor, Function | Cofunction) + and any(set(tensor.dat).intersection(set(input_.dat)) + for input_ in inputs) + ): + output = tensor + tensor = assembler.allocate() + copyout = (partial(tensor.dat.copy, output.dat),) + elif tensor is None and self.access in {op2.MIN, op2.MAX}: + tensor = assembler.allocate() + finfo = numpy.finfo(tensor.dat.dtype) + value = finfo.max if self.access == op2.MIN else finfo.min + tensor.assign(Constant(value)) + + assembler_tensor = None if self.rank == 2 else tensor + + def callable(): + if self._needs_adjoint_weighting: + self._update_weighted_dual_arg() + result = assembler.assemble(tensor=assembler_tensor) + for copy in copyout: + copy() + if isinstance(result, MatrixBase): + return result.petscmat + return output if copyout else result + + return callable + + def _get_transmesh_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None): mat_type = mat_type or "aij" if (isinstance(tensor, Cofunction) and isinstance(self.dual_arg, Cofunction)) and set(tensor.dat).intersection(set(self.dual_arg.dat)): # adjoint one-form case: we need an empty tensor, so if it shares dats with @@ -827,7 +962,7 @@ def callable() -> Function | Cofunction | PETSc.Mat | Number: @property def _allowed_mat_types(self): - return {"aij", "baij", "matfree", None} + return {"aij", "baij", "nest", "matfree", None} class VomOntoVomInterpolator(SameMeshInterpolator): diff --git a/pyop2/parloop.py b/pyop2/parloop.py index 6392bc889f..6773e6ccf3 100644 --- a/pyop2/parloop.py +++ b/pyop2/parloop.py @@ -121,6 +121,12 @@ def _kernel_args_(self): @property def map_kernel_args(self): rmap, cmap = self.maps + if isinstance(rmap, MixedMap): + return tuple(itertools.chain.from_iterable( + itertools.chain(rmap_._kernel_args_, cmap_._kernel_args_) + for rmap_ in rmap.split + for cmap_ in cmap.split + )) return tuple(itertools.chain(rmap._kernel_args_, cmap._kernel_args_)) @@ -143,7 +149,11 @@ def _kernel_args_(self): @property def map_kernel_args(self): rmap, cmap = self.maps - return tuple(itertools.chain(rmap._kernel_args_, cmap._kernel_args_)) + return tuple(itertools.chain.from_iterable( + itertools.chain(rmap_._kernel_args_, cmap_._kernel_args_) + for rmap_ in rmap.split + for cmap_ in cmap.split + )) @dataclass diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index cc8fd7a61e..15567bbc9e 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -4,7 +4,7 @@ from ufl import Coefficient, FunctionSpace from ufl.domain import MeshSequence -from finat.ufl import MixedElement as ufl_MixedElement, FiniteElement +from finat.ufl import FiniteElement import gem from gem.flop_count import count_flops @@ -206,13 +206,7 @@ def set_coefficients(self, coefficients): self.coefficient_split = {} for i, coefficient in enumerate(coefficients): - if type(coefficient.ufl_element()) == ufl_MixedElement: - subcoeffs = coefficient.subfunctions # Firedrake-specific - self.coefficient_split[coefficient] = subcoeffs - for j, subcoeff in enumerate(subcoeffs): - self._coefficient(subcoeff, f"w_{i}_{j}") - else: - self._coefficient(coefficient, f"w_{i}") + self._coefficient(coefficient, f"w_{i}") def set_constants(self, constants): for i, const in enumerate(constants): diff --git a/tsfc/ufl_utils.py b/tsfc/ufl_utils.py index e5f7af13ea..1740248d51 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -426,6 +426,22 @@ def apply_mapping(expression, element, domain): mapping = element.mapping().lower() if mapping == "identity": rexpression = expression + elif isinstance(element.pullback, ufl.MixedPullback): + flat = [expression[index] for index in numpy.ndindex(expression.ufl_shape)] + reference_components = [] + offset = 0 + for subelement, subdomain in zip(element.sub_elements, mesh.iterable_like(element)): + physical_shape = subelement.pullback.physical_value_shape(subelement, subdomain) + size = int(numpy.prod(physical_shape, dtype=int)) + piece = as_tensor(numpy.asarray(flat[offset:offset + size]).reshape(physical_shape)) + mapped = apply_mapping(piece, subelement, subdomain) + reference_components.extend( + mapped[index] for index in numpy.ndindex(mapped.ufl_shape) + ) + offset += size + rexpression = as_tensor( + numpy.asarray(reference_components).reshape(element.reference_value_shape) + ) elif mapping == "covariant piola": J = Jacobian(mesh) *k, i, j = indices(len(expression.ufl_shape) + 1) From 0845c66fe69657d832ec53de94e5e559ef2339a2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 27 Aug 2026 16:43:40 +0100 Subject: [PATCH 03/10] Unify Interpolate and Form assembly Route a same-mesh Interpolate through the ordinary Form assembly path, so one code generator serves both: the kernel builder compiles an Interpolate the way it compiles an integrand, and the assembler allocates and applies boundary conditions for it the way it does for a Form. --- firedrake/assemble.py | 299 ++++++++--------- firedrake/bcs.py | 24 ++ firedrake/formmanipulation.py | 34 +- firedrake/interpolation.py | 313 +----------------- firedrake/tsfc_interface.py | 16 +- firedrake/ufl_expr.py | 2 + pyop2/parloop.py | 12 +- .../firedrake/regression/test_interp_dual.py | 71 ++++ .../submesh/test_submesh_interpolate.py | 40 +++ tsfc/driver.py | 142 +++++++- tsfc/fem.py | 73 +++- tsfc/kernel_interface/__init__.py | 4 + tsfc/kernel_interface/common.py | 72 +++- tsfc/kernel_interface/firedrake_loopy.py | 11 +- tsfc/spectral.py | 4 +- tsfc/ufl_utils.py | 7 +- 16 files changed, 608 insertions(+), 516 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 0e32a3d0b5..08df0ce215 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -19,12 +19,13 @@ from firedrake import (extrusion_utils as eutils, parameters, solving, tsfc_interface, utils) from firedrake.adjoint_utils import annotate_assemble -from firedrake.ufl_expr import extract_domains as _extract_domains +from firedrake.ufl_expr import extract_domains from firedrake.bcs import DirichletBC, EquationBC, EquationBCSplit from firedrake.matrix import MatrixBase, Matrix, ImplicitMatrix +from firedrake.mesh import VertexOnlyMeshTopology from firedrake.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace from firedrake.functionspacedata import entity_dofs_key, entity_permutations_key -from firedrake.interpolation import get_interpolator +from firedrake.interpolation import get_interp_node_map, get_interpolator from firedrake.petsc import PETSc from firedrake.slate import slac, slate from firedrake.slate.slac.kernel_builder import CellFacetKernelArg, LayerCountKernelArg @@ -43,17 +44,6 @@ """Entry used in form cache to try and reuse assemblers where possible.""" -def extract_domains(expression): - """Extract domains, including interpolation argument slots.""" - domains = list(_extract_domains(expression)) - if isinstance(expression, ufl.Interpolate): - dual_arg, _ = expression.argument_slots() - for domain in _extract_domains(dual_arg): - if domain not in domains: - domains.append(domain) - return tuple(domains) - - @PETSc.Log.EventDecorator() @annotate_assemble def assemble(expr, *args, **kwargs): @@ -183,14 +173,22 @@ def get_assembler(form, *args, **kwargs): form = BaseFormAssembler.preprocess_base_form(form, mat_type=mat_type, form_compiler_parameters=fc_params) base_form_operands = BaseFormAssembler.base_form_operands(form) can_compile = not base_form_operands - if isinstance(form, ufl.form.Form) and len(form.ufl_domains()) == 1: - domain, = form.ufl_domains() - can_compile = all( - isinstance(op, ufl.Interpolate) - and op.ufl_function_space().ufl_domain() == domain - and set(extract_domains(op)) <= {domain} - for op in base_form_operands - ) + if isinstance(form, ufl.form.Form): + can_compile = True + for integral in form.integrals(): + valid_domains = set(integral.extra_domain_integral_type_map()) + valid_domains.add(integral.ufl_domain()) + for op in ufl.algorithms.extract_base_form_operators( + integral.integrand() + ): + if ( + not isinstance(op, ufl.Interpolate) + or not set(extract_domains(op)) <= valid_domains + ): + can_compile = False + break + if not can_compile: + break if isinstance(form, (ufl.form.Form, slate.TensorBase)) and can_compile: diagonal = kwargs.pop('diagonal', False) @@ -1035,9 +1033,11 @@ class ParloopFormAssembler(FormAssembler): Should ``tensor`` be zeroed before assembling? """ - def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True): + def __init__(self, form, bcs=None, form_compiler_parameters=None, + needs_zeroing=True, access=op2.INC): super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) self._needs_zeroing = needs_zeroing + self._access = access def assemble(self, tensor=None, current_state=None): """Assemble the form. @@ -1129,13 +1129,12 @@ def local_kernels(self): each possible combination. """ - if isinstance(self._form, ufl.Form): + if isinstance(self._form, (ufl.Form, ufl.Interpolate)): kernels = tsfc_interface.compile_form( self._form, "form", diagonal=self.diagonal, - parameters=self._form_compiler_params + parameters=self._form_compiler_params, + access=self._access, ) - elif isinstance(self._form, ufl.Interpolate): - kernels = self._compile_interpolate() elif isinstance(self._form, slate.TensorBase): kernels = slac.compile_expression( self._form, @@ -1147,91 +1146,6 @@ def local_kernels(self): (k, subdomain_id) for k in kernels for subdomain_id in k.kinfo.subdomain_id ) - def _compile_interpolate(self): - """Compile an interpolation as a form kernel.""" - from firedrake.interpolation import compile_expression - - all_meshes = extract_domains(self._form) - all_constants = extract_firedrake_constants(self._form) - parameters = {"scalar_type": ScalarType} - parameters.update(self._form_compiler_params) - - dual_arg, operand = self._form.argument_slots() - target_mesh = dual_arg.function_space().mesh().unique() - source_mesh = (ufl.domain.extract_unique_domain(operand) or target_mesh).unique() - target_element = dual_arg.function_space().dual().ufl_element() - expression_kernel = compile_expression( - target_mesh.comm, self._form, target_element, - domain=source_mesh, parameters=parameters, - ) - - arguments = list(expression_kernel.arguments) - active_domain = all_meshes.index(source_mesh) - coefficient_offset = ( - 1 - + expression_kernel.oriented - + expression_kernel.needs_cell_sizes - ) - if expression_kernel.needs_external_coords: - coordinate_arg = arguments[coefficient_offset] - arguments[coefficient_offset] = kernel_args.CoordinatesKernelArg( - coordinate_arg.loopy_arg - ) - - active_domain_numbers = tsfc_interface.ActiveDomainNumbers( - coordinates=( - (active_domain,) - if expression_kernel.needs_external_coords - else () - ), - cell_orientations=( - (active_domain,) if expression_kernel.oriented else () - ), - cell_sizes=( - (active_domain,) if expression_kernel.needs_cell_sizes else () - ), - exterior_facets=(), - interior_facets=(), - orientations_cell=(), - orientations_exterior_facet=(), - orientations_interior_facet=(), - ) - - coefficient_numbers = tuple( - (i, None) for i in expression_kernel.coefficient_numbers - ) - constants = extract_firedrake_constants(self._form) - constant_numbers = tuple( - all_constants.index(constant) for constant in constants - ) - access = self._form.options.access or op2.WRITE - kernel = tsfc_interface.as_pyop2_local_kernel( - expression_kernel.ast, - expression_kernel.name, - len(arguments), - access=access, - flop_count=expression_kernel.flop_count, - events=(expression_kernel.event,), - ) - kinfo = tsfc_interface.KernelInfo( - kernel=kernel, - integral_type="cell", - subdomain_id=("everywhere",), - domain_number=all_meshes.index(target_mesh), - active_domain_numbers=active_domain_numbers, - coefficient_numbers=coefficient_numbers, - constant_numbers=constant_numbers, - needs_cell_facets=False, - pass_layer_arg=False, - arguments=tuple(arguments), - events=(expression_kernel.event,), - ) - return ( - tsfc_interface.SplitKernel( - (None,) * len(self._form.arguments()), kinfo - ), - ) - @property @abc.abstractmethod def diagonal(self): @@ -1320,14 +1234,21 @@ class OneFormAssembler(ParloopFormAssembler): @classmethod def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=True, diagonal=False, weight=1.0): + zero_bc_nodes=True, diagonal=False, weight=1.0, access=op2.INC): bcs = solving._extract_bcs(bcs) - return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal, weight + return (tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, + zero_bc_nodes, diagonal, weight, access) @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=True, diagonal=False, weight=1.0): - super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + zero_bc_nodes=True, diagonal=False, weight=1.0, access=op2.INC): + super().__init__( + form, + bcs=bcs, + form_compiler_parameters=form_compiler_parameters, + needs_zeroing=needs_zeroing, + access=access, + ) self._weight = weight self._diagonal = diagonal self._zero_bc_nodes = zero_bc_nodes @@ -1389,7 +1310,7 @@ def _as_pyop2_type(tensor, indices=None): return tensor.dat def execute_parloops(self, tensor): - if isinstance(self._form, ufl.Interpolate) and self._form.options.access is not op2.INC: + if self._access is not op2.INC: for parloop in self.parloops(tensor): parloop() return @@ -1480,8 +1401,14 @@ def _cache_key(cls, *args, **kwargs): @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, mat_type=None, sub_mat_type=None, options_prefix=None, appctx=None, weight=1.0, - allocation_integral_types=None): - super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + allocation_integral_types=None, access=op2.INC): + super().__init__( + form, + bcs=bcs, + form_compiler_parameters=form_compiler_parameters, + needs_zeroing=needs_zeroing, + access=access, + ) self._mat_type = mat_type self._sub_mat_type = sub_mat_type self._options_prefix = options_prefix @@ -1546,8 +1473,14 @@ def _make_maps_and_regions(self): # Make Sparsity independent of the subdomain of integration for better reusability; # subdomain_id is passed here only to determine the integration_type on the target domain # (see ``entity_node_map``). - rmap_ = test.function_space().topological[i].entity_node_map(mesh.topology, integral_type, subdomain_id, all_subdomain_ids) - cmap_ = trial.function_space().topological[j].entity_node_map(mesh.topology, integral_type, subdomain_id, all_subdomain_ids) + rmap_ = _get_entity_node_map( + mesh, test.function_space()[i], + integral_type, subdomain_id, all_subdomain_ids, + ) + cmap_ = _get_entity_node_map( + mesh, trial.function_space()[j], + integral_type, subdomain_id, all_subdomain_ids, + ) region = ExplicitMatrixAssembler._integral_type_region_map[integral_type] maps_and_regions[(i, j)][(rmap_, cmap_)].add(region) return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()] @@ -1566,8 +1499,12 @@ def _make_maps_and_regions_default(test, trial, allocation_integral_types): for i, Vrow in enumerate(test.function_space()): for j, Vcol in enumerate(trial.function_space()): mesh = Vrow.mesh() - rmap_ = Vrow.topological.entity_node_map(mesh.topology, integral_type, None, None) - cmap_ = Vcol.topological.entity_node_map(mesh.topology, integral_type, None, None) + rmap_ = _get_entity_node_map( + mesh, Vrow, integral_type, None, None + ) + cmap_ = _get_entity_node_map( + mesh, Vcol, integral_type, None, None + ) maps_and_regions[(i, j)][(rmap_, cmap_)].add(region) return {block_indices: [map_pair + (tuple(region_set), ) for map_pair, region_set in map_pair_to_region_set.items()] for block_indices, map_pair_to_region_set in maps_and_regions.items()} @@ -1598,8 +1535,6 @@ def _all_assemblers(self): def _apply_bc(self, tensor, bc, u=None): assert u is None - if isinstance(self._form, ufl.Interpolate): - return op2tensor = tensor.M spaces = tuple(a.function_space() for a in tensor.a.arguments()) V = bc.function_space() @@ -1609,7 +1544,7 @@ def _apply_bc(self, tensor, bc, u=None): index = 0 if V.index is None else V.index space = V if V.parent is None else V.parent if isinstance(bc, DirichletBC): - if not any(space == fs for fs in spaces): + if not any(bc.function_space_match(fs) for fs in spaces): raise TypeError("bc space does not match the test or trial function space") if spaces[0] != spaces[1]: # Not on a diagonal block, we cannot set diagonal entries @@ -1734,29 +1669,24 @@ def _global_kernel_cache_key(form, local_knl, subdomain_id, all_integer_subdomai all_meshes = extract_domains(form) domain_ids = tuple(mesh.ufl_id() for mesh in all_meshes) - if isinstance(form, ufl.Form): + if isinstance(form, (ufl.Form, ufl.Interpolate)): sig = form.signature() - elif isinstance(form, ufl.Interpolate): - sig = hash(form) elif isinstance(form, slate.TensorBase): sig = form.expression_hash # The form signature does not store this information. This should be accessible from # the UFL so we don't need this nasty hack. subdomain_key = [] - if isinstance(form, ufl.Interpolate): - subdomain_key.append(("cell", 0, form.options.subset is not None)) - else: - for val in form.subdomain_data().values(): - for k, v in val.items(): - for i, vi in enumerate(v): - if vi is not None: - extruded = vi._extruded - constant_layers = extruded and vi.constant_layers - subset = isinstance(vi, op2.Subset) - subdomain_key.append((k, i, extruded, constant_layers, subset)) - else: - subdomain_key.append((k, i)) + for val in form.subdomain_data().values(): + for k, v in val.items(): + for i, vi in enumerate(v): + if vi is not None: + extruded = vi._extruded + constant_layers = extruded and vi.constant_layers + subset = isinstance(vi, op2.Subset) + subdomain_key.append((k, i, extruded, constant_layers, subset)) + else: + subdomain_key.append((k, i)) return (domain_ids + (sig, subdomain_id) @@ -1770,6 +1700,18 @@ def _make_global_kernel(*args, **kwargs): return _GlobalKernelBuilder(*args, **kwargs).build() +def _get_entity_node_map( + mesh, function_space, integral_type, subdomain_id, + all_integer_subdomain_ids, +): + if isinstance(mesh.topology, VertexOnlyMeshTopology): + return get_interp_node_map(function_space.mesh(), mesh, function_space) + return function_space.topological.entity_node_map( + mesh.topology, integral_type, subdomain_id, + all_integer_subdomain_ids, + ) + + class _GlobalKernelBuilder: """Class that builds a :class:`op2.GlobalKernel`. @@ -1858,8 +1800,6 @@ def _mesh(self): @cached_property def _needs_subset(self): - if isinstance(self._form, ufl.Interpolate): - return self._form.options.subset is not None subdomain_data = self._form.subdomain_data()[self._mesh] if not all(sd is None for sd in subdomain_data.get(self._integral_type, [None])): return True @@ -1887,7 +1827,11 @@ def _get_dim(self, finat_element): def _make_dat_global_kernel_arg(self, V, index=None): finat_element = create_element(V.ufl_element()) - map_arg = V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg + map_ = _get_entity_node_map( + self._mesh, V, self._integral_type, + self._subdomain_id, self._all_integer_subdomain_ids, + ) + map_arg = map_._global_kernel_arg if isinstance(finat_element, finat.EnrichedElement) and finat_element.is_mixed: assert index is None subargs = tuple(self._make_dat_global_kernel_arg(Vsub, index=index) @@ -1902,10 +1846,17 @@ def _make_mat_global_kernel_arg(self, Vrow, Vcol): if any(isinstance(e, finat.EnrichedElement) and e.is_mixed for e in {relem, celem}): subargs = tuple(self._make_mat_global_kernel_arg(Vrow_sub, Vcol_sub) for Vrow_sub, Vcol_sub in product(Vrow, Vcol)) - shape = len(Vrow), len(Vcol) + shape = len(relem.elements), len(celem.elements) return op2.MixedMatKernelArg(subargs, shape) else: - rmap_arg, cmap_arg = (V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids)._global_kernel_arg for V in [Vrow, Vcol]) + rmap, cmap = ( + _get_entity_node_map( + self._mesh, V, self._integral_type, + self._subdomain_id, self._all_integer_subdomain_ids, + ) + for V in (Vrow, Vcol) + ) + rmap_arg, cmap_arg = rmap._global_kernel_arg, cmap._global_kernel_arg # PyOP2 matrix objects have scalar dims so we flatten them here rdim = numpy.prod(self._get_dim(relem), dtype=int) cdim = numpy.prod(self._get_dim(celem), dtype=int) @@ -2007,6 +1958,18 @@ def _as_global_kernel_arg_constant(_, self): return op2.GlobalKernelArg((value_size,)) +@_as_global_kernel_arg.register(kernel_args.TabulationKernelArg) +def _as_global_kernel_arg_tabulation(arg, self): + if ( + arg.loopy_arg.name != "rt_X" + or not isinstance(self._mesh.topology, VertexOnlyMeshTopology) + ): + raise NotImplementedError("Unknown runtime tabulation argument") + return self._make_dat_global_kernel_arg( + self._mesh.reference_coordinates.function_space() + ) + + @_as_global_kernel_arg.register(kernel_args.ExteriorFacetKernelArg) def _as_global_kernel_arg_exterior_facet(_, self): mesh = next(self._active_exterior_facets) @@ -2166,18 +2129,11 @@ def get_indicess(self): def _filter_bcs(self, row, col): assert len(self._form.arguments()) == 2 and not self._diagonal - def matches_space(bc, V): - V = V.dual() if isinstance(V, FiredrakeDualSpace) else V - bc_space = bc.function_space() - if V.parent is None: - return bc_space.parent is None and bc_space == V - return bc_space.parent == V.parent and bc_space.index == V.index - test_space = self.test_function_space if len(test_space) > 1: test_space = test_space[row] bcrow = tuple( - bc for bc in self._bcs if matches_space(bc, test_space) + bc for bc in self._bcs if bc.function_space_match(test_space) ) trial_space = self.trial_function_space @@ -2185,7 +2141,8 @@ def matches_space(bc, V): trial_space = trial_space[col] bccol = tuple( bc for bc in self._bcs - if isinstance(bc, DirichletBC) and matches_space(bc, trial_space) + if isinstance(bc, DirichletBC) + and bc.function_space_match(trial_space) ) return bcrow, bccol @@ -2218,6 +2175,7 @@ def collect_lgmaps(self): if len(self._form.arguments()) == 2 and not self._diagonal: if not self._bcs: return None + if any(i is not None for i in self._local_knl.indices): i, j = self._local_knl.indices row_bcs, col_bcs = self._filter_bcs(i, j) @@ -2263,8 +2221,6 @@ def _mesh(self): @cached_property def _iterset(self): - if isinstance(self._form, ufl.Interpolate): - return self._form.options.subset or self._mesh.cell_set try: subdomain_data = self._form.subdomain_data()[self._mesh][self._integral_type] except KeyError: @@ -2288,7 +2244,10 @@ def _iterset(self): def _get_map(self, V): """Return the appropriate PyOP2 map for a given function space.""" assert isinstance(V, (WithGeometry, FiredrakeDualSpace, FunctionSpace)) - return V.topological.entity_node_map(self._mesh.topology, self._integral_type, self._subdomain_id, self._all_integer_subdomain_ids) + return _get_entity_node_map( + self._mesh, V, self._integral_type, + self._subdomain_id, self._all_integer_subdomain_ids, + ) def _as_parloop_arg(self, tsfc_arg): """Return a :class:`op2.ParloopArg` corresponding to the provided @@ -2368,6 +2327,18 @@ def _as_parloop_arg_constant(arg, self): return op2.GlobalParloopArg(const.dat) +@_as_parloop_arg.register(kernel_args.TabulationKernelArg) +def _as_parloop_arg_tabulation(arg, self): + if ( + arg.loopy_arg.name != "rt_X" + or not isinstance(self._mesh.topology, VertexOnlyMeshTopology) + ): + raise NotImplementedError("Unknown runtime tabulation argument") + reference_coordinates = self._mesh.reference_coordinates + map_ = self._get_map(reference_coordinates.function_space()) + return op2.DatParloopArg(reference_coordinates.dat, map_) + + @_as_parloop_arg.register(kernel_args.ExteriorFacetKernelArg) def _as_parloop_arg_exterior_facet(_, self): mesh = next(self._active_exterior_facets) @@ -2466,16 +2437,10 @@ def iter_active_cell_sizes(form, kinfo): @staticmethod def iter_active_coefficients(form, kinfo): """Yield the form coefficients referenced in ``kinfo``.""" - if isinstance(form, ufl.Interpolate): - all_coefficients = ufl.algorithms.extract_coefficients(form) - else: - all_coefficients = form.coefficients() + all_coefficients = form.coefficients() for idx, subidxs in kinfo.coefficient_numbers: - if subidxs is None: - yield all_coefficients[idx] - else: - for subidx in subidxs: - yield all_coefficients[idx].subfunctions[subidx] + for subidx in subidxs: + yield all_coefficients[idx].subfunctions[subidx] @staticmethod def iter_constants(form, kinfo): diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 16bed26f60..9f4f6b4ec8 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -90,6 +90,30 @@ def function_space_index(self): raise RuntimeError("This function should only be called when function space is indexed") return fs.index + def function_space_match( + self, V: ufl.functionspace.AbstractFunctionSpace + ) -> bool: + """Return whether this boundary condition is defined on ``V``. + + Parameters + ---------- + V : ufl.functionspace.AbstractFunctionSpace + Function space to compare with this boundary condition's space. + + Returns + ------- + bool + Whether the spaces represent the same indexed field. + """ + if ufl.duals.is_dual(V): + V = V.dual() + bc_space = self.function_space() + if V.parent is None: + if bc_space.parent is None: + return bc_space == V + return bc_space.parent == V + return bc_space.parent == V.parent and bc_space.index == V.index + @cached_property def domain_args(self): r"""The sub_domain the BC applies to.""" diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 45c3ef5c5e..a5a184c2e8 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -3,7 +3,7 @@ import collections from ufl import as_tensor, as_vector, split -from ufl.classes import Form, Zero, FixedIndex, ListTensor, ZeroBaseForm +from ufl.classes import Form, Interpolate, Zero, FixedIndex, ListTensor, ZeroBaseForm from ufl.algorithms.map_integrands import map_integrand_dags from ufl.algorithms import expand_derivatives from ufl.corealg.map_dag import MultiFunction, map_expr_dags @@ -30,6 +30,12 @@ class ExtractSubBlock(MultiFunction): """Extract a sub-block from a form.""" + def __init__(self): + super().__init__() + self._arg_cache = {} + self.blocks = {} + self._splitting_interpolate = False + class IndexInliner(MultiFunction): """Inline fixed index of list tensors""" expr = MultiFunction.reuse_if_untouched @@ -82,6 +88,7 @@ def split(self, form, argument_indices): args = form.arguments() self._arg_cache = {} self.blocks = dict(enumerate(map(as_tuple, argument_indices))) + self._splitting_interpolate = isinstance(form, Interpolate) if len(args) == 0: # Functional can't be split return form @@ -231,7 +238,9 @@ def zero_base_form(self, o): def interpolate(self, o, operand): if isinstance(operand, Zero): - return self(ZeroBaseForm(o.arguments())) + if self._splitting_interpolate: + return self(ZeroBaseForm(o.arguments())) + return Zero(o.ufl_shape) dual_arg, _ = o.argument_slots() if len(dual_arg.arguments()) == 1 or len(dual_arg.arguments()[-1].function_space()) == 1: @@ -258,9 +267,26 @@ def interpolate(self, o, operand): operand = as_tensor(numpy.reshape(components, W.value_shape)) if isinstance(operand, Zero): - return self(ZeroBaseForm(o.arguments())) + if self._splitting_interpolate: + return self(ZeroBaseForm(o.arguments())) + return Zero(o.ufl_shape) - return o._ufl_expr_reconstruct_(operand, sub_dual_arg) + interpolation = o._ufl_expr_reconstruct_(operand, sub_dual_arg) + if self._splitting_interpolate: + return interpolation + + interpolation_components = iter( + interpolation[j] for j in numpy.ndindex(interpolation.ufl_shape) + ) if interpolation.ufl_shape else iter((interpolation,)) + components = [] + for i, Vi in enumerate(V): + if i in indices: + components.extend( + next(interpolation_components) for _ in range(Vi.value_size) + ) + else: + components.extend(Zero() for _ in range(Vi.value_size)) + return as_tensor(numpy.reshape(components, V.value_shape)) SplitForm = collections.namedtuple("SplitForm", ["indices", "form"]) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 0b0b58f829..311c329bb8 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -1,33 +1,22 @@ import numpy -import os -import tempfile import abc from functools import cached_property, partial -from typing import Hashable, Literal, Callable, Iterable +from typing import Literal, Callable, Iterable from dataclasses import asdict, dataclass from numbers import Number from ufl.algorithms import extract_arguments, extract_coefficients, replace from ufl.domain import extract_unique_domain from ufl.classes import Expr -from ufl.duals import is_dual -from ufl.constantvalue import zero, as_ufl +from ufl.constantvalue import as_ufl from ufl.form import ZeroBaseForm, BaseForm from ufl.core.interpolate import Interpolate as UFLInterpolate from pyop2 import op2 -from pyop2.caching import memory_and_disk_cache - from finat.ufl import TensorElement, VectorElement, MixedElement, FiniteElementBase -from finat.element_factory import create_element - -from tsfc.driver import compile_expression_dual_evaluation -from tsfc.ufl_utils import extract_firedrake_constants, hash_expr -from firedrake.utils import IntType, ScalarType, known_pyop2_safe, tuplify -from firedrake.pointeval_utils import runtime_quadrature_element -from firedrake.tsfc_interface import extract_numbered_coefficients, _cachedir +from firedrake.utils import IntType, ScalarType from firedrake.ufl_expr import Argument, Coargument, TrialFunction, TestFunction, action from firedrake.mesh import MissingPointsBehaviour, VertexOnlyMeshTopology, MeshGeometry, MeshTopology, VertexOnlyMesh from firedrake.petsc import PETSc @@ -151,6 +140,11 @@ def options(self) -> InterpolateOptions: """ return self._options + def subdomain_data(self): + """Return cell-iteration subdomain data for the target mesh.""" + domain = self.target_space.mesh().unique() + return {domain: {"cell": [self.options.subset]}} + @cached_property def _interpolator(self): """Access the numerical interpolator. @@ -735,6 +729,7 @@ def _adjoint_weight(self): source_mesh = self.source_mesh.unique() target_mesh = self.target_mesh.unique() + iterset = target_mesh.cell_set if self.subset is None else self.subset for i, (V, component_weight) in enumerate(spaces_and_weights): node_map = get_interp_node_map(source_mesh, target_mesh, V) size = V.finat_element.space_dimension() * V.block_size @@ -744,7 +739,7 @@ def _adjoint_weight(self): }}""" kernel = op2.Kernel(kernel_code, f"multiplicity_{i}") op2.par_loop( - kernel, target_mesh.cell_set, + kernel, iterset, component_weight(op2.INC, node_map), ) with weight.vec as weight_vec: @@ -841,6 +836,7 @@ def _get_form_assembler(self, bcs=None, mat_type=None, sub_mat_type=None): self._assembler_form, bcs=bcs, needs_zeroing=self.access is op2.INC, + access=self.access, ) elif self.rank == 2: return TwoFormAssembler( @@ -849,7 +845,7 @@ def _get_form_assembler(self, bcs=None, mat_type=None, sub_mat_type=None): mat_type=mat_type, sub_mat_type=sub_mat_type, needs_zeroing=True, - allocation_integral_types=("cell",), + access=self.access, ) else: raise ValueError( @@ -857,14 +853,6 @@ def _get_form_assembler(self, bcs=None, mat_type=None, sub_mat_type=None): ) def _get_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None): - if self.source_mesh.unique() is not self.target_mesh.unique(): - return self._get_transmesh_callable( - tensor=tensor, - bcs=bcs, - mat_type=mat_type, - sub_mat_type=sub_mat_type, - ) - assembler = self._get_form_assembler( bcs=bcs, mat_type=mat_type, sub_mat_type=sub_mat_type, ) @@ -903,63 +891,6 @@ def callable(): return callable - def _get_transmesh_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None): - mat_type = mat_type or "aij" - if (isinstance(tensor, Cofunction) and isinstance(self.dual_arg, Cofunction)) and set(tensor.dat).intersection(set(self.dual_arg.dat)): - # adjoint one-form case: we need an empty tensor, so if it shares dats with - # the dual_arg we cannot use it directly, so we store it - f = self._get_tensor(mat_type) - copyout = (partial(f.dat.copy, tensor.dat),) - else: - f = tensor or self._get_tensor(mat_type) - copyout = () - - op2_tensor = f if isinstance(f, op2.Mat) else f.dat - loops = [] - if self.access is op2.INC: - loops.append(op2_tensor.zero) - - # Arguments in the operand are allowed to be from a MixedFunctionSpace - # We need to split the target space V and generate separate kernels - if self.rank == 2: - expressions = {(0,): self.ufl_interpolate} - elif isinstance(self.dual_arg, Coargument): - # Split in the coargument - expressions = dict(split_form(self.ufl_interpolate)) - else: - assert isinstance(self.dual_arg, Cofunction) - # Split in the cofunction: split_form can only split in the coargument - # Replace the cofunction with a coargument to construct the Jacobian - interp = self.ufl_interpolate._ufl_expr_reconstruct_(self.operand, self.target_space) - # Split the Jacobian into blocks - interp_split = dict(split_form(interp)) - # Split the cofunction - dual_split = dict(split_form(self.dual_arg)) - # Combine the splits by taking their action - expressions = {i: action(interp_split[i], dual_split[i[-1:]]) for i in interp_split} - - # Interpolate each sub expression into each function space - for indices, sub_expr in expressions.items(): - sub_op2_tensor = op2_tensor[indices[0]] if self.rank == 1 else op2_tensor - loops.extend(_build_interpolation_callables(sub_expr, sub_op2_tensor, self.access, self.subset, bcs)) - - if bcs and self.rank == 1: - loops.extend(partial(bc.apply, f) for bc in bcs) - - loops.extend(copyout) - - def callable() -> Function | Cofunction | PETSc.Mat | Number: - for l in loops: - l() - if self.rank == 0: - return f.dat.data.item() - elif self.rank == 2: - return f.handle # In this case f is an op2.Mat - else: - return f - - return callable - @property def _allowed_mat_types(self): return {"aij", "baij", "nest", "matfree", None} @@ -1108,204 +1039,6 @@ def _allowed_mat_types(self): return {"aij", "baij", "matfree", None} -@known_pyop2_safe -def _build_interpolation_callables( - expr: Interpolate | ZeroBaseForm, - tensor: op2.Dat | op2.Mat | op2.Global, - access: Literal[op2.WRITE, op2.MIN, op2.MAX, op2.INC], - subset: op2.Subset | None = None, - bcs: Iterable[DirichletBC] | None = None -) -> tuple[Callable, ...]: - """Return a tuple of callables which calculate the interpolation. - - Parameters - ---------- - expr : ufl.Interpolate | ufl.ZeroBaseForm - The symbolic interpolation expression, or a ZeroBaseForm. ZeroBaseForms - are simplified here to avoid code generation when access is WRITE or INC. - tensor : op2.Dat | op2.Mat | op2.Global - Object to hold the result of the interpolation. - access : Literal[op2.WRITE, op2.MIN, op2.MAX, op2.INC] - op2 access descriptor - subset : op2.Subset | None - An optional subset to apply the interpolation over, by default None. - bcs : Iterable[DirichletBC] | None - An optional list of boundary conditions to zero-out in the - output function space. Interpolator rows or columns which are - associated with boundary condition nodes are zeroed out when this is - specified. By default None, by default None. - - Returns - ------- - tuple[Callable, ...] - Tuple of callables which perform the interpolation. - """ - if isinstance(expr, ZeroBaseForm): - # Zero simplification, avoid code-generation - if access is op2.INC: - return () - elif access is op2.WRITE: - return (partial(tensor.zero, subset=subset),) - # Unclear how to avoid codegen for MIN and MAX - # Reconstruct the expression as an Interpolate - V = expr.arguments()[-1].function_space().dual() - expr = interpolate(zero(V.value_shape), V) - - if not isinstance(expr, Interpolate): - raise ValueError("Expecting to interpolate a symbolic Interpolate expression.") - - dual_arg, operand = expr.argument_slots() - assert isinstance(dual_arg, Cofunction | Coargument) - V = dual_arg.function_space().dual() - - if access is op2.READ: - raise ValueError("Can't have READ access for output function") - - # NOTE: The par_loop is always over the target mesh cells. - target_mesh = V.mesh() - source_mesh = extract_unique_domain(operand) or target_mesh - target_element = V.ufl_element() - if isinstance(target_mesh.topology, VertexOnlyMeshTopology): - # For interpolation onto a VOM, we use a FInAT QuadratureElement as the - # target element with runtime point set expressions as their - # quadrature rule point set. - rt_var_name = "rt_X" - target_element = runtime_quadrature_element(source_mesh, target_element, - rt_var_name=rt_var_name) - - cell_set = target_mesh.cell_set - if subset is not None: - assert subset.superset == cell_set - cell_set = subset - - parameters = {} - parameters['scalar_type'] = ScalarType - - copyin = () - copyout = () - - # For the matfree adjoint 1-form and the 0-form, the cellwise kernel will add multiple - # contributions from the facet DOFs of the dual argument. - # The incoming Cofunction needs to be weighted by the reciprocal of the DOF multiplicity. - if isinstance(dual_arg, Cofunction) and not create_element(target_element).is_dg(): - # Create a buffer for the weighted Cofunction - W = dual_arg.function_space() - v = Function(W) - expr = expr._ufl_expr_reconstruct_(operand, v=v) - copyin += (partial(dual_arg.dat.copy, v.dat),) - - # Compute the reciprocal of the DOF multiplicity - wdat = W.make_dat() - m_ = get_interp_node_map(source_mesh, target_mesh, W) - wsize = W.finat_element.space_dimension() * W.block_size - kernel_code = f""" - void multiplicity(PetscScalar *restrict w) {{ - for (PetscInt i=0; i<{wsize}; i++) w[i] += 1; - }}""" - kernel = op2.Kernel(kernel_code, "multiplicity") - op2.par_loop(kernel, cell_set, wdat(op2.INC, m_)) - with wdat.vec as w: - w.reciprocal() - - # Create a callable to apply the weight - with wdat.vec_ro as w, v.dat.vec as y: - copyin += (partial(y.pointwiseMult, y, w),) - - kernel = compile_expression(cell_set.comm, expr, target_element, - domain=source_mesh, parameters=parameters) - ast = kernel.ast - oriented = kernel.oriented - needs_cell_sizes = kernel.needs_cell_sizes - coefficient_numbers = kernel.coefficient_numbers - needs_external_coords = kernel.needs_external_coords - name = kernel.name - kernel = op2.Kernel(ast, name, requires_zeroed_output_arguments=(access is not op2.INC), - flop_count=kernel.flop_count, events=(kernel.event,)) - - parloop_args = [kernel, cell_set] - - coefficients = extract_numbered_coefficients(expr, coefficient_numbers) - if needs_external_coords: - coefficients = [source_mesh.coordinates] + coefficients - - if any(c.dat == tensor for c in coefficients): - output = tensor - tensor = op2.Dat(tensor.dataset) - if access is not op2.WRITE: - copyin += (partial(output.copy, tensor), ) - copyout += (partial(tensor.copy, output), ) - - arguments = expr.arguments() - if isinstance(tensor, op2.Global): - parloop_args.append(tensor(access)) - elif isinstance(tensor, op2.Dat): - V_dest = arguments[-1].function_space() - m_ = get_interp_node_map(source_mesh, target_mesh, V_dest) - parloop_args.append(tensor(access, m_)) - else: - assert access == op2.WRITE # Other access descriptors not done for Matrices. - Vrow = arguments[0].function_space() - Vcol = arguments[1].function_space() - assert tensor.handle.getSize() == (Vrow.dim(), Vcol.dim()) - rows_map = get_interp_node_map(source_mesh, target_mesh, Vrow) - columns_map = get_interp_node_map(source_mesh, target_mesh, Vcol) - lgmaps = None - if bcs: - if is_dual(Vrow): - Vrow = Vrow.dual() - if is_dual(Vcol): - Vcol = Vcol.dual() - bc_rows = [bc for bc in bcs if bc.function_space() == Vrow] - bc_cols = [bc for bc in bcs if bc.function_space() == Vcol] - lgmaps = [(Vrow.local_to_global_map(bc_rows), Vcol.local_to_global_map(bc_cols))] - parloop_args.append(tensor(access, (rows_map, columns_map), lgmaps=lgmaps)) - - if oriented: - co = source_mesh.cell_orientations() - parloop_args.append(co.dat(op2.READ, co.cell_node_map())) - - if needs_cell_sizes: - cs = source_mesh.cell_sizes - parloop_args.append(cs.dat(op2.READ, cs.cell_node_map())) - - for coefficient in coefficients: - m_ = get_interp_node_map(source_mesh, target_mesh, coefficient.function_space()) - parloop_args.append(coefficient.dat(op2.READ, m_)) - - for const in extract_firedrake_constants(expr): - parloop_args.append(const.dat(op2.READ)) - - # Finally, add the target mesh reference coordinates if they appear in the kernel - if isinstance(target_mesh.topology, VertexOnlyMeshTopology): - if target_mesh is not source_mesh: - # NOTE: TSFC will sometimes drop run-time arguments in generated - # kernels if they are deemed not-necessary. - # FIXME: Checking for argument name in the inner kernel to decide - # whether to add an extra coefficient is a stopgap until - # compile_expression_dual_evaluation - # (a) outputs a coefficient map to indicate argument ordering in - # parloops as `compile_form` does and - # (b) allows the dual evaluation related coefficients to be supplied to - # them rather than having to be added post-hoc (likely by - # replacing `to_element` with a CoFunction/CoArgument as the - # target `dual` which would contain `dual` related - # coefficient(s)) - if any(arg.name == rt_var_name for arg in kernel.code[name].args): - # Add the coordinates of the target mesh quadrature points in the - # source mesh's reference cell as an extra argument for the inner - # loop. (With a vertex only mesh this is a single point for each - # vertex cell.) - target_ref_coords = target_mesh.reference_coordinates - m_ = target_ref_coords.cell_node_map() - parloop_args.append(target_ref_coords.dat(op2.READ, m_)) - - parloop = op2.ParLoop(*parloop_args) - if isinstance(tensor, op2.Mat): - return parloop, tensor.assemble - else: - return copyin + (parloop, ) + copyout - - def get_interp_node_map(source_mesh: MeshGeometry, target_mesh: MeshGeometry, fs: WithGeometry) -> op2.Map | None: """Return the map between cells of the target mesh and nodes of the function space. @@ -1341,28 +1074,6 @@ def get_interp_node_map(source_mesh: MeshGeometry, target_mesh: MeshGeometry, fs return m_ -try: - _expr_cachedir = os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] -except KeyError: - _expr_cachedir = os.path.join(tempfile.gettempdir(), - f"firedrake-tsfc-expression-kernel-cache-uid{os.getuid()}") - - -def _compile_expression_key(comm, expr, ufl_element, domain, parameters) -> tuple[Hashable, ...]: - """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`.""" - dual_arg, operand = expr.argument_slots() - return (hash_expr(operand), type(dual_arg), hash(ufl_element), tuplify(parameters)) - - -@memory_and_disk_cache( - hashkey=_compile_expression_key, - cachedir=_cachedir -) -@PETSc.Log.EventDecorator() -def compile_expression(comm, *args, **kwargs): - return compile_expression_dual_evaluation(*args, **kwargs) - - def compose_map_and_cache(map1: op2.Map, map2: op2.Map | None) -> op2.ComposedMap | None: """ Retrieve a :class:`pyop2.ComposedMap` map from the cache of map1 diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index cde7f678a1..73f6772e41 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -84,6 +84,7 @@ def __init__( coefficient_numbers, constant_numbers, dont_split_numbers, + access=op2.INC, diagonal=False ): """A wrapper object for one or more TSFC kernels compiled from a given :class:`~ufl.classes.Form`. @@ -131,6 +132,7 @@ def __init__( events = (kernel.event,) pyop2_kernel = as_pyop2_local_kernel(kernel.ast, kernel.name, len(kernel.arguments), + access=access, flop_count=kernel.flop_count, events=events) kernels.append(KernelInfo(kernel=pyop2_kernel, @@ -150,7 +152,8 @@ def __init__( SplitKernel = collections.namedtuple("SplitKernel", ["indices", "kinfo"]) -def _compile_form_hashkey(form, name, parameters=None, split=True, dont_split=(), diagonal=False): +def _compile_form_hashkey(form, name, parameters=None, split=True, dont_split=(), + diagonal=False, access=op2.INC): return ( form.signature(), name, @@ -158,6 +161,7 @@ def _compile_form_hashkey(form, name, parameters=None, split=True, dont_split=() split, _make_dont_split_numbers(dont_split, form), diagonal, + access, ) @@ -168,7 +172,8 @@ def _compile_form_hashkey(form, name, parameters=None, split=True, dont_split=() cachedir=_cachedir ) @PETSc.Log.EventDecorator() -def compile_form(form, name, parameters=None, split=True, dont_split=(), diagonal=False): +def compile_form(form, name, parameters=None, split=True, dont_split=(), + diagonal=False, access=op2.INC): """Compile a form using TSFC. Parameters @@ -188,6 +193,8 @@ def compile_form(form, name, parameters=None, split=True, dont_split=(), diagona Coefficients that are not to be split into components by form compiler. diagonal : bool If assembling a matrix is it diagonal? + access : pyop2.Access + Access mode for the output tensor. Returns ------- @@ -207,7 +214,7 @@ def compile_form(form, name, parameters=None, split=True, dont_split=(), diagona """ # Check that we get a Form - if not isinstance(form, Form): + if not isinstance(form, (Form, ufl.Interpolate)): raise RuntimeError("Unable to convert object to a UFL form: %s" % repr(form)) if parameters is None: @@ -230,12 +237,12 @@ def compile_form(form, name, parameters=None, split=True, dont_split=(), diagona nargs = 1 iterable = ([(None, )*nargs, form], ) for idx, f in iterable: - f = _real_mangle(f) if isinstance(f, ZeroBaseForm) or f.empty(): # If we're assembling the R space component of a mixed argument, # and that component doesn't actually appear in the form then we # have an empty form, which we should not attempt to assemble. continue + f = _real_mangle(f) # Map local domain/coefficient/constant numbers (as seen inside the # compiler) to the global coefficient/constant numbers meshes = extract_domains(f) @@ -256,6 +263,7 @@ def compile_form(form, name, parameters=None, split=True, dont_split=(), diagona coefficient_numbers, constant_numbers, dont_split_numbers, + access, diagonal, ) for kinfo in tsfc_kernel.kernels: diff --git a/firedrake/ufl_expr.py b/firedrake/ufl_expr.py index f71d111981..7503b820b6 100644 --- a/firedrake/ufl_expr.py +++ b/firedrake/ufl_expr.py @@ -391,6 +391,8 @@ def extract_domains(f): return list(set(mesh._meshes)) else: return [mesh] + elif isinstance(f, ufl.core.base_form_operator.BaseFormOperator): + return f.ufl_domains() elif isinstance(f, (ufl.form.FormSum, ufl.Action)): # ufl.domain.extract_domains does not work. if f._domains is None: diff --git a/pyop2/parloop.py b/pyop2/parloop.py index 6773e6ccf3..6392bc889f 100644 --- a/pyop2/parloop.py +++ b/pyop2/parloop.py @@ -121,12 +121,6 @@ def _kernel_args_(self): @property def map_kernel_args(self): rmap, cmap = self.maps - if isinstance(rmap, MixedMap): - return tuple(itertools.chain.from_iterable( - itertools.chain(rmap_._kernel_args_, cmap_._kernel_args_) - for rmap_ in rmap.split - for cmap_ in cmap.split - )) return tuple(itertools.chain(rmap._kernel_args_, cmap._kernel_args_)) @@ -149,11 +143,7 @@ def _kernel_args_(self): @property def map_kernel_args(self): rmap, cmap = self.maps - return tuple(itertools.chain.from_iterable( - itertools.chain(rmap_._kernel_args_, cmap_._kernel_args_) - for rmap_ in rmap.split - for cmap_ in cmap.split - )) + return tuple(itertools.chain(rmap._kernel_args_, cmap._kernel_args_)) @dataclass diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index d9bd4c08d4..a84133cc55 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -397,6 +397,77 @@ def test_assemble_action_adjoint(V1, V2): assert np.allclose(res.dat.data, res4.dat.data) +def test_assemble_interp_vector_matrix(): + mesh = UnitSquareMesh(2, 2) + V = VectorFunctionSpace(mesh, "CG", 1, dim=2) + W = VectorFunctionSpace(mesh, "DG", 1, dim=2) + x, y = SpatialCoordinate(mesh) + f = Function(V).interpolate(as_vector((x + 2*y, 2*x - y))) + + operator = assemble(interpolate(TrialFunction(V), W)) + actual = assemble(action(operator, f)) + expected = assemble(interpolate(f, W)) + + assert np.allclose(actual.dat.data_ro, expected.dat.data_ro) + + +def test_assemble_interp_mixed_vector_matrix(): + mesh = UnitSquareMesh(2, 2) + X = VectorFunctionSpace(mesh, "CG", 2) + Y = VectorFunctionSpace(mesh, "DG", 1) + V = FunctionSpace(mesh, "CG", 1) + Z = X * V + W = Y * V + x, y = SpatialCoordinate(mesh) + f = Function(Z) + f.sub(0).interpolate(as_vector((x + 2*y, 2*x - y))) + f.sub(1).interpolate(x - y) + + operator = assemble(interpolate(TrialFunction(Z), W), mat_type="nest") + actual = assemble(action(operator, f)) + expected = assemble(interpolate(f, W)) + + for actual_subfunction, expected_subfunction in zip( + actual.subfunctions, expected.subfunctions + ): + assert np.allclose( + actual_subfunction.dat.data_ro, + expected_subfunction.dat.data_ro, + ) + + +def test_interpolate_mixed_vector_in_bilinear_form(): + from firedrake.assemble import ExplicitMatrixAssembler, get_assembler + + mesh = UnitSquareMesh(2, 2) + X = VectorFunctionSpace(mesh, "CG", 2) + Y = VectorFunctionSpace(mesh, "DG", 1) + V = FunctionSpace(mesh, "CG", 1) + Z = X * V + W = Y * V + x, y = SpatialCoordinate(mesh) + f = Function(Z) + f.sub(0).interpolate(as_vector((x + 2*y, 2*x - y))) + f.sub(1).interpolate(x - y) + v = TestFunction(W) + form = inner(interpolate(TrialFunction(Z), W), v) * dx + + assembler = get_assembler(form, mat_type="nest") + assert isinstance(assembler, ExplicitMatrixAssembler) + operator = assembler.assemble() + actual = assemble(action(operator, f)) + interpolated = assemble(interpolate(f, W)) + expected = assemble(inner(interpolated, v) * dx) + + for actual_subfunction, expected_subfunction in zip( + actual.subfunctions, expected.subfunctions + ): + assert np.allclose( + actual_subfunction.dat.data_ro, + expected_subfunction.dat.data_ro, + ) + + @pytest.mark.parallel(2) def test_interpolate_in_form_compiled_reuse(): from firedrake.assemble import OneFormAssembler, get_assembler diff --git a/tests/firedrake/submesh/test_submesh_interpolate.py b/tests/firedrake/submesh/test_submesh_interpolate.py index c3b084a878..04a787fa63 100644 --- a/tests/firedrake/submesh/test_submesh_interpolate.py +++ b/tests/firedrake/submesh/test_submesh_interpolate.py @@ -57,6 +57,46 @@ def _test_submesh_interpolate_cell_cell(mesh, subdomain_cond, fe_fesub): assert assemble(inner(g - f, g - f) * dx(label_value)).real < 1e-14 +@pytest.mark.parallel([1, 3]) +def test_submesh_interpolate_compile_form(): + from firedrake.assemble import OneFormAssembler, get_assembler + + mesh = UnitSquareMesh(4, 4) + x, y = SpatialCoordinate(mesh) + submesh = make_submesh(mesh, conditional(x < 0.51, 1, 0), 999) + V = FunctionSpace(mesh, "CG", 2) + W = FunctionSpace(submesh, "CG", 1) + f = Function(V).interpolate(x + 2*y) + xs, ys = SpatialCoordinate(submesh) + expected = Function(W).interpolate(xs + 2*ys) + + actual = assemble(interpolate(f, W)) + assert np.allclose( + actual.dat.data_ro_with_halos, + expected.dat.data_ro_with_halos, + ) + + operator = assemble(interpolate(TrialFunction(V), W)) + actual = assemble(action(operator, f)) + assert np.allclose( + actual.dat.data_ro_with_halos, + expected.dat.data_ro_with_halos, + ) + + v = TestFunction(W) + subdx = Measure( + "dx", + submesh, + intersect_measures=(Measure("dx", mesh),), + ) + form = inner(interpolate(f, W), v) * subdx + assembler = get_assembler(form) + assert isinstance(assembler, OneFormAssembler) + actual = assembler.assemble() + expected = assemble(inner(expected, v) * dx(submesh)) + assert np.allclose(actual.dat.data_ro, expected.dat.data_ro) + + @pytest.mark.parametrize('nelem', [2, 4, 8, None]) @pytest.mark.parametrize('fe_fesub', [[("DQ", 0), ("DQ", 0)], [("Q", 4), ("Q", 5)]]) diff --git a/tsfc/driver.py b/tsfc/driver.py index 71efba33d0..8cbe7b5fa8 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -9,13 +9,16 @@ from ufl.algorithms.analysis import has_type from ufl.algorithms.apply_coefficient_split import CoefficientSplitter from ufl.classes import Form, GeometricQuantity -from ufl.domain import extract_unique_domain, extract_domains +from ufl.domain import MeshSequence, extract_unique_domain, extract_domains import gem import gem.impero_utils as impero_utils import finat from finat.element_factory import as_fiat_cell +from finat.point_set import UnknownPointSet +from finat.quadrature import QuadratureRule +from finat.ufl import FiniteElement, TensorElement from tsfc import fem, ufl_utils from tsfc.logging import logger @@ -52,6 +55,21 @@ """ +TSFCInterpolationData = collections.namedtuple( + "TSFCInterpolationData", + ["domain", "iteration_domain", "integral_type", "subdomain_id", + "domain_integral_type_map", "enabled_coefficients", "integrals", + "expression"], +) + +TSFCInterpolationFormData = collections.namedtuple( + "TSFCInterpolationFormData", + ["original_form", "preprocessed_form", "reduced_coefficients", + "function_replace_map", "coefficient_split", + "original_coefficient_positions", "constants"], +) + + def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), diagonal=False): """Compiles a UFL form into a set of assembly kernels. @@ -76,6 +94,9 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di """ cpu_time = time.time() + if isinstance(form, ufl.Interpolate): + return compile_interpolate(form, prefix=prefix, parameters=parameters) + assert isinstance(form, Form) GREEN = "\033[1;37;32m%s\033[0m" @@ -110,6 +131,101 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di return kernels +def compile_interpolate(expression, prefix="interpolate", parameters=None): + """Compile an interpolation using the integral kernel builder.""" + parameters = preprocess_parameters(parameters) + complex_mode = is_complex(parameters["scalar_type"]) + original_expression = expression + original_coefficients = expression.coefficients() + dual_arg, operand = expression.argument_slots() + target_domain = dual_arg.ufl_function_space().ufl_domain() + if isinstance(target_domain, MeshSequence): + target_domains = set(target_domain.meshes) + if len(target_domains) != 1: + raise NotImplementedError( + "Interpolation onto multiple distinct meshes is not supported" + ) + target_domain, = target_domains + source_domain = ( + extract_unique_domain(operand) + or target_domain + ) + all_domains = expression.ufl_domains() + + target_element = expression.ufl_element() + if ( + target_domain.topological_dimension == 0 + and source_domain.topological_dimension > 0 + ): + cell = source_domain.ufl_cell() + point_expr = gem.Variable("rt_X", (1, cell.topological_dimension)) + point_set = UnknownPointSet(point_expr) + rule = QuadratureRule( + point_set, weights=[1.0], ref_el=as_fiat_cell(cell) + ) + shape = target_element.pullback.physical_value_shape( + target_element, target_domain + ) + target_element = FiniteElement( + "Quadrature", cell=cell, degree=0, quad_scheme=rule + ) + if shape: + symmetry = None if len(shape) < 2 else expression.ufl_element().symmetry() + target_element = TensorElement( + target_element, shape=shape, symmetry=symmetry + ) + + operand = apply_mapping(operand, target_element, source_domain) + operand = ufl_utils.preprocess_expression( + operand, complex_mode=complex_mode + ) + operand = simplify_abs(operand, complex_mode) + expression = ufl.Interpolate(operand, dual_arg) + + coefficients = expression.coefficients() + coefficient_split = {} + for coefficient in coefficients: + element = coefficient.ufl_element() + if type(element) is finat.ufl.MixedElement: + domain = extract_unique_domain( + coefficient, expand_mesh_sequence=False + ) + coefficient_split[coefficient] = [ + ufl.Coefficient(ufl.FunctionSpace(mesh, subelement)) + for mesh, subelement in zip( + domain.iterable_like(element), element.sub_elements + ) + ] + + form_data = TSFCInterpolationFormData( + original_form=original_expression, + preprocessed_form=expression, + reduced_coefficients=coefficients, + function_replace_map={coefficient: coefficient for coefficient in coefficients}, + coefficient_split=coefficient_split, + original_coefficient_positions=tuple( + original_coefficients.index(coefficient) + for coefficient in coefficients + ), + constants=extract_firedrake_constants(expression), + ) + integral_data = TSFCInterpolationData( + domain=source_domain, + iteration_domain=target_domain, + integral_type="cell", + subdomain_id=("everywhere",), + domain_integral_type_map={domain: "cell" for domain in all_domains}, + enabled_coefficients=(True,) * len(coefficients), + integrals=(), + expression=expression, + ) + return [ + compile_integral( + integral_data, form_data, prefix, parameters, diagonal=False + ) + ] + + def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=False): """Compiles a UFL integral into an assembly kernel. @@ -142,8 +258,13 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F coefficient_split[coeff] = form_data.coefficient_split[coeff] coefficient_numbers.append(form_data.original_coefficient_positions[i]) mesh = integral_data.domain - all_meshes = extract_domains(form_data.original_form) - domain_number = all_meshes.index(mesh) + if isinstance(integral_data, TSFCInterpolationData): + iteration_domain = integral_data.iteration_domain + all_meshes = tuple(integral_data.domain_integral_type_map) + else: + iteration_domain = mesh + all_meshes = extract_domains(form_data.original_form) + domain_number = all_meshes.index(iteration_domain) integral_data_info = TSFCIntegralDataInfo( domain=integral_data.domain, @@ -172,12 +293,17 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F # so we should attach the constants to integral data instead builder.set_constants(form_data.constants) ctx = builder.create_context() - for integral in integral_data.integrals: + if isinstance(integral_data, TSFCInterpolationData): params = parameters.copy() - params.update(integral.metadata()) # integral metadata overrides - integrand_exprs = builder.compile_integrand(integral.integrand(), params, ctx) - integral_exprs = builder.construct_integrals(integrand_exprs, params) - builder.stash_integrals(integral_exprs, params, ctx) + interpolate_exprs = builder.compile_interpolate(integral_data.expression, params, ctx) + builder.stash_integrals(interpolate_exprs, params, ctx) + else: + for integral in integral_data.integrals: + params = parameters.copy() + params.update(integral.metadata()) # integral metadata overrides + integrand_exprs = builder.compile_integrand(integral.integrand(), params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) return builder.construct_kernel(kernel_name, ctx, parameters["add_petsc_events"]) diff --git a/tsfc/fem.py b/tsfc/fem.py index 3657885a4b..a6840a7997 100644 --- a/tsfc/fem.py +++ b/tsfc/fem.py @@ -176,6 +176,12 @@ def coefficient(self, ufl_coefficient, r): assert r is None return self._wrapee.coefficient(ufl_coefficient, self.restriction) + def coefficient_components(self, ufl_coefficient, r): + assert r is None + return self._wrapee.coefficient_components( + ufl_coefficient, self.restriction + ) + class CoordinateMapping(PhysicalGeometry): """Callback class that provides physical geometry to FInAT elements. @@ -359,18 +365,44 @@ def dual_evaluate(expression: ufl.Interpolate, to_element: FiniteElementBase, ke coordinate_mapping = CoordinateMapping(analyse_modified_terminal(coefficient), ctx) else: coordinate_mapping = None - evaluation, point_indices, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) - quadrature_multiindex = tuple(point_indices) - if isinstance(dual_arg, ufl.Cofunction): - gem_dual = kernel_cfg["interface"].coefficient_map[dual_arg] - if is_complex(kernel_cfg["scalar_type"]): - evaluation = gem.MathFunction("conj", evaluation) - # The dual argument contracts over the nodes, so the basis indices are - # reduction indices like the points, not indices of the return value. - evaluation = evaluation * gem_dual[basis_indices] - quadrature_multiindex += tuple(basis_indices) + gem_duals = kernel_cfg["interface"].coefficient_components( + dual_arg, None + ) + else: + gem_duals = () + + if len(gem_duals) > 1: + assert to_element.is_mixed + assert len(to_element.elements) == len(gem_duals) + # The summands do not share their points, so each one contracts on its own. + evaluations = [] + for element, gem_dual in zip(to_element.elements, gem_duals): + evaluation, point_indices, basis_indices = element.dual_evaluation( + fn, coordinate_mapping + ) + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + evaluations.append(gem.IndexSum( + evaluation * gem_dual[basis_indices], + tuple(point_indices) + basis_indices + )) + evaluation = gem.optimise.make_sum(evaluations) + quadrature_multiindex = () basis_indices = () + else: + evaluation, point_indices, basis_indices = to_element.dual_evaluation( + fn, coordinate_mapping + ) + quadrature_multiindex = tuple(point_indices) + if gem_duals: + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + # The dual argument contracts over the nodes, so the basis indices + # are reduction indices like the points, not return value indices. + evaluation = evaluation * gem_duals[0][basis_indices] + quadrature_multiindex += tuple(basis_indices) + basis_indices = () return evaluation, quadrature_multiindex, basis_indices @@ -820,7 +852,11 @@ def translate_constant_value(terminal, mt, ctx): @translate.register(ufl.Interpolate) def translate_interpolate(terminal: ufl.Interpolate, mt: ModifiedTerminal, ctx: ContextBase) -> gem.Node: - domain = extract_unique_domain(terminal) + dual_arg, operand = terminal.argument_slots() + domain = ( + extract_unique_domain(operand) + or dual_arg.ufl_function_space().ufl_domain() + ) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) kernel_cfg = { "interface": CellVolumeKernelInterface( @@ -899,8 +935,12 @@ def take_singleton(xs): unsummed_indices = set(chain(*argument_multiindices)) unsummed_indices.update(ctx.unsummed_coefficient_indices) for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): - indices = tuple(i for i in var.index_ordering() if i not in unsummed_indices) - value = gem.IndexSum(gem.Product(expr, var), indices) + product = gem.Product(expr, var) + indices = tuple( + i for i in dict.fromkeys(chain(var.index_ordering(), beta)) + if i not in unsummed_indices and i in product.free_indices + ) + value = gem.IndexSum(product, indices) summands.append(gem.optimise.contraction(value)) optimised_value = gem.optimise.make_sum(summands) value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) @@ -912,7 +952,12 @@ def take_singleton(xs): if hasattr(argument_multiindices, "values"): argument_multiindices = argument_multiindices.values() allowed_indices = set(chain(ctx.point_indices, *argument_multiindices)) - assert set(result.free_indices) - ctx.unsummed_coefficient_indices <= allowed_indices + unexpected_indices = ( + set(result.free_indices) + - ctx.unsummed_coefficient_indices + - allowed_indices + ) + assert not unexpected_indices, unexpected_indices # Detect Jacobian of affine cells if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2 diff --git a/tsfc/kernel_interface/__init__.py b/tsfc/kernel_interface/__init__.py index 3c20720c33..abd59fd892 100644 --- a/tsfc/kernel_interface/__init__.py +++ b/tsfc/kernel_interface/__init__.py @@ -17,6 +17,10 @@ def coefficient(self, ufl_coefficient, restriction): """A function that maps :class:`ufl.Coefficient`s to GEM expressions.""" + @abstractmethod + def coefficient_components(self, ufl_coefficient, restriction): + """Return GEM expressions for a coefficient's stored components.""" + @abstractmethod def constant(self, const): """Return the GEM expression corresponding to the constant.""" diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 5d61a916aa..0f01f1fafd 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -5,9 +5,12 @@ from itertools import chain, product import copy +from ufl.classes import Cofunction from ufl.utils.sequences import max_degree from ufl.domain import extract_unique_domain +from ufl.algorithms.apply_coefficient_split import CoefficientSplitter +import finat import gem import gem.impero_utils as impero_utils import petsctools @@ -41,6 +44,7 @@ def __init__(self, scalar_type): # Coefficients self.coefficient_map = collections.OrderedDict() + self.coefficient_split = {} # Constants self.constant_map = collections.OrderedDict() @@ -63,6 +67,16 @@ def coefficient(self, ufl_coefficient, restriction): else: return kernel_arg[{'+': 0, '-': 1}[restriction]] + def coefficient_components(self, ufl_coefficient, restriction): + """Return GEM expressions for a coefficient's stored components.""" + coefficients = self.coefficient_split.get( + ufl_coefficient, (ufl_coefficient,) + ) + return tuple( + self.coefficient(coefficient, restriction) + for coefficient in coefficients + ) + def constant(self, const): return self.constant_map[const] @@ -136,6 +150,54 @@ def domain_integral_type_map(self): class KernelBuilderMixin(object): """Mixin for KernelBuilder classes.""" + def compile_interpolate(self, expression, params, ctx): + """Compile UFL interpolate. + + :arg expression: UFL interpolate. + :arg params: a dict containing "quadrature_rule". + :arg ctx: context created with :meth:`create_context` method. + + See :meth:`create_context` for typical calling sequence. + """ + expression = CoefficientSplitter(self.coefficient_split)( + expression + ) + target_element = self.create_element(expression.ufl_element()) + config = self.fem_config() + config.update( + argument_multiindices=self.argument_multiindices, + index_cache=ctx["index_cache"], + ) + if isinstance(target_element, finat.QuadratureElement): + config["quadrature_rule"] = target_element._rule + evaluation, quadrature_multiindex, basis_indices = fem.dual_evaluate( + expression, target_element, config + ) + dual_arg, _ = expression.argument_slots() + if not isinstance(dual_arg, Cofunction): + arguments = expression.arguments() + argument_number = arguments.index(dual_arg) + output_indices = self.argument_multiindices[argument_number] + if basis_indices != output_indices: + if tuple(i.extent for i in basis_indices) != tuple( + i.extent for i in output_indices + ): + raise ValueError("Interpolation output index shape mismatch") + mapper = gem.node.MemoizerArg( + gem.optimise.filtered_replace_indices + ) + evaluation = mapper( + evaluation, tuple(zip(basis_indices, output_indices)) + ) + + mode = pick_mode(params["mode"]) + ctx["quadrature_indices"].extend(quadrature_multiindex) + # Argument factorisation does not cancel every Delta here, so lower them. + ctx["finalise_options"]["replace_delta"] = True + return mode.Integrals( + [evaluation], quadrature_multiindex, self.argument_multiindices, params + ) + def compile_integrand(self, integrand, params, ctx): """Compile UFL integrand. @@ -220,6 +282,7 @@ def compile_gem(self, ctx): options = dict(reduce(operator.and_, [mode.finalise_options.items() for mode in mode_irs.keys()])) + options.update(ctx['finalise_options']) expressions = impero_utils.preprocess_gem(expressions, **options) # Let the kernel interface inspect the optimised IR to register @@ -277,6 +340,11 @@ def create_context(self): Dict for mode representations. + *finalise_options* + + Options overriding the modes' own :func:`impero_utils.preprocess_gem` + options. + For each set of integrals to make a kernel for (i,e., `integral_data.integrals`), one must first create a ctx object by calling :meth:`create_context` method. @@ -299,6 +367,7 @@ def create_context(self): """ return {'index_cache': {}, 'quadrature_indices': [], + 'finalise_options': {}, 'mode_irs': collections.OrderedDict()} @@ -577,7 +646,8 @@ def expression(restricted): c_shape = copy.deepcopy(u_shape) rs_tuples = [] for arg_num, arg in enumerate(arguments): - integral_type = domain_integral_type_map[extract_unique_domain(arg)] + domain = arg.ufl_function_space().ufl_domain() + integral_type = domain_integral_type_map[domain] if integral_type is None: raise RuntimeError(f"Can not determine integral_type on {arg}") if integral_type.startswith("interior_facet"): diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index 15567bbc9e..49f795d845 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -4,7 +4,7 @@ from ufl import Coefficient, FunctionSpace from ufl.domain import MeshSequence -from finat.ufl import FiniteElement +from finat.ufl import MixedElement as ufl_MixedElement, FiniteElement import gem from gem.flop_count import count_flops @@ -206,7 +206,13 @@ def set_coefficients(self, coefficients): self.coefficient_split = {} for i, coefficient in enumerate(coefficients): - self._coefficient(coefficient, f"w_{i}") + if type(coefficient.ufl_element()) == ufl_MixedElement: + subcoeffs = coefficient.subfunctions # Firedrake-specific + self.coefficient_split[coefficient] = subcoeffs + for j, subcoeff in enumerate(subcoeffs): + self._coefficient(subcoeff, f"w_{i}_{j}") + else: + self._coefficient(coefficient, f"w_{i}") def set_constants(self, constants): for i, const in enumerate(constants): @@ -287,6 +293,7 @@ def __init__(self, integral_data_info, scalar_type, self.local_tensor = None self.coefficient_number_index_map = OrderedDict() self.integral_data_info = integral_data_info + self.coefficient_split = integral_data_info.coefficient_split self._domain_integral_type_map = integral_data_info.domain_integral_type_map # For consistency with ExpressionKernelBuilder. self.set_arguments() diff --git a/tsfc/spectral.py b/tsfc/spectral.py index a521fdb2fd..ac270fafd2 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -2,7 +2,7 @@ from functools import partial from itertools import chain, zip_longest -from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem.gem import Delta, FlexiblyIndexed, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination @@ -123,7 +123,7 @@ def classify(argument_indices, expression, delta_inside): if n == 0: return OTHER elif n == 1: - if isinstance(expression, (Delta, Indexed)) and not delta_inside(expression): + if isinstance(expression, (Delta, FlexiblyIndexed, Indexed)) and not delta_inside(expression): return ATOMIC else: return COMPOUND diff --git a/tsfc/ufl_utils.py b/tsfc/ufl_utils.py index 1740248d51..eaf8bf4ffd 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -45,8 +45,11 @@ class InterpolateMapper(MultiFunction): def interpolate(self, o: ufl.Interpolate, operand: Expr) -> Expr: dual_arg, _ = o.argument_slots() - domain = extract_unique_domain(o) - element = o.ufl_function_space().ufl_element() + domain = ( + extract_unique_domain(operand) + or dual_arg.ufl_function_space().ufl_domain() + ) + element = o.ufl_element() operand = apply_mapping(operand, element, domain) expr = o._ufl_expr_reconstruct_(operand, v=dual_arg) return element.pullback.apply(ReferenceValue(expr), domain) From 8b01b0b2b5df6b3e7b00b2799fc06daffeeb93a2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 11 Aug 2026 00:58:49 +0100 Subject: [PATCH 04/10] Match boundary conditions to a space without function_space_match function_space_match folded two questions into one predicate: whether a boundary condition belongs to a form's space, and which block of a mixed space it belongs to. Its callers then had to pre-index the space they passed, so the block comparison happened in two places at once and the predicate had to reconstruct what the caller had already discarded. Give function_space a `parent` argument, which walks up through indexed and component subspaces to the top-level space, and let each caller ask the two questions separately: the top-level spaces must be equal, and for a mixed space the block index must match the row or column. This surfaces the TypeError that ExplicitMatrixAssembler already raises for a boundary condition on an unrelated space, so the two tests that assert on it now name that type. Co-Authored-By: Claude Sonnet 5 --- firedrake/assemble.py | 19 +++++++---- firedrake/bcs.py | 45 +++++++++++--------------- tests/firedrake/regression/test_bcs.py | 4 +-- 3 files changed, 32 insertions(+), 36 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 08df0ce215..9a446a32cd 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1544,7 +1544,7 @@ def _apply_bc(self, tensor, bc, u=None): index = 0 if V.index is None else V.index space = V if V.parent is None else V.parent if isinstance(bc, DirichletBC): - if not any(bc.function_space_match(fs) for fs in spaces): + if not any(bc.function_space(parent=True) == fs for fs in spaces): raise TypeError("bc space does not match the test or trial function space") if spaces[0] != spaces[1]: # Not on a diagonal block, we cannot set diagonal entries @@ -2129,20 +2129,25 @@ def get_indicess(self): def _filter_bcs(self, row, col): assert len(self._form.arguments()) == 2 and not self._diagonal + def block_index(bc): + fs = bc.function_space() + if fs.component is not None: + fs = fs.parent + return fs.index + test_space = self.test_function_space - if len(test_space) > 1: - test_space = test_space[row] bcrow = tuple( - bc for bc in self._bcs if bc.function_space_match(test_space) + bc for bc in self._bcs + if bc.function_space(parent=True) == test_space + and (len(test_space) == 1 or block_index(bc) == row) ) trial_space = self.trial_function_space - if len(trial_space) > 1: - trial_space = trial_space[col] bccol = tuple( bc for bc in self._bcs if isinstance(bc, DirichletBC) - and bc.function_space_match(trial_space) + and bc.function_space(parent=True) == trial_space + and (len(trial_space) == 1 or block_index(bc) == col) ) return bcrow, bccol diff --git a/firedrake/bcs.py b/firedrake/bcs.py index 9f4f6b4ec8..586d1707fe 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -76,11 +76,26 @@ def __iter__(self): yield self yield from itertools.chain(*self.bcs) - def function_space(self): + def function_space(self, parent=False): '''The :class:`.FunctionSpace` on which this boundary condition should - be applied.''' + be applied. - return self._function_space + Parameters + ---------- + parent : bool + If ``True``, walk up through any indexed or component subspaces + and return the top-level function space instead. + + Returns + ------- + firedrake.functionspaceimpl.WithGeometry + The function space. + ''' + V = self._function_space + if parent: + while V.parent is not None: + V = V.parent + return V def function_space_index(self): fs = self._function_space @@ -90,30 +105,6 @@ def function_space_index(self): raise RuntimeError("This function should only be called when function space is indexed") return fs.index - def function_space_match( - self, V: ufl.functionspace.AbstractFunctionSpace - ) -> bool: - """Return whether this boundary condition is defined on ``V``. - - Parameters - ---------- - V : ufl.functionspace.AbstractFunctionSpace - Function space to compare with this boundary condition's space. - - Returns - ------- - bool - Whether the spaces represent the same indexed field. - """ - if ufl.duals.is_dual(V): - V = V.dual() - bc_space = self.function_space() - if V.parent is None: - if bc_space.parent is None: - return bc_space == V - return bc_space.parent == V - return bc_space.parent == V.parent and bc_space.index == V.index - @cached_property def domain_args(self): r"""The sub_domain the BC applies to.""" diff --git a/tests/firedrake/regression/test_bcs.py b/tests/firedrake/regression/test_bcs.py index 9e43ba805b..ccd4477e99 100644 --- a/tests/firedrake/regression/test_bcs.py +++ b/tests/firedrake/regression/test_bcs.py @@ -56,7 +56,7 @@ def test_assemble_bcs_wrong_fs(V, measure): u, v = TrialFunction(V), TestFunction(V) W = FunctionSpace(V.mesh(), "CG", 2) - with pytest.raises(RuntimeError): + with pytest.raises(TypeError): assemble(inner(u, v)*measure, bcs=[DirichletBC(W, 32, 1)]) @@ -65,7 +65,7 @@ def test_assemble_bcs_wrong_fs_interior(V): u, v = TrialFunction(V), TestFunction(V) W = FunctionSpace(V.mesh(), "CG", 2) n = FacetNormal(V.mesh()) - with pytest.raises(RuntimeError): + with pytest.raises(TypeError): assemble(inner(jump(u, n), jump(v, n))*dS, bcs=[DirichletBC(W, 32, 1)]) From edc17c3022d669db9b775f64c5a11883f978fed3 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Tue, 11 Aug 2026 00:59:04 +0100 Subject: [PATCH 05/10] TSFC: resolve an unsplit mixed argument's integral type prepare_arguments looks an argument's domain up in the integral type map. An argument that the form compiler has not split, which is what compile_form(split=False) hands it for a mixed space, carries a MeshSequence rather than one of the meshes the map is keyed on, so the lookup raises KeyError. PatchPC builds its patch operator that way, so every mixed-space Vanka or star patch failed there. Fall back to the meshes the sequence holds. They are all iterated by the same integral, so a single integral type has to come back, and the tuple unpacking says so. Co-Authored-By: Claude Sonnet 5 --- tsfc/kernel_interface/common.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 0f01f1fafd..4cb817bef0 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -7,7 +7,7 @@ from ufl.classes import Cofunction from ufl.utils.sequences import max_degree -from ufl.domain import extract_unique_domain +from ufl.domain import MeshSequence, extract_unique_domain from ufl.algorithms.apply_coefficient_split import CoefficientSplitter import finat @@ -647,7 +647,15 @@ def expression(restricted): rs_tuples = [] for arg_num, arg in enumerate(arguments): domain = arg.ufl_function_space().ufl_domain() - integral_type = domain_integral_type_map[domain] + try: + integral_type = domain_integral_type_map[domain] + except KeyError: + # An unsplit argument (e.g. a mixed-space patch argument) reports + # its domain as a MeshSequence rather than a single mesh: every + # mesh it sequences is the same iteration, so they must agree. + if not isinstance(domain, MeshSequence): + raise + integral_type, = {domain_integral_type_map[m] for m in domain.meshes} if integral_type is None: raise RuntimeError(f"Can not determine integral_type on {arg}") if integral_type.startswith("interior_facet"): From 682ba86755d15651901664a02304bd4c2c8740c6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 26 Aug 2026 18:11:08 +0100 Subject: [PATCH 06/10] Interpolate onto a point cloud through the integral kernel builder A point cloud target has no element to evaluate a dual basis against, so the interpolation compiles against a quadrature element on the source cell whose points are only known at run time. That element reached apply_mapping but not the kernel builder, which rebuilt the target element from the dual argument and tabulated the source on the vertex cell's empty point set. Carry it on the interpolation data instead. The domain numbering was inconsistent either side of the kernel: the point cloud counts as a domain of the interpolation, so it takes a slot in the cell sizes it cannot fill, and extract_domains has to reach the argument slots to see it at all. Co-Authored-By: Claude Opus 5 --- tsfc/driver.py | 7 +++++-- tsfc/kernel_interface/common.py | 8 ++++++-- tsfc/kernel_interface/firedrake_loopy.py | 14 ++++++++------ 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/tsfc/driver.py b/tsfc/driver.py index 8cbe7b5fa8..638a73d87f 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -59,7 +59,7 @@ "TSFCInterpolationData", ["domain", "iteration_domain", "integral_type", "subdomain_id", "domain_integral_type_map", "enabled_coefficients", "integrals", - "expression"], + "expression", "target_element"], ) TSFCInterpolationFormData = collections.namedtuple( @@ -218,6 +218,7 @@ def compile_interpolate(expression, prefix="interpolate", parameters=None): enabled_coefficients=(True,) * len(coefficients), integrals=(), expression=expression, + target_element=target_element, ) return [ compile_integral( @@ -295,7 +296,9 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F ctx = builder.create_context() if isinstance(integral_data, TSFCInterpolationData): params = parameters.copy() - interpolate_exprs = builder.compile_interpolate(integral_data.expression, params, ctx) + interpolate_exprs = builder.compile_interpolate( + integral_data.expression, integral_data.target_element, params, ctx + ) builder.stash_integrals(interpolate_exprs, params, ctx) else: for integral in integral_data.integrals: diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 4cb817bef0..128e7643a0 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -150,10 +150,14 @@ def domain_integral_type_map(self): class KernelBuilderMixin(object): """Mixin for KernelBuilder classes.""" - def compile_interpolate(self, expression, params, ctx): + def compile_interpolate(self, expression, target_element, params, ctx): """Compile UFL interpolate. :arg expression: UFL interpolate. + :arg target_element: UFL element of the interpolation target. This is + not the dual argument's own element when the target is a point + cloud: the points are then only known at run time, so the target + is a quadrature element on the source cell. :arg params: a dict containing "quadrature_rule". :arg ctx: context created with :meth:`create_context` method. @@ -162,7 +166,7 @@ def compile_interpolate(self, expression, params, ctx): expression = CoefficientSplitter(self.coefficient_split)( expression ) - target_element = self.create_element(expression.ufl_element()) + target_element = self.create_element(target_element) config = self.fem_config() config.update( argument_multiindices=self.argument_multiindices, diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index 49f795d845..c4c4d91bb8 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -152,18 +152,20 @@ def set_cell_sizes(self, domains): measure of the mesh size around each vertex (hence this lives in P1). - Should the domain have topological dimension 0 this does - nothing. + A domain of topological dimension 0 gets a ``None`` entry: every + domain must keep its slot, since the active domain numbers index + this dict positionally. """ self._cell_sizes = {} for i, domain in enumerate(domains): if domain.ufl_cell().topological_dimension > 0: - # Can't create P1 since only P0 is a valid finite element if - # topological_dimension is 0 and the concept of "cell size" - # is not useful for a vertex. f = Coefficient(FunctionSpace(domain, FiniteElement("P", domain.ufl_cell(), 1))) expr = prepare_coefficient(f, f"cell_sizes_{i}", self._domain_integral_type_map) - self._cell_sizes[domain] = expr + else: + # Only P0 is a valid finite element on a vertex, and the + # concept of "cell size" is not useful there. + expr = None + self._cell_sizes[domain] = expr def create_element(self, element, **kwargs): """Create a FInAT element (suitable for tabulating with) given From 1028f51203f6ecee6c63ec59534ca00019c5b1f6 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 27 Aug 2026 12:16:59 +0100 Subject: [PATCH 07/10] Compare boundary conditions against primal spaces An Interpolate takes its test function from the dual space, and a dual space compares unequal to its primal. Once same-mesh interpolation assembles through ExplicitMatrixAssembler, that made two comparisons in the boundary condition path answer the wrong question. Asking which argument a condition belongs to, a condition named on the target space matched neither, and assembly raised a TypeError instead of applying it. Asking whether the matrix is square, an interpolation from a space to itself looked rectangular, so its boundary rows were zeroed without ever taking the diagonal entry that keeps them solvable. Take the primal space of each argument once, where the arguments are read, and let both comparisons stand as they are. The matrix a mixed interpolation assembles is now the one Form assembly builds, so the tests and the manual describe that matrix: every block of a nest is allocated, the off-diagonal ones assembling to zero rather than staying NULL; a blocked space defaults its nest blocks to baij; an unmixed space collapses its one-block nest instead of refusing it; and a mixed space rejects baij with the ValueError a Form raises. Co-Authored-By: Claude Opus 5 --- docs/source/interpolation.rst | 3 +- firedrake/assemble.py | 24 +++++++++++-- .../firedrake/regression/test_interpolate.py | 18 ++++++++++ .../regression/test_interpolation_manual.py | 2 +- .../test_interpolation_operators.py | 34 +++++++++++-------- 5 files changed, 62 insertions(+), 19 deletions(-) diff --git a/docs/source/interpolation.rst b/docs/source/interpolation.rst index 50441902ff..43ca8f3147 100644 --- a/docs/source/interpolation.rst +++ b/docs/source/interpolation.rst @@ -458,7 +458,8 @@ each block given by \end{pmatrix} The off-diagonal blocks are zero since the dofs are applied component-wise. Firedrake's form -compiler recognises this and avoids assembling the zero blocks. +compiler recognises this and avoids assembling the zero blocks, which the nest +still allocates. We can assemble more general interpolation matrices between mixed function spaces by interpolating vector expressions with arguments. For example, by doing diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 9a446a32cd..e85190e05c 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -1377,6 +1377,24 @@ def _get_mat_type(mat_type, sub_mat_type, arguments): return mat_type, sub_mat_type +def _primal_space(V): + """Return the primal space of a form argument's function space. + + Parameters + ---------- + V : firedrake.functionspaceimpl.WithGeometry + The function space of a form argument, primal or dual. + + Returns + ------- + firedrake.functionspaceimpl.WithGeometry + The primal space. An `~ufl.Interpolate` takes its test function from + the dual space, which compares unequal to the space a boundary + condition names. + """ + return V.dual() if ufl.duals.is_dual(V) else V + + class ExplicitMatrixAssembler(ParloopFormAssembler): """Class for assembling a matrix. @@ -1536,7 +1554,7 @@ def _all_assemblers(self): def _apply_bc(self, tensor, bc, u=None): assert u is None op2tensor = tensor.M - spaces = tuple(a.function_space() for a in tensor.a.arguments()) + spaces = tuple(_primal_space(a.function_space()) for a in tensor.a.arguments()) V = bc.function_space() component = V.component if component is not None: @@ -2135,14 +2153,14 @@ def block_index(bc): fs = fs.parent return fs.index - test_space = self.test_function_space + test_space = _primal_space(self.test_function_space) bcrow = tuple( bc for bc in self._bcs if bc.function_space(parent=True) == test_space and (len(test_space) == 1 or block_index(bc) == row) ) - trial_space = self.trial_function_space + trial_space = _primal_space(self.trial_function_space) bccol = tuple( bc for bc in self._bcs if isinstance(bc, DirichletBC) diff --git a/tests/firedrake/regression/test_interpolate.py b/tests/firedrake/regression/test_interpolate.py index 4539c20504..0eca9d52f8 100644 --- a/tests/firedrake/regression/test_interpolate.py +++ b/tests/firedrake/regression/test_interpolate.py @@ -641,6 +641,24 @@ def test_interpolator_reuse(family, degree, mode): assert np.allclose(result.dat.data, expected) +@pytest.mark.parallel([1, 3]) +def test_square_space_bcs(): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 1) + rg = RandomGenerator(PCG64(seed=123456789)) + w = rg.uniform(V) + + # Source and target agree, so the interpolation has a diagonal to carry + # the boundary rows, just as a Form on the same spaces does. + I = assemble(interpolate(2 * TrialFunction(V), V), bcs=[DirichletBC(V, 0, 1)]) + result = assemble(action(I, w)) + + expected = Function(V).assign(2 * w) + DirichletBC(V, w, 1).apply(expected) + + assert np.allclose(result.dat.data, expected.dat.data) + + def test_mixed_space_bcs(): mesh = UnitSquareMesh(2, 2) V = FunctionSpace(mesh, "CG", 1) diff --git a/tests/firedrake/regression/test_interpolation_manual.py b/tests/firedrake/regression/test_interpolation_manual.py index cedaee54be..3c32838a8d 100644 --- a/tests/firedrake/regression/test_interpolation_manual.py +++ b/tests/firedrake/regression/test_interpolation_manual.py @@ -260,7 +260,7 @@ def test_mixed_space_interpolation(): for j in range(2): sub_mat = I.petscmat.getNestSubMatrix(i, j) if i != j: - assert not sub_mat + assert sub_mat.norm() == 0.0 continue else: res_block = assemble(interpolate(TrialFunction(U.sub(j)), W.sub(i))) diff --git a/tests/firedrake/regression/test_interpolation_operators.py b/tests/firedrake/regression/test_interpolation_operators.py index 6c97264ded..bcaa4d56f8 100644 --- a/tests/firedrake/regression/test_interpolation_operators.py +++ b/tests/firedrake/regression/test_interpolation_operators.py @@ -1,6 +1,6 @@ from firedrake import * from firedrake.interpolation import ( - MixedInterpolator, SameMeshInterpolator, CrossMeshInterpolator, + SameMeshInterpolator, CrossMeshInterpolator, get_interpolator, VomOntoVomInterpolator, ) from firedrake.matrix import ImplicitMatrix, Matrix @@ -73,9 +73,12 @@ def test_same_mesh_mattype(value_shape, mat_type, mode): res2 = assemble(action(adjoint(forward_I_mat), f)) assert np.allclose(res2.dat.data, exact.dat.data) - with pytest.raises(NotImplementedError): - # MatNest only implemented for interpolation between MixedFunctionSpaces - assemble(interp, mat_type="nest") + # A nest of the one block these unmixed spaces make collapses to that + # block's own type, just as it does for a Form. + nest_mat = assemble(interp, mat_type="nest") + assert nest_mat.petscmat.type == prefix + ("baij" if value_shape == "vector" else "aij") + res = assemble(action(nest_mat, f)) + assert np.allclose(res.dat.data, exact.dat.data) @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") @@ -188,7 +191,7 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): expr = as_vector([x**2, x**2, y**2, y**2]) interp = interpolate(TrialFunction(U), W) - assert isinstance(get_interpolator(interp), MixedInterpolator) + assert isinstance(get_interpolator(interp), SameMeshInterpolator) I_mat = assemble(interp, mat_type=mat_type, sub_mat_type=sub_mat_type) assert isinstance(I_mat, ImplicitMatrix if mat_type == "matfree" else Matrix) @@ -198,17 +201,19 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): assert I_mat.petscmat.type == "seqaij" else: assert I_mat.petscmat.type == "nest" + if value_shape == "scalar": + # Always seqaij for scalar + sub_type = "seqaij" + else: + # A blocked space makes matnest default to baij + sub_type = "seq" + (sub_mat_type if sub_mat_type else "baij") for (i, j) in [(0, 0), (0, 1), (1, 0), (1, 1)]: + # Every block is assembled, as it is for a Form. The components + # do not mix, so the off-diagonal ones assemble to zero. sub_mat = I_mat.petscmat.getNestSubMatrix(i, j) + assert sub_mat.type == sub_type if i != j: - assert not sub_mat - continue - if value_shape == "scalar": - # Always seqaij for scalar - assert sub_mat.type == "seqaij" - else: - # matnest sub_mat_type defaults to aij - assert sub_mat.type == "seq" + (sub_mat_type if sub_mat_type else "aij") + assert sub_mat.norm() == 0.0 f = Function(U).interpolate(expr) exact = Function(W).interpolate(expr) @@ -216,5 +221,6 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): for resi, exi in zip(res.subfunctions, exact.subfunctions): assert np.allclose(resi.dat.data, exi.dat.data) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): + # A mixed space has no block structure to give BAIJ, as for a Form. assemble(interp, mat_type="baij") From e11294d5da6a1fd8a1f42f8fa039c32db882b15d Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 27 Aug 2026 12:49:47 +0100 Subject: [PATCH 08/10] Leave a Coargument out of the Real mangling A Real-space argument is replaced by the literal 1 before a form reaches TSFC, because the Real element has a single global degree of freedom whose basis function is constant. That reasoning applies to a primal argument the form integrates against. An Interpolate's Coargument instead names the space the result lands in, and TSFC dual-evaluates it, so replacing it with 1 hands Interpolate.__init__ an IntValue where a space belongs. Mangle before the empty check, as the components are split before they are mangled and an empty component carries no Real argument either way. --- firedrake/tsfc_interface.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index 73f6772e41..1585686d8e 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -237,12 +237,12 @@ def compile_form(form, name, parameters=None, split=True, dont_split=(), nargs = 1 iterable = ([(None, )*nargs, form], ) for idx, f in iterable: + f = _real_mangle(f) if isinstance(f, ZeroBaseForm) or f.empty(): # If we're assembling the R space component of a mixed argument, # and that component doesn't actually appear in the form then we # have an empty form, which we should not attempt to assemble. continue - f = _real_mangle(f) # Map local domain/coefficient/constant numbers (as seen inside the # compiler) to the global coefficient/constant numbers meshes = extract_domains(f) @@ -277,7 +277,10 @@ def _real_mangle(form): """If the form contains arguments in the Real function space, replace these with literal 1 before passing to tsfc.""" a = form.arguments() - reals = [x.ufl_element().family() == "Real" for x in a] + # A Coargument names the space the result lands in rather than something to + # integrate against, so TSFC dual-evaluates it instead. + reals = [x.ufl_element().family() == "Real" and not isinstance(x, ufl.Coargument) + for x in a] if not any(reals): return form replacements = {} From 065316fea04ac1f7d26b32552339ea31cc85194c Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 27 Aug 2026 15:02:10 +0100 Subject: [PATCH 09/10] Count the coordinates among an interpolation's inputs An interpolation whose output is the mesh coordinate field reads those same coordinates through the kernel's coordinate argument, so it has to compute into a tensor of its own and copy out. The coordinates never appear among the form's coefficients, so the check for that overlap missed them and the interpolation overwrote, cell by cell, the coordinates it was still reading. Building the octahedral sphere is the case that notices: it remaps its own coordinates twice, and the second remap collapsed cells it had already written. Co-Authored-By: Claude Opus 5 --- firedrake/interpolation.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index 311c329bb8..978bd950cf 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -863,6 +863,9 @@ def _get_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None) ) if isinstance(self.dual_arg, Cofunction): inputs += (self.dual_arg,) + # The coordinates reach the kernel without being coefficients of the form. + inputs += (self.source_mesh.unique().coordinates, + self.target_mesh.unique().coordinates) if ( isinstance(tensor, Function | Cofunction) and any(set(tensor.dat).intersection(set(input_.dat)) From a620b1c3406729b71ce5fc313adffc5078ad69e5 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 20:16:47 +0100 Subject: [PATCH 10/10] DROP BEFORE MERGE --- .github/actions/install/action.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index ddd9d583b1..5089bef54b 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -164,6 +164,8 @@ runs: --extra-index-url https://download.pytorch.org/whl/cpu \ "./firedrake-repo[${{ inputs.deps }}]" + pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/fiat.git@pbrubeck/interp-mixed + pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/ufl.git@pbrubeck/interpolate-holes firedrake-clean pip list