From f9ec6d574b54f8572a8ff8ba23390deb39507441 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 16 Jul 2026 09:43:28 +0100 Subject: [PATCH 1/8] 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 | 90 +------------ tsfc/fem.py | 124 +++++++++++++++++- tsfc/modified_terminals.py | 8 +- tsfc/ufl_utils.py | 33 ++++- 8 files changed, 343 insertions(+), 100 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index de62cc1d29..f6fb08f064 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 @@ -77,6 +92,29 @@ toolchain: * **Docstrings:** All public-facing APIs must include properly formatted `numpydoc`-style docstrings. * **Type Hints:** New code should include type hints on function/method signatures. +## 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 @@ -104,6 +142,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 @@ -140,12 +184,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. @@ -156,7 +213,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) @@ -165,6 +222,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 6ee929eced..e8c6139f39 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 @@ -21,7 +20,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,29 +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, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) - - # 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) - evaluation = gem.IndexSum(evaluation * gem_dual[basis_indices], basis_indices) - basis_indices = () - else: + evaluation, 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())) @@ -384,64 +363,3 @@ def predicate(index): 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..761f82f01f 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,75 @@ 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 + GEM expression for the local interpolated values and its basis indices. + """ + 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, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) + + 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) + evaluation = gem.IndexSum(evaluation * gem_dual[basis_indices], basis_indices) + basis_indices = () + + return evaluation, 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 +813,38 @@ 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, basis_indices = dual_evaluate(terminal, element, kernel_cfg) + 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 +872,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 +900,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 cf47fb8974..41e67c4c6c 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 @@ -21,7 +22,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 @@ -35,6 +36,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 00aa3ee39655ad4b689939e836306047710e90e4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Thu, 16 Jul 2026 15:16:45 +0100 Subject: [PATCH 2/8] FormAssembler: handle same-mesh interpolation --- firedrake/assemble.py | 200 +++++++++++++++++++---- firedrake/interpolation.py | 143 +++++++++++++++- pyop2/parloop.py | 12 +- tsfc/driver.py | 3 + tsfc/kernel_interface/firedrake_loopy.py | 10 +- tsfc/ufl_utils.py | 16 ++ 6 files changed, 336 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 52b4ff7020..5414da48b3 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() @@ -709,6 +707,68 @@ def __init__(self, expr): # 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. @@ -770,7 +830,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 @@ -829,7 +964,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/driver.py b/tsfc/driver.py index e8c6139f39..d39c5e9e9a 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -356,6 +356,9 @@ def predicate(index): # but we don't for now. evaluation, = impero_utils.preprocess_gem([evaluation]) pairs = unconcatenate([(return_expr, evaluation)]) + return_indices = tuple(dict.fromkeys( + chain(return_indices, *(variable.free_indices for variable, _ in pairs)) + )) impero_c = impero_utils.compile_gem(pairs, return_indices) index_names = {idx: f"p{i}" for (i, idx) in enumerate(basis_indices)} # Handle kernel interface requirements 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 41e67c4c6c..dc4e6b807c 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -414,6 +414,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 5973b7b6aa1bd5af73b7801f55bcc61293e29797 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 10:26:46 +0100 Subject: [PATCH 3/8] unify Interpolate and Form assembly --- firedrake/assemble.py | 222 +++++------------- firedrake/bcs.py | 24 ++ firedrake/formmanipulation.py | 34 ++- firedrake/interpolation.py | 76 +----- 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 | 89 ++++++- tsfc/fem.py | 64 ++++- tsfc/kernel_interface/__init__.py | 4 + tsfc/kernel_interface/common.py | 11 + tsfc/kernel_interface/firedrake_loopy.py | 11 +- tsfc/ufl_utils.py | 7 +- 15 files changed, 422 insertions(+), 261 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 0e32a3d0b5..707cc33a88 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 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.functionspaceimpl import WithGeometry, FunctionSpace, FiredrakeDualSpace @@ -43,17 +43,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 +172,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 +1032,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 +1128,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 +1145,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 +1233,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 +1309,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 +1400,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 @@ -1598,8 +1524,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 +1533,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 +1658,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) @@ -1858,8 +1777,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 @@ -1902,7 +1819,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(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]) @@ -2166,18 +2083,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 +2095,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 +2129,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 +2175,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: @@ -2466,16 +2376,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 60ab31e179..affae4cbc8 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 5414da48b3..e1007b2100 100644 --- a/firedrake/interpolation.py +++ b/firedrake/interpolation.py @@ -151,6 +151,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. @@ -737,6 +742,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 @@ -746,7 +752,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: @@ -843,6 +849,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( @@ -851,7 +858,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( @@ -859,14 +866,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, ) @@ -905,63 +904,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} 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 d39c5e9e9a..b4fc8ea6e1 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -9,7 +9,7 @@ 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 @@ -18,7 +18,7 @@ import finat from finat.element_factory import as_fiat_cell -from tsfc import fem, ufl_utils +from tsfc import fem, kernel_args, ufl_utils from tsfc.logging import logger from tsfc.parameters import default_parameters, is_complex from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants, simplify_abs @@ -76,6 +76,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 +113,88 @@ 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 into a cell-iteration assembly kernel.""" + parameters = preprocess_parameters(parameters) + 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() + + expression_kernel = compile_expression_dual_evaluation( + expression, + expression.ufl_element(), + domain=source_domain, + parameters=parameters, + name=prefix, + ) + + arguments = list(expression_kernel.arguments) + active_domain = all_domains.index(source_domain) + 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 + ) + + coefficients = expression.coefficients() + coefficient_numbers = [] + for number in expression_kernel.coefficient_numbers: + coefficient = coefficients[number] + if type(coefficient.ufl_element()) is finat.ufl.MixedElement: + subindices = tuple(range(len(coefficient.ufl_element().sub_elements))) + else: + subindices = (0,) + coefficient_numbers.append((number, subindices)) + + active_domain_numbers = firedrake_interface_loopy.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=(), + ) + return [ + firedrake_interface_loopy.Kernel( + ast=expression_kernel.ast, + arguments=tuple(arguments), + integral_type="cell", + subdomain_id=("everywhere",), + domain_number=all_domains.index(target_domain), + active_domain_numbers=active_domain_numbers, + coefficient_numbers=tuple(coefficient_numbers), + tabulations=expression_kernel.tabulations, + flop_count=expression_kernel.flop_count, + name=expression_kernel.name, + event=expression_kernel.event, + ) + ] + + def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=False): """Compiles a UFL integral into an assembly kernel. diff --git a/tsfc/fem.py b/tsfc/fem.py index 761f82f01f..cdcee98a2f 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. @@ -358,14 +364,39 @@ def dual_evaluate(expression: ufl.Interpolate, to_element: FiniteElementBase, ke coordinate_mapping = CoordinateMapping(analyse_modified_terminal(coefficient), ctx) else: coordinate_mapping = None - evaluation, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) - 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) - evaluation = gem.IndexSum(evaluation * gem_dual[basis_indices], 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) + evaluations = [] + for element, gem_dual in zip(to_element.elements, gem_duals): + evaluation, 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], basis_indices + )) + evaluation = gem.optimise.make_sum(evaluations) basis_indices = () + else: + evaluation, basis_indices = to_element.dual_evaluation( + fn, coordinate_mapping + ) + if gem_duals: + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + evaluation = gem.IndexSum( + evaluation * gem_duals[0][basis_indices], basis_indices + ) + basis_indices = () return evaluation, basis_indices @@ -815,7 +846,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( @@ -891,8 +926,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) @@ -904,7 +943,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..0cb09a5fa8 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -41,6 +41,7 @@ def __init__(self, scalar_type): # Coefficients self.coefficient_map = collections.OrderedDict() + self.coefficient_split = {} # Constants self.constant_map = collections.OrderedDict() @@ -63,6 +64,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] 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/ufl_utils.py b/tsfc/ufl_utils.py index dc4e6b807c..ddbbbf82b6 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -43,8 +43,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 86a425c13098278b28c58c6360c8fdde12fdf2e1 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 12:02:51 +0100 Subject: [PATCH 4/8] Plate Bending Demo --- .../plate_bending_mitc.py.rst | 155 ++++++++++++++++++ docs/source/advanced_tut.rst | 1 + tests/firedrake/demos/test_demos_run.py | 1 + 3 files changed, 157 insertions(+) create mode 100644 demos/plate_bending_mitc/plate_bending_mitc.py.rst diff --git a/demos/plate_bending_mitc/plate_bending_mitc.py.rst b/demos/plate_bending_mitc/plate_bending_mitc.py.rst new file mode 100644 index 0000000000..33c1b48f94 --- /dev/null +++ b/demos/plate_bending_mitc/plate_bending_mitc.py.rst @@ -0,0 +1,155 @@ +Plate Bending Using Mixed Interpolation of Tensorial Components (MITC) +====================================================================== + +This demo illustrates how to solve a Reissner-Mindlin plate problem using the +Mixed Interpolation of Tensorial Components (MITC) formulation. The main goal of +the MITC method is to prevent shear locking in the thin-plate limit (:math:`t \to 0`) +by projecting the rotation field into an edge-conforming Nédélec space. + +Background and Formulation +-------------------------- + +We model the plate using the Reissner-Mindlin equations. Given a domain +:math:`\Omega` representing the plate's mid-surface, we seek the transverse deflection +:math:`w` and the rotations :math:`\boldsymbol{\beta}` of the mid-surface normal. + +To avoid shear locking, the standard shear strain term :math:`(\nabla w - \boldsymbol{\beta})` +is modified. We project the rotation field :math:`\boldsymbol{\beta}` onto an +:math:`H(\text{curl})`-conforming space :math:`\boldsymbol{R}_h` (a Nédélec space) using a reduction +operator :math:`\Pi_h`. The shear strain term is then evaluated as: + +.. math:: + + \bar{\boldsymbol{\gamma}}_h = \nabla w - \Pi_h \boldsymbol{\beta} + +Thanks to native compilation support for symbolic interpolation nodes inside +variational forms, we can define the reduction operator directly using +Firedrake's symbolic ``interpolate`` function within the UFL expression. Under the +hood, this bypasses the legacy BaseFormAssembler specifically for the case where +we have a Form with an Interpolate node onto the same mesh. + +We begin by importing the Firedrake namespace. + +:: + + from firedrake import * + +Mesh and Geometry +----------------- + +First, we set up a simple mesh of a unit square representing our plate. + +:: + + n = 16 + mesh = UnitSquareMesh(n, n) + +Material and Physical Parameters +-------------------------------- + +We set up standard parameters for a thin plate. Here, we define the thickness +:math:`t`, Young's modulus :math:`E`, Poisson's ratio :math:`\nu`, and the shear correction factor +:math:`k_s`. + +:: + + t = Constant(0.01) + E = Constant(1e3) + nu = Constant(0.3) + k_s = Constant(5.0/6.0) + +The bending stiffness :math:`D` and shear stiffness :math:`G` are derived below. + +:: + + D = E * t**3 / (12 * (1 - nu**2)) + G = E / (2 * (1 + nu)) + G_shear = k_s * G * t + +Function Spaces +--------------- + +We construct a mixed function space for the deflection and the rotation. We use +quadratic Lagrange elements for both variables. Crucially, we also define a +Nédélec space :math:`R` of degree 1 to serve as our target edge-conforming space for +the MITC projection. + +:: + + W = FunctionSpace(mesh, "Lagrange", 2) + B = VectorFunctionSpace(mesh, "Lagrange", 2) + R = FunctionSpace(mesh, "N1curl", 1) + + V = MixedFunctionSpace([W, B]) + +Variational Formulation +----------------------- + +Next, we declare the trial and test functions. We write the isotropic bending +stress tensor :math:`\boldsymbol{\sigma}(\boldsymbol{\beta})` and establish our bilinear forms. + +:: + + u = Function(V) + w, beta = TrialFunctions(V) + v, theta = TestFunctions(V) + + def sigma(phi): + return D * ((1 - nu) * sym(grad(phi)) + nu * div(phi) * Identity(2)) + + a_bending = inner(sigma(beta), sym(grad(theta))) * dx + +We implement the MITC projection using Firedrake's symbolic ``interpolate()`` +function. This acts as a true symbolic UFL operator embedded +directly within the form definition. TSFC handles the assembly +pipeline seamlessly by evaluating the node on the same mesh. + +:: + + Pi_beta = interpolate(beta, R) + Pi_theta = interpolate(theta, R) + + a_shear = G_shear * inner(grad(w) - Pi_beta, grad(v) - Pi_theta) * dx + + a = a_bending + a_shear + +We impose a uniform transverse downward load :math:`f` acting on the plate. + +:: + + f = Constant(1.0) + L = inner(f, v) * dx + +Boundary Conditions +------------------- + +We apply fully clamped boundary conditions on all boundaries, meaning both the +deflection :math:`w` and the rotation :math:`\boldsymbol{\beta}` vanish. + +:: + + bcs = [DirichletBC(V.sub(0), 0, "on_boundary"), + DirichletBC(V.sub(1), 0, "on_boundary")] + +Computation +----------- + +We solve the problem in the usual way. + +:: + + solve(a == L, u, bcs=bcs) + +We recover the split deflection and rotation solutions for analysis. + +:: + + w_sol, beta_sol = u.subfunctions + max_w = w_sol.dat.data.max() + print(f"Max deflection: {max_w:.6e}") + +Finally, we output the deflection to a PVD file for visualization in ParaView. + +:: + + VTKFile("mitc_plate.pvd").write(w_sol) diff --git a/docs/source/advanced_tut.rst b/docs/source/advanced_tut.rst index 6c6dc19e46..544accc181 100644 --- a/docs/source/advanced_tut.rst +++ b/docs/source/advanced_tut.rst @@ -38,3 +38,4 @@ element systems. Coupled volume-surface reaction-diffusion on a torus using Submesh and geometric multigrid. Nonlinear preconditioning using an auxiliary SNES for the Allen-Cahn equation. Reynolds-robust preconditioning of the stationary Navier-Stokes equations. + A plate bending problem solved with Mixed Interpolation of Tensorial Components. diff --git a/tests/firedrake/demos/test_demos_run.py b/tests/firedrake/demos/test_demos_run.py index f53a611c2b..76e0a7c6d2 100644 --- a/tests/firedrake/demos/test_demos_run.py +++ b/tests/firedrake/demos/test_demos_run.py @@ -56,6 +56,7 @@ Demo(('submesh_reaction_diffusion', 'submesh_reaction_diffusion'), ["netgen", "vtk"]), Demo(('nonlinear_pc', 'nonlinear_pc_allen_cahn'), []), Demo(('reynolds_robust_navier_stokes_hdiv', 'reynolds_robust_navier_stokes_hdiv'), ["vtk"]), + Demo(('plate_bending_mitc', 'plate_bending_mitc'), ["vtk"]), ] PARALLEL_DEMOS = [ Demo(("full_waveform_inversion", "full_waveform_inversion"), ["adjoint"]), From f7f5ee3257362c0fb2aaf42e1f8245f95fc172ec Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 15:42:00 +0100 Subject: [PATCH 5/8] merge Form and Interpolate codegen --- firedrake/assemble.py | 77 ++++++++-- firedrake/interpolation.py | 239 +------------------------------- tsfc/driver.py | 202 ++++++++++++++++++--------- tsfc/kernel_interface/common.py | 3 +- tsfc/spectral.py | 4 +- 5 files changed, 218 insertions(+), 307 deletions(-) diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 707cc33a88..08df0ce215 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -22,9 +22,10 @@ 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 @@ -1472,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()] @@ -1492,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()} @@ -1689,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`. @@ -1804,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) @@ -1822,7 +1849,14 @@ def _make_mat_global_kernel_arg(self, Vrow, 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) @@ -1924,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) @@ -2198,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 @@ -2278,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) diff --git a/firedrake/interpolation.py b/firedrake/interpolation.py index e1007b2100..50eab90fcd 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 @@ -175,6 +164,8 @@ def _interpolator(self): try: source_mesh = extract_unique_domain(operand) or target_mesh except ValueError: + if has_mixed_arguments or len(self.target_space) > 1: + return MixedInterpolator(self) raise NotImplementedError( "Interpolating an expression with no arguments defined on multiple meshes is not implemented yet." ) @@ -1050,204 +1041,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. @@ -1283,28 +1076,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/tsfc/driver.py b/tsfc/driver.py index b4fc8ea6e1..021ed4d76f 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -17,8 +17,12 @@ 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, kernel_args, ufl_utils +from tsfc import fem, ufl_utils +from tsfc.kernel_interface.common import pick_mode from tsfc.logging import logger from tsfc.parameters import default_parameters, is_complex from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants, simplify_abs @@ -52,6 +56,21 @@ """ +TSFCInterpolationData = collections.namedtuple( + "TSFCInterpolationData", + ["domain", "iteration_domain", "integral_type", "subdomain_id", + "domain_integral_type_map", "enabled_coefficients", "integrals", + "expression", "target_element"], +) + +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. @@ -114,8 +133,11 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di def compile_interpolate(expression, prefix="interpolate", parameters=None): - """Compile an interpolation into a cell-iteration assembly kernel.""" + """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): @@ -131,66 +153,77 @@ def compile_interpolate(expression, prefix="interpolate", parameters=None): ) all_domains = expression.ufl_domains() - expression_kernel = compile_expression_dual_evaluation( - expression, - expression.ufl_element(), - domain=source_domain, - parameters=parameters, - name=prefix, - ) + 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 + ) - arguments = list(expression_kernel.arguments) - active_domain = all_domains.index(source_domain) - coefficient_offset = ( - 1 - + expression_kernel.oriented - + expression_kernel.needs_cell_sizes + operand = apply_mapping(operand, target_element, source_domain) + operand = ufl_utils.preprocess_expression( + operand, complex_mode=complex_mode ) - if expression_kernel.needs_external_coords: - coordinate_arg = arguments[coefficient_offset] - arguments[coefficient_offset] = kernel_args.CoordinatesKernelArg( - coordinate_arg.loopy_arg - ) + operand = simplify_abs(operand, complex_mode) + expression = ufl.Interpolate(operand, dual_arg) coefficients = expression.coefficients() - coefficient_numbers = [] - for number in expression_kernel.coefficient_numbers: - coefficient = coefficients[number] - if type(coefficient.ufl_element()) is finat.ufl.MixedElement: - subindices = tuple(range(len(coefficient.ufl_element().sub_elements))) - else: - subindices = (0,) - coefficient_numbers.append((number, subindices)) - - active_domain_numbers = firedrake_interface_loopy.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 () + 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 ), - exterior_facets=(), - interior_facets=(), - orientations_cell=(), - orientations_exterior_facet=(), - orientations_interior_facet=(), + 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, + target_element=target_element, ) return [ - firedrake_interface_loopy.Kernel( - ast=expression_kernel.ast, - arguments=tuple(arguments), - integral_type="cell", - subdomain_id=("everywhere",), - domain_number=all_domains.index(target_domain), - active_domain_numbers=active_domain_numbers, - coefficient_numbers=tuple(coefficient_numbers), - tabulations=expression_kernel.tabulations, - flop_count=expression_kernel.flop_count, - name=expression_kernel.name, - event=expression_kernel.event, + compile_integral( + integral_data, form_data, prefix, parameters, diagonal=False ) ] @@ -227,8 +260,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, @@ -257,12 +295,52 @@ 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): + expression = CoefficientSplitter(builder.coefficient_split)( + integral_data.expression + ) + target_element = builder.create_element(integral_data.target_element) + config = builder.fem_config() + config.update( + argument_multiindices=builder.argument_multiindices, + index_cache=ctx["index_cache"], + ) + if isinstance(target_element, finat.QuadratureElement): + config["quadrature_rule"] = target_element._rule + evaluation, basis_indices = fem.dual_evaluate( + expression, target_element, config + ) + dual_arg, _ = expression.argument_slots() + if not isinstance(dual_arg, ufl.Cofunction): + arguments = expression.arguments() + argument_number = arguments.index(dual_arg) + output_indices = builder.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)) + ) + 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) + params["mode"] = "vanilla" + mode = pick_mode(params["mode"]) + representations = mode.Integrals( + [evaluation], (), builder.argument_multiindices, params + ) + builder.stash_integrals(representations, 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/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 0cb09a5fa8..020aeb77e8 100644 --- a/tsfc/kernel_interface/common.py +++ b/tsfc/kernel_interface/common.py @@ -588,7 +588,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/spectral.py b/tsfc/spectral.py index 69e471104e..79474b9013 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 @@ -129,7 +129,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 From 1d4637867dc25a76c771b3a4a490ea1c31647c5b Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 16:13:55 +0100 Subject: [PATCH 6/8] Tidy --- tsfc/driver.py | 42 +++-------------------------- tsfc/kernel_interface/common.py | 48 +++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+), 39 deletions(-) diff --git a/tsfc/driver.py b/tsfc/driver.py index 021ed4d76f..1491fcc4b0 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -22,7 +22,6 @@ from finat.ufl import FiniteElement, TensorElement from tsfc import fem, ufl_utils -from tsfc.kernel_interface.common import pick_mode from tsfc.logging import logger from tsfc.parameters import default_parameters, is_complex from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants, simplify_abs @@ -60,7 +59,7 @@ "TSFCInterpolationData", ["domain", "iteration_domain", "integral_type", "subdomain_id", "domain_integral_type_map", "enabled_coefficients", "integrals", - "expression", "target_element"], + "expression"], ) TSFCInterpolationFormData = collections.namedtuple( @@ -219,7 +218,6 @@ def compile_interpolate(expression, prefix="interpolate", parameters=None): enabled_coefficients=(True,) * len(coefficients), integrals=(), expression=expression, - target_element=target_element, ) return [ compile_integral( @@ -296,44 +294,10 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F builder.set_constants(form_data.constants) ctx = builder.create_context() if isinstance(integral_data, TSFCInterpolationData): - expression = CoefficientSplitter(builder.coefficient_split)( - integral_data.expression - ) - target_element = builder.create_element(integral_data.target_element) - config = builder.fem_config() - config.update( - argument_multiindices=builder.argument_multiindices, - index_cache=ctx["index_cache"], - ) - if isinstance(target_element, finat.QuadratureElement): - config["quadrature_rule"] = target_element._rule - evaluation, basis_indices = fem.dual_evaluate( - expression, target_element, config - ) - dual_arg, _ = expression.argument_slots() - if not isinstance(dual_arg, ufl.Cofunction): - arguments = expression.arguments() - argument_number = arguments.index(dual_arg) - output_indices = builder.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)) - ) - params = parameters.copy() params["mode"] = "vanilla" - mode = pick_mode(params["mode"]) - representations = mode.Integrals( - [evaluation], (), builder.argument_multiindices, params - ) - builder.stash_integrals(representations, 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() diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py index 020aeb77e8..17c0b0eb2a 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 @@ -147,6 +150,51 @@ 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, 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"]) + return mode.Integrals( + [evaluation], (), self.argument_multiindices, params + ) + def compile_integrand(self, integrand, params, ctx): """Compile UFL integrand. From 1162fee79a24971ec79b83bf91fee640a590c6c2 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 20:14:08 +0100 Subject: [PATCH 7/8] tweak demo --- demos/plate_bending_mitc/elements.svg | 325 ++++++++++++++++++ demos/plate_bending_mitc/elements.tex | 38 ++ .../plate_bending_mitc.py.rst | 33 +- 3 files changed, 388 insertions(+), 8 deletions(-) create mode 100644 demos/plate_bending_mitc/elements.svg create mode 100644 demos/plate_bending_mitc/elements.tex diff --git a/demos/plate_bending_mitc/elements.svg b/demos/plate_bending_mitc/elements.svg new file mode 100644 index 0000000000..618a7b0d67 --- /dev/null +++ b/demos/plate_bending_mitc/elements.svg @@ -0,0 +1,325 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + P1 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + P1-iso-P2 + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Nedelec1 + + + + + + + + + + + + diff --git a/demos/plate_bending_mitc/elements.tex b/demos/plate_bending_mitc/elements.tex new file mode 100644 index 0000000000..915e47d0a5 --- /dev/null +++ b/demos/plate_bending_mitc/elements.tex @@ -0,0 +1,38 @@ +\documentclass[tikz,border=2pt]{standalone} +\usepackage[utf8]{inputenc} +\begin{document} +\begin{tikzpicture} + + % P1 Element - Filled + \begin{scope}[shift={(0,0)}] + \filldraw[fill=blue!20, draw=black, thick] (0,0) -- (2,0) -- (1,1.732) -- cycle; + \foreach \p in {(0,0), (2,0), (1,1.732)} \fill[black] \p circle (2pt); + \node at (1,-0.5) {P1}; + \end{scope} + + % P1-iso-P2 Element - Filled + \begin{scope}[shift={(4,0)}] + \filldraw[fill=green!20, draw=black, thick] (0,0) -- (2,0) -- (1,1.732) -- cycle; + \draw[draw=black, thick] (1,0) -- (1.5,0.866) -- (0.5,0.866) -- cycle; + \draw[draw=black, thick] (0,0) -- (0.5,0.866); + \draw[draw=black, thick] (2,0) -- (1.5,0.866); + \draw[draw=black, thick] (1,1.732) -- (0.5,0.866); + \draw[draw=black, thick] (1,1.732) -- (1.5,0.866); + \draw[draw=black, thick] (1,0) -- (0.5,0.866); + \draw[draw=black, thick] (1,0) -- (1.5,0.866); + \foreach \p in {(0,0), (2,0), (1,1.732), (1,0), (0.5,0.866), (1.5,0.866)} \fill[black] \p circle (2pt); + \node at (1,-0.5) {P1-iso-P2}; + \end{scope} + + % Nédélec (1st kind) Element - Filled + \begin{scope}[shift={(8,0)}] + \filldraw[fill=red!20, draw=black, thick] (0,0) -- (2,0) -- (1,1.732) -- cycle; + % Vectors centered along the edges, always black + \draw[->, ultra thick, black] (0.25, 0) -- (1.75, 0); + \draw[->, ultra thick, black] (1.75, 0.433) -- (1.25, 1.299); + \draw[->, ultra thick, black] (0.75, 1.299) -- (0.25, 0.433); + \node at (1,-0.5) {Nedelec 1}; + \end{scope} + +\end{tikzpicture} +\end{document} diff --git a/demos/plate_bending_mitc/plate_bending_mitc.py.rst b/demos/plate_bending_mitc/plate_bending_mitc.py.rst index 33c1b48f94..ca7cd92b39 100644 --- a/demos/plate_bending_mitc/plate_bending_mitc.py.rst +++ b/demos/plate_bending_mitc/plate_bending_mitc.py.rst @@ -22,11 +22,26 @@ operator :math:`\Pi_h`. The shear strain term is then evaluated as: \bar{\boldsymbol{\gamma}}_h = \nabla w - \Pi_h \boldsymbol{\beta} +Element Spaces +-------------- + +The following diagram illustrates the element spaces used: + +.. image:: elements.svg + :align: center + +* **P1**: Standard linear Lagrange element for the deflection :math:`w`. +* **P1-iso-P2**: Constructed as macroelement, where the master triangle is divided into four smaller sub-triangles. This structure is utilized for the rotation field :math:`\boldsymbol{\beta}`. +* **Nédélec 1**: An :math:`H(\text{curl})`-conforming space used for the MITC projection. The blue arrows represent the edge-based degrees of freedom. + +Implementation +-------------- + Thanks to native compilation support for symbolic interpolation nodes inside variational forms, we can define the reduction operator directly using Firedrake's symbolic ``interpolate`` function within the UFL expression. Under the -hood, this bypasses the legacy BaseFormAssembler specifically for the case where -we have a Form with an Interpolate node onto the same mesh. +hood, this bypasses ``BaseFormAssembler`` specifically for the case where +we have a ``Form`` with an ``Interpolate`` node onto the same mesh. We begin by importing the Firedrake namespace. @@ -70,15 +85,17 @@ Function Spaces --------------- We construct a mixed function space for the deflection and the rotation. We use -quadratic Lagrange elements for both variables. Crucially, we also define a -Nédélec space :math:`R` of degree 1 to serve as our target edge-conforming space for -the MITC projection. +linear Lagrange elements for the deflection and P1-iso-P2 elements for the rotation. +Crucially, we also define a Nédélec space :math:`R` of degree 1 to serve as our +target edge-conforming space for the MITC projection. Because the facets are +split by the P1-iso-P2 macro-triangulation, the Nédélec integral moments must +employ a composite quadrature scheme (``quad_scheme="iso"``). :: - W = FunctionSpace(mesh, "Lagrange", 2) - B = VectorFunctionSpace(mesh, "Lagrange", 2) - R = FunctionSpace(mesh, "N1curl", 1) + W = FunctionSpace(mesh, "Lagrange", 1) + B = VectorFunctionSpace(mesh, "Lagrange", 1, variant="iso") + R = FunctionSpace(mesh, "N1curl", 1, quad_scheme="iso") V = MixedFunctionSpace([W, B]) From e8ad626408c6f92c046bb2868580055c3664dfa4 Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Fri, 17 Jul 2026 20:16:47 +0100 Subject: [PATCH 8/8] 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 8ba3dbbe57..813f6774cd 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -152,6 +152,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