diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index ddd9d583b1..5089bef54b 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -164,6 +164,8 @@ runs: --extra-index-url https://download.pytorch.org/whl/cpu \ "./firedrake-repo[${{ inputs.deps }}]" + pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/fiat.git@pbrubeck/interp-mixed + pip install -v --no-deps --ignore-installed git+https://github.com/firedrakeproject/ufl.git@pbrubeck/interpolate-holes firedrake-clean pip list diff --git a/AGENTS.md b/AGENTS.md index ffc2ca2295..c5568c4e5e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,8 +1,10 @@ # Firedrake Firedrake is an automated system for the portable solution of partial differential equations using -the finite element method (FEM). The codebase is primarily Python, relying heavily on code generation -and high-performance C backends to achieve scalability and speed. +the finite element method. The codebase is primarily Python, relying heavily on code generation +and high-performance C backends to achieve scalability and speed. Firedrake is highly composable and +fully differentiable, enabling the automatic generation of tangent linear and adjoint models for +PDE-constrained optimization. Firedrake's full contribution process is documented at [Contributing to Firedrake](https://firedrakeproject.org/contribute.html). In short, for AI-assisted @@ -30,6 +32,12 @@ toolchain: 2. **Lowering to Loopy:** The GEM expressions are then lowered into **loopy** kernels. * **PyOP2:** Finally, the generated loopy kernels are wrapped and executed by PyOP2, which handles the parallel execution of loops over mesh cells and facets. +* **pyadjoint:** Firedrake integrates with `pyadjoint` for algorithmic differentiation. Annotated + operations (assembly, interpolation, variational solves, boundary condition application, ...) are + recorded on a `Tape` as a DAG of `Block`s; a `ReducedFunctional` composed with one or more `Control`s + can then be evaluated, differentiated (adjoint or tangent-linear), and checked with a `taylor_test`. + Taping happens at the level of Firedrake's own operations, not the underlying numerics, so any new + feature assembled purely from already-annotated building blocks is differentiable automatically. ## Core Working Rules @@ -56,6 +64,13 @@ toolchain: prose explaining what the removed, incorrect approach used to do or why it was wrong. Keep comments and documentation focused on the current, correct code; a reader should never need the history of what used to be there to understand why the present code is right. +* **Composability And Differentiability:** New features are expected to compose with existing ones + without special-casing, and, via `pyadjoint`, to remain differentiable when built from already-taped + operations. Prefer the annotated, top-level API (e.g. `firedrake.assemble`) over a lower-level + equivalent that bypasses it (e.g. calling an `Interpolator`'s own `.assemble()` directly) even when + both give the same forward numbers — the lower-level call silently drops out of the tape, and no test + will notice unless it specifically exercises pyadjoint. When a change could plausibly sit on the tape, + verify differentiability explicitly with a `taylor_test`, not just a forward-value check. ## Coding Style And Conventions @@ -90,6 +105,29 @@ toolchain: excluded from that rule, so its snippet renders in the docs but never runs. Prefer `::` — reach for the directive only for an illustrative fragment naming things the demo never defines. +## Design And Debugging Method + +How to spend effort on any feature or bug: the first three shape a design before writing code, +the last two localize a failure before reading code. + +* **Design by nearest working neighbor.** Some existing feature already solves a structurally + identical problem. Grep for the invariant yours must satisfy (the type handled, the hook fired, + the kwarg accepted), read how the neighbor earns it, and implement only the delta. +* **Write the state contract before the code.** For anything flowing through a cached, replayed, + or lazily-refreshed system, answer up front: who owns it, when is it refreshed, what happens on + reuse or in-place mutation, how does a replay recover it. Stale-state bugs fail far from their + cause and only under reuse. +* **Classify the mathematical structure before choosing machinery.** Affine and linear + dependencies have closed-form contributions (identities, fixed operators, exact zeros for higher + derivatives); recognizing an exact zero lets you skip a code path instead of fixing it. +* **Attribute before you analyze.** Shrink *where* with one-delta experiments — each ingredient + toggled in isolation against a known-good baseline, consecutive CI failure sets diffed, a single + hunk reverted, the merge base rerun — then read code for *why*. +* **Census the consumers before changing a lifecycle.** Changing *when* or *how often* something + is computed (rather than its value) is safe only after grepping every access site: some consumer + calls it in a context you did not design for (per nonlinear iteration, inside a PETSc callback, + on every attribute read). + ## Testing Requirements * **Pull Requests:** All PRs must include comprehensive tests demonstrating that the new feature works @@ -117,6 +155,12 @@ toolchain: in the install docs to get a component installed in editable mode so source edits take effect without reinstalling, and check which branch/commit of each component is actually active before assuming a fix belongs in Firedrake itself. +* **Branch pairing across the stack:** Firedrake's `main` and `release` branches go hand in hand with + the `main` and `release` branches of its components (FIAT, UFL, ...). A CI failure may be + unreproducible locally simply because a component checkout in the venv sits on some other branch — + check `git -C $VIRTUAL_ENV/src/ branch`, switch to the branch matching the Firedrake branch + under test, run `firedrake-clean`, and reproduce again before hunting for the bug in Firedrake + itself. * **`petsc4py`/PETSc version skew:** `petsc4py` is a compiled extension built against one specific PETSc checkout. If you switch the PETSc branch/commit underneath an existing venv (e.g. to bisect a PETSc-side issue) without rebuilding `petsc4py` against it, `import firedrake` fails with a confusing @@ -153,12 +197,25 @@ toolchain: conclude a parallel code path is untested just because a plain, unmarked `pytest` run was green. * **Splitting for CI:** `firedrake-run-split-tests` shards the suite by process count for CI; look at it (and `.github/workflows/pr.yml`/`core.yml`) if a failure only reproduces in CI and not locally. +* **CI triage:** `gh pr checks ` lists job statuses. When `gh run view --log` returns nothing + (it does for very large logs), download the log with + `gh api repos///actions/jobs//logs` and grep for `FAILED`. Before debugging + anything, fetch the failure list of the *previous* run of the same PR: the difference between the + two failure sets attributes each failure to the commits pushed in between. * **Narrow reproduction first:** Run the single failing test node (`pytest path::test_name -k ...`) before the full module; the suite is large and full-module reruns are slow to iterate against. +* **Test mathematical correctness, not just that it runs or looks structurally right.** Neither "no + exception was raised" nor `==` agreement between two expressions proves the result is + correct — two independently-built expressions can match structurally while sharing the same wrong + derivative or simplification rule. Verify the actual mathematical claim: evaluate numerically and + compare against a hand-computed or finite-difference value, or use a Taylor test for anything + claiming to be a derivative. +* **Taylor-test-everything is the immune system:** Taylor-test a `ReducedFunctional`, ensuring that any + new feature built from existing, annotated Firedrake operations is automatically differentiable. ### Debugging -* **Generated kernels (niche, rarely needed):** By default, generated C is compiled optimized and +* **Generated kernels:** By default, generated C is compiled optimized and without debug symbols, so a debugger attached to the Python process cannot meaningfully step through it. Set `PYOP2_DEBUG=1` to compile with `-O0 -g` instead, which is the prerequisite for using `gdb`/`cgdb` on the compiled kernel at all. @@ -169,7 +226,7 @@ toolchain: computed differently per rank and fed into code generation (e.g. a rank-local decision that should be a collective/global one) — make that decision the same on every rank, rather than patching the generated source or the difference itself. -* **Parallel deadlocks (niche, rarely needed):** `PYOP2_SPMD_STRICT=1` adds barriers around calls +* **Parallel deadlocks:** `PYOP2_SPMD_STRICT=1` adds barriers around calls marked `@collective` and around cache access, trading overhead for a much narrower failure point when ranks disagree about control flow. * **Logging:** `firedrake.logging.set_log_level()` (or the `PYOP2_LOG_LEVEL` environment variable) @@ -178,6 +235,23 @@ toolchain: standard PETSc options (`-ksp_view`, `-snes_view`, `-ksp_monitor`, `-log_view`, `-start_in_debugger`) can be passed through Firedrake's `solver_parameters` or the command line exactly as in a plain PETSc application. +* **Errors inside PETSc callbacks do not surface as their own traceback:** under pytest they often + appear as a bare `Segmentation fault` with no Firedrake frames; standalone they appear as + `petsc4py.PETSc.Error: error code 101` whose *first* chained traceback (e.g. a `TypeError` about + the callback context in `petscsnes.pxi`) describes the corrupted callback state, not the cause. + Rerun the failing test as a standalone script to expose the chained tracebacks, and read the PETSc + call stack inside the error (`PCSetUp_MG` → `SNESComputeFunction`, ...) to identify *which* + callback was executing. +* **Construction-time code re-runs inside solver callbacks:** geometric multigrid coarsens the + entire problem lazily inside `PCSetUp`, through the `coarsen` singledispatch in + `firedrake/mg/ufl_utils.py` — whatever a feature does at construction time (e.g. `DirichletBC` + interpolating or projecting its boundary value into the space) is re-executed per level inside + that PETSc callback. Objects that carry solvers or attach DM hooks (a `Projector`, a variational + solver) must be built once and cached, never rebuilt on each call of an accessor that may fire + there: repeatedly constructing and garbage-collecting a solver stack inside `PCSetUp_MG` corrupts + the DM callback state and segfaults far from the allocation site. Extruded (hexahedral) + hierarchies are the stress test: tensor-product elements (NCE/NCF) have no dual-basis + interpolation, so paths that interpolate on simplices take the projection fallback there. ### Reproducible Environments diff --git a/docs/source/interpolation.rst b/docs/source/interpolation.rst index 50441902ff..43ca8f3147 100644 --- a/docs/source/interpolation.rst +++ b/docs/source/interpolation.rst @@ -458,7 +458,8 @@ each block given by \end{pmatrix} The off-diagonal blocks are zero since the dofs are applied component-wise. Firedrake's form -compiler recognises this and avoids assembling the zero blocks. +compiler recognises this and avoids assembling the zero blocks, which the nest +still allocates. We can assemble more general interpolation matrices between mixed function spaces by interpolating vector expressions with arguments. For example, by doing diff --git a/firedrake/assemble.py b/firedrake/assemble.py index 6dd7c4f03f..e85190e05c 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 @@ -170,7 +171,26 @@ 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): + 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) if len(form.arguments()) == 0: return ZeroFormAssembler(form, form_compiler_parameters=fc_params) @@ -893,8 +913,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 @@ -949,8 +969,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 @@ -971,9 +991,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 @@ -1012,9 +1033,11 @@ class ParloopFormAssembler(FormAssembler): Should ``tensor`` be zeroed before assembling? """ - def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True): + def __init__(self, form, bcs=None, form_compiler_parameters=None, + needs_zeroing=True, access=op2.INC): super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters) self._needs_zeroing = needs_zeroing + self._access = access def assemble(self, tensor=None, current_state=None): """Assemble the form. @@ -1106,10 +1129,11 @@ 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, slate.TensorBase): kernels = slac.compile_expression( @@ -1210,14 +1234,21 @@ class OneFormAssembler(ParloopFormAssembler): @classmethod def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=True, diagonal=False, weight=1.0): + zero_bc_nodes=True, diagonal=False, weight=1.0, access=op2.INC): bcs = solving._extract_bcs(bcs) - return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal, weight + return (tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, + zero_bc_nodes, diagonal, weight, access) @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=True, diagonal=False, weight=1.0): - super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + zero_bc_nodes=True, diagonal=False, weight=1.0, access=op2.INC): + super().__init__( + form, + bcs=bcs, + form_compiler_parameters=form_compiler_parameters, + needs_zeroing=needs_zeroing, + access=access, + ) self._weight = weight self._diagonal = diagonal self._zero_bc_nodes = zero_bc_nodes @@ -1279,6 +1310,11 @@ def _as_pyop2_type(tensor, indices=None): return tensor.dat def execute_parloops(self, tensor): + if self._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): @@ -1294,7 +1330,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()) @@ -1341,6 +1377,24 @@ def _get_mat_type(mat_type, sub_mat_type, arguments): return mat_type, sub_mat_type +def _primal_space(V): + """Return the primal space of a form argument's function space. + + Parameters + ---------- + V : firedrake.functionspaceimpl.WithGeometry + The function space of a form argument, primal or dual. + + Returns + ------- + firedrake.functionspaceimpl.WithGeometry + The primal space. An `~ufl.Interpolate` takes its test function from + the dual space, which compares unequal to the space a boundary + condition names. + """ + return V.dual() if ufl.duals.is_dual(V) else V + + class ExplicitMatrixAssembler(ParloopFormAssembler): """Class for assembling a matrix. @@ -1365,8 +1419,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 @@ -1431,8 +1491,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()] @@ -1451,8 +1517,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()} @@ -1484,7 +1554,7 @@ def _all_assemblers(self): def _apply_bc(self, tensor, bc, u=None): assert u is None op2tensor = tensor.M - spaces = tuple(a.function_space() for a in tensor.a.arguments()) + spaces = tuple(_primal_space(a.function_space()) for a in tensor.a.arguments()) V = bc.function_space() component = V.component if component is not None: @@ -1492,7 +1562,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(parent=True) == fs for fs in spaces): raise TypeError("bc space does not match the test or trial function space") if spaces[0] != spaces[1]: # Not on a diagonal block, we cannot set diagonal entries @@ -1617,7 +1687,7 @@ 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, slate.TensorBase): sig = form.expression_hash @@ -1648,6 +1718,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`. @@ -1763,7 +1845,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) @@ -1781,7 +1867,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) @@ -1883,6 +1976,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) @@ -2041,18 +2146,27 @@ 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 block_index(bc): + fs = bc.function_space() + if fs.component is not None: + fs = fs.parent + return fs.index + + test_space = _primal_space(self.test_function_space) + bcrow = tuple( + bc for bc in self._bcs + if bc.function_space(parent=True) == test_space + and (len(test_space) == 1 or block_index(bc) == row) + ) + + trial_space = _primal_space(self.trial_function_space) + bccol = tuple( + bc for bc in self._bcs + if isinstance(bc, DirichletBC) + and bc.function_space(parent=True) == trial_space + and (len(trial_space) == 1 or block_index(bc) == col) + ) return bcrow, bccol def needs_unrolling(self): @@ -2153,7 +2267,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 @@ -2233,6 +2350,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/bcs.py b/firedrake/bcs.py index 16bed26f60..586d1707fe 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -76,11 +76,26 @@ def __iter__(self): yield self yield from itertools.chain(*self.bcs) - def function_space(self): + def function_space(self, parent=False): '''The :class:`.FunctionSpace` on which this boundary condition should - be applied.''' - - return self._function_space + be applied. + + Parameters + ---------- + parent : bool + If ``True``, walk up through any indexed or component subspaces + and return the top-level function space instead. + + Returns + ------- + firedrake.functionspaceimpl.WithGeometry + The function space. + ''' + V = self._function_space + if parent: + while V.parent is not None: + V = V.parent + return V def function_space_index(self): fs = self._function_space 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 58823f5451..2179998901 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, 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 -from ufl.constantvalue import zero, as_ufl +from ufl.constantvalue import as_ufl from ufl.form import ZeroBaseForm, BaseForm from ufl.core.interpolate import Interpolate as UFLInterpolate from pyop2 import op2 -from pyop2.caching import memory_and_disk_cache - from finat.ufl import TensorElement, VectorElement, MixedElement, FiniteElementBase -from finat.element_factory import create_element - -from tsfc.driver import compile_expression_dual_evaluation -from tsfc.ufl_utils import extract_firedrake_constants, hash_expr -from firedrake.utils import IntType, ScalarType, known_pyop2_safe, tuplify -from firedrake.pointeval_utils import runtime_quadrature_element -from firedrake.tsfc_interface import extract_numbered_coefficients, _cachedir +from firedrake.utils import IntType, ScalarType from firedrake.ufl_expr import Argument, Coargument, TrialFunction, TestFunction, action from firedrake.mesh import MissingPointsBehaviour, VertexOnlyMeshTopology, MeshGeometry, MeshTopology, VertexOnlyMesh from firedrake.petsc import PETSc @@ -151,6 +140,11 @@ def options(self) -> InterpolateOptions: """ return self._options + def subdomain_data(self): + """Return cell-iteration subdomain data for the target mesh.""" + domain = self.target_space.mesh().unique() + return {domain: {"cell": [self.options.subset]}} + @cached_property def _interpolator(self): """Access the numerical interpolator. @@ -163,8 +157,6 @@ def _interpolator(self): """ arguments = self.arguments() has_mixed_arguments = any(len(arg.function_space()) > 1 for arg in arguments) - if len(arguments) == 2 and has_mixed_arguments: - return MixedInterpolator(self) operand, = self.ufl_operands target_mesh = self.target_space.mesh() @@ -707,6 +699,69 @@ def __init__(self, expr, source_mesh, target_mesh): # Default access for forward 1-form or 2-form (forward and adjoint) self.access = op2.WRITE + @cached_property + def _form_interpolate(self): + options = asdict(self.ufl_interpolate.options) + options.update(subset=self.subset, access=self.access) + return self.ufl_interpolate._ufl_expr_reconstruct_( + self.operand, v=self.dual_arg, **options + ) + + @property + def _needs_adjoint_weighting(self): + return ( + isinstance(self.dual_arg, Cofunction) + and any(not V.finat_element.is_dg() for V in self.target_space) + ) + + @cached_property + def _weighted_dual_arg(self): + return Function(self.dual_arg.function_space()) + + @cached_property + def _adjoint_weight(self): + W = self.dual_arg.function_space() + weight = W.make_dat() + if len(W) > 1: + spaces_and_weights = zip(W, weight) + else: + spaces_and_weights = ((W, weight),) + + source_mesh = self.source_mesh.unique() + target_mesh = self.target_mesh.unique() + 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 + 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, iterset, + component_weight(op2.INC, node_map), + ) + with weight.vec as weight_vec: + weight_vec.reciprocal() + return weight + + @cached_property + def _assembler_form(self): + if self._needs_adjoint_weighting: + return self._form_interpolate._ufl_expr_reconstruct_( + self.operand, v=self._weighted_dual_arg + ) + return self._form_interpolate + + def _update_weighted_dual_arg(self): + self.dual_arg.dat.copy(self._weighted_dual_arg.dat) + with ( + self._adjoint_weight.vec_ro as weight, + self._weighted_dual_arg.dat.vec as dual, + ): + dual.pointwiseMult(dual, weight) + def _get_tensor(self, mat_type: Literal["aij", "baij"]) -> op2.Mat | Function | Cofunction: """Return a suitable tensor to interpolate into. @@ -768,66 +823,91 @@ def _get_monolithic_sparsity(self, mat_type: Literal["aij", "baij"]) -> op2.Spar block_sparse=(mat_type == "baij")) return sparsity - def _get_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),) + 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, + access=self.access, + ) + elif self.rank == 2: + return TwoFormAssembler( + self._assembler_form, + bcs=bcs, + mat_type=mat_type, + sub_mat_type=sub_mat_type, + needs_zeroing=True, + access=self.access, + ) else: - f = tensor or self._get_tensor(mat_type) - copyout = () + raise ValueError( + f"Cannot interpolate an expression with {self.rank} arguments" + ) - op2_tensor = f if isinstance(f, op2.Mat) else f.dat - loops = [] - if self.access is op2.INC: - loops.append(op2_tensor.zero) + def _get_callable(self, tensor=None, bcs=None, mat_type=None, sub_mat_type=None): + from firedrake.assemble import ParloopFormAssembler - # 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 + assembler = self._get_form_assembler( + bcs=bcs, mat_type=mat_type, sub_mat_type=sub_mat_type, + ) + # Compile here rather than inside the callable below, so that a target + # element the code generator cannot dual evaluate is reported to + # whoever asked for the callable. DirichletBC asks for one to find out + # whether it can interpolate its value, and projects instead when it + # cannot; a kernel left to the callable hides that from it until the + # boundary condition is first applied. + if isinstance(assembler, ParloopFormAssembler): + assembler.local_kernels + + 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,) + # The coordinates reach the kernel without being coefficients of the form. + inputs += (self.source_mesh.unique().coordinates, + self.target_mesh.unique().coordinates) + if ( + isinstance(tensor, Function | Cofunction) + and any(set(tensor.dat).intersection(set(input_.dat)) + 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 @property def _allowed_mat_types(self): - return {"aij", "baij", "matfree", None} + return {"aij", "baij", "nest", "matfree", None} class VomOntoVomInterpolator(SameMeshInterpolator): @@ -973,204 +1053,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. @@ -1206,28 +1088,6 @@ def get_interp_node_map(source_mesh: MeshGeometry, target_mesh: MeshGeometry, fs return m_ -try: - _expr_cachedir = os.environ["FIREDRAKE_TSFC_KERNEL_CACHE_DIR"] -except KeyError: - _expr_cachedir = os.path.join(tempfile.gettempdir(), - f"firedrake-tsfc-expression-kernel-cache-uid{os.getuid()}") - - -def _compile_expression_key(comm, expr, ufl_element, domain, parameters) -> tuple[Hashable, ...]: - """Generate a cache key suitable for :func:`tsfc.compile_expression_dual_evaluation`.""" - dual_arg, operand = expr.argument_slots() - return (hash_expr(operand), type(dual_arg), hash(ufl_element), tuplify(parameters)) - - -@memory_and_disk_cache( - hashkey=_compile_expression_key, - cachedir=_cachedir -) -@PETSc.Log.EventDecorator() -def compile_expression(comm, *args, **kwargs): - return compile_expression_dual_evaluation(*args, **kwargs) - - def compose_map_and_cache(map1: op2.Map, map2: op2.Map | None) -> op2.ComposedMap | None: """ Retrieve a :class:`pyop2.ComposedMap` map from the cache of map1 diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index cde7f678a1..1585686d8e 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: @@ -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: @@ -269,7 +277,10 @@ def _real_mangle(form): """If the form contains arguments in the Real function space, replace these with literal 1 before passing to tsfc.""" a = form.arguments() - reals = [x.ufl_element().family() == "Real" for x in a] + # A Coargument names the space the result lands in rather than something to + # integrate against, so TSFC dual-evaluates it instead. + reals = [x.ufl_element().family() == "Real" and not isinstance(x, ufl.Coargument) + for x in a] if not any(reals): return form replacements = {} 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/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_bcs.py b/tests/firedrake/regression/test_bcs.py index 9e43ba805b..ccd4477e99 100644 --- a/tests/firedrake/regression/test_bcs.py +++ b/tests/firedrake/regression/test_bcs.py @@ -56,7 +56,7 @@ def test_assemble_bcs_wrong_fs(V, measure): u, v = TrialFunction(V), TestFunction(V) W = FunctionSpace(V.mesh(), "CG", 2) - with pytest.raises(RuntimeError): + with pytest.raises(TypeError): assemble(inner(u, v)*measure, bcs=[DirichletBC(W, 32, 1)]) @@ -65,7 +65,7 @@ def test_assemble_bcs_wrong_fs_interior(V): u, v = TrialFunction(V), TestFunction(V) W = FunctionSpace(V.mesh(), "CG", 2) n = FacetNormal(V.mesh()) - with pytest.raises(RuntimeError): + with pytest.raises(TypeError): assemble(inner(jump(u, n), jump(v, n))*dS, bcs=[DirichletBC(W, 32, 1)]) diff --git a/tests/firedrake/regression/test_interp_dual.py b/tests/firedrake/regression/test_interp_dual.py index b0ce971a95..a84133cc55 100644 --- a/tests/firedrake/regression/test_interp_dual.py +++ b/tests/firedrake/regression/test_interp_dual.py @@ -395,3 +395,153 @@ 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) + + +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 + + 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/tests/firedrake/regression/test_interpolate.py b/tests/firedrake/regression/test_interpolate.py index 4539c20504..0eca9d52f8 100644 --- a/tests/firedrake/regression/test_interpolate.py +++ b/tests/firedrake/regression/test_interpolate.py @@ -641,6 +641,24 @@ def test_interpolator_reuse(family, degree, mode): assert np.allclose(result.dat.data, expected) +@pytest.mark.parallel([1, 3]) +def test_square_space_bcs(): + mesh = UnitSquareMesh(2, 2) + V = FunctionSpace(mesh, "CG", 1) + rg = RandomGenerator(PCG64(seed=123456789)) + w = rg.uniform(V) + + # Source and target agree, so the interpolation has a diagonal to carry + # the boundary rows, just as a Form on the same spaces does. + I = assemble(interpolate(2 * TrialFunction(V), V), bcs=[DirichletBC(V, 0, 1)]) + result = assemble(action(I, w)) + + expected = Function(V).assign(2 * w) + DirichletBC(V, w, 1).apply(expected) + + assert np.allclose(result.dat.data, expected.dat.data) + + def test_mixed_space_bcs(): mesh = UnitSquareMesh(2, 2) V = FunctionSpace(mesh, "CG", 1) diff --git a/tests/firedrake/regression/test_interpolation_manual.py b/tests/firedrake/regression/test_interpolation_manual.py index cedaee54be..3c32838a8d 100644 --- a/tests/firedrake/regression/test_interpolation_manual.py +++ b/tests/firedrake/regression/test_interpolation_manual.py @@ -260,7 +260,7 @@ def test_mixed_space_interpolation(): for j in range(2): sub_mat = I.petscmat.getNestSubMatrix(i, j) if i != j: - assert not sub_mat + assert sub_mat.norm() == 0.0 continue else: res_block = assemble(interpolate(TrialFunction(U.sub(j)), W.sub(i))) diff --git a/tests/firedrake/regression/test_interpolation_operators.py b/tests/firedrake/regression/test_interpolation_operators.py index 6c97264ded..bcaa4d56f8 100644 --- a/tests/firedrake/regression/test_interpolation_operators.py +++ b/tests/firedrake/regression/test_interpolation_operators.py @@ -1,6 +1,6 @@ from firedrake import * from firedrake.interpolation import ( - MixedInterpolator, SameMeshInterpolator, CrossMeshInterpolator, + SameMeshInterpolator, CrossMeshInterpolator, get_interpolator, VomOntoVomInterpolator, ) from firedrake.matrix import ImplicitMatrix, Matrix @@ -73,9 +73,12 @@ def test_same_mesh_mattype(value_shape, mat_type, mode): res2 = assemble(action(adjoint(forward_I_mat), f)) assert np.allclose(res2.dat.data, exact.dat.data) - with pytest.raises(NotImplementedError): - # MatNest only implemented for interpolation between MixedFunctionSpaces - assemble(interp, mat_type="nest") + # A nest of the one block these unmixed spaces make collapses to that + # block's own type, just as it does for a Form. + nest_mat = assemble(interp, mat_type="nest") + assert nest_mat.petscmat.type == prefix + ("baij" if value_shape == "vector" else "aij") + res = assemble(action(nest_mat, f)) + assert np.allclose(res.dat.data, exact.dat.data) @pytest.mark.parametrize("value_shape", ["scalar", "vector"], ids=lambda v: f"fs_type={v}") @@ -188,7 +191,7 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): expr = as_vector([x**2, x**2, y**2, y**2]) interp = interpolate(TrialFunction(U), W) - assert isinstance(get_interpolator(interp), MixedInterpolator) + assert isinstance(get_interpolator(interp), SameMeshInterpolator) I_mat = assemble(interp, mat_type=mat_type, sub_mat_type=sub_mat_type) assert isinstance(I_mat, ImplicitMatrix if mat_type == "matfree" else Matrix) @@ -198,17 +201,19 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): assert I_mat.petscmat.type == "seqaij" else: assert I_mat.petscmat.type == "nest" + if value_shape == "scalar": + # Always seqaij for scalar + sub_type = "seqaij" + else: + # A blocked space makes matnest default to baij + sub_type = "seq" + (sub_mat_type if sub_mat_type else "baij") for (i, j) in [(0, 0), (0, 1), (1, 0), (1, 1)]: + # Every block is assembled, as it is for a Form. The components + # do not mix, so the off-diagonal ones assemble to zero. sub_mat = I_mat.petscmat.getNestSubMatrix(i, j) + assert sub_mat.type == sub_type if i != j: - assert not sub_mat - continue - if value_shape == "scalar": - # Always seqaij for scalar - assert sub_mat.type == "seqaij" - else: - # matnest sub_mat_type defaults to aij - assert sub_mat.type == "seq" + (sub_mat_type if sub_mat_type else "aij") + assert sub_mat.norm() == 0.0 f = Function(U).interpolate(expr) exact = Function(W).interpolate(expr) @@ -216,5 +221,6 @@ def test_mixed_same_mesh_mattype(value_shape, mat_type, sub_mat_type): for resi, exi in zip(res.subfunctions, exact.subfunctions): assert np.allclose(resi.dat.data, exi.dat.data) - with pytest.raises(NotImplementedError): + with pytest.raises(ValueError): + # A mixed space has no block structure to give BAIJ, as for a Form. assemble(interp, mat_type="baij") 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 2c480f9c55..638a73d87f 100644 --- a/tsfc/driver.py +++ b/tsfc/driver.py @@ -3,24 +3,25 @@ import sys from itertools import chain import numpy -from finat.physically_mapped import NeedsCoordinateMappingElement import ufl from ufl.algorithms import extract_coefficients from ufl.algorithms.analysis import has_type from ufl.algorithms.apply_coefficient_split import CoefficientSplitter from ufl.classes import Form, GeometricQuantity -from ufl.domain import extract_unique_domain, extract_domains +from ufl.domain import MeshSequence, extract_unique_domain, extract_domains import gem import gem.impero_utils as impero_utils import finat from finat.element_factory import as_fiat_cell +from finat.point_set import UnknownPointSet +from finat.quadrature import QuadratureRule +from finat.ufl import FiniteElement, TensorElement from tsfc import fem, ufl_utils from tsfc.logging import logger -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 @@ -54,6 +55,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. @@ -78,9 +94,13 @@ 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" + form = ufl_utils.lower_form_interpolations(form) # Determine whether in complex mode: complex_mode = parameters and is_complex(parameters.get("scalar_type")) @@ -111,6 +131,102 @@ def compile_form(form, prefix="form", parameters=None, dont_split_numbers=(), di return kernels +def compile_interpolate(expression, prefix="interpolate", parameters=None): + """Compile an interpolation using the integral kernel builder.""" + parameters = preprocess_parameters(parameters) + complex_mode = is_complex(parameters["scalar_type"]) + original_expression = expression + original_coefficients = expression.coefficients() + dual_arg, operand = expression.argument_slots() + target_domain = dual_arg.ufl_function_space().ufl_domain() + if isinstance(target_domain, MeshSequence): + target_domains = set(target_domain.meshes) + if len(target_domains) != 1: + raise NotImplementedError( + "Interpolation onto multiple distinct meshes is not supported" + ) + target_domain, = target_domains + source_domain = ( + extract_unique_domain(operand) + or target_domain + ) + all_domains = expression.ufl_domains() + + target_element = expression.ufl_element() + if ( + target_domain.topological_dimension == 0 + and source_domain.topological_dimension > 0 + ): + cell = source_domain.ufl_cell() + point_expr = gem.Variable("rt_X", (1, cell.topological_dimension)) + point_set = UnknownPointSet(point_expr) + rule = QuadratureRule( + point_set, weights=[1.0], ref_el=as_fiat_cell(cell) + ) + shape = target_element.pullback.physical_value_shape( + target_element, target_domain + ) + target_element = FiniteElement( + "Quadrature", cell=cell, degree=0, quad_scheme=rule + ) + if shape: + symmetry = None if len(shape) < 2 else expression.ufl_element().symmetry() + target_element = TensorElement( + target_element, shape=shape, symmetry=symmetry + ) + + operand = apply_mapping(operand, target_element, source_domain) + operand = ufl_utils.preprocess_expression( + operand, complex_mode=complex_mode + ) + operand = simplify_abs(operand, complex_mode) + expression = ufl.Interpolate(operand, dual_arg) + + coefficients = expression.coefficients() + coefficient_split = {} + for coefficient in coefficients: + element = coefficient.ufl_element() + if type(element) is finat.ufl.MixedElement: + domain = extract_unique_domain( + coefficient, expand_mesh_sequence=False + ) + coefficient_split[coefficient] = [ + ufl.Coefficient(ufl.FunctionSpace(mesh, subelement)) + for mesh, subelement in zip( + domain.iterable_like(element), element.sub_elements + ) + ] + + form_data = TSFCInterpolationFormData( + original_form=original_expression, + preprocessed_form=expression, + reduced_coefficients=coefficients, + function_replace_map={coefficient: coefficient for coefficient in coefficients}, + coefficient_split=coefficient_split, + original_coefficient_positions=tuple( + original_coefficients.index(coefficient) + for coefficient in coefficients + ), + constants=extract_firedrake_constants(expression), + ) + integral_data = TSFCInterpolationData( + domain=source_domain, + iteration_domain=target_domain, + integral_type="cell", + subdomain_id=("everywhere",), + domain_integral_type_map={domain: "cell" for domain in all_domains}, + enabled_coefficients=(True,) * len(coefficients), + integrals=(), + expression=expression, + target_element=target_element, + ) + return [ + compile_integral( + integral_data, form_data, prefix, parameters, diagonal=False + ) + ] + + def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=False): """Compiles a UFL integral into an assembly kernel. @@ -143,8 +259,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, @@ -173,12 +294,19 @@ def compile_integral(integral_data, form_data, prefix, parameters, *, diagonal=F # so we should attach the constants to integral data instead builder.set_constants(form_data.constants) ctx = builder.create_context() - for integral in integral_data.integrals: + if isinstance(integral_data, TSFCInterpolationData): params = parameters.copy() - params.update(integral.metadata()) # integral metadata overrides - integrand_exprs = builder.compile_integrand(integral.integrand(), params, ctx) - integral_exprs = builder.construct_integrals(integrand_exprs, params) - builder.stash_integrals(integral_exprs, params, ctx) + interpolate_exprs = builder.compile_interpolate( + integral_data.expression, integral_data.target_element, params, ctx + ) + builder.stash_integrals(interpolate_exprs, params, ctx) + else: + for integral in integral_data.integrals: + params = parameters.copy() + params.update(integral.metadata()) # integral metadata overrides + integrand_exprs = builder.compile_integrand(integral.integrand(), params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) return builder.construct_kernel(kernel_name, ctx, parameters["add_petsc_events"]) @@ -332,33 +460,9 @@ def compile_expression_dual_evaluation(expression, ufl_element, *, if isinstance(to_element, finat.QuadratureElement): kernel_cfg["quadrature_rule"] = to_element._rule - dual_arg, operand = expression.argument_slots() - - # Create callable for translation of UFL expression to gem - fn = DualEvaluationCallable(operand, kernel_cfg) - - # Get the gem expression for dual evaluation and corresponding basis - # indices needed for compilation of the expression - if isinstance(to_element, NeedsCoordinateMappingElement): - ctx = fem.PointSetContext(**kernel_cfg) - mt = analyse_modified_terminal(ufl.Coefficient(dual_arg.ufl_function_space().dual())) - coordinate_mapping = fem.CoordinateMapping(mt, ctx) - else: - coordinate_mapping = None - evaluation, point_indices, basis_indices = to_element.dual_evaluation(fn, coordinate_mapping) - quadrature_multiindex = tuple(point_indices) - - # Compute the action against the dual argument - if isinstance(dual_arg, ufl.Cofunction): - gem_dual = builder.coefficient_map[dual_arg] - if complex_mode: - evaluation = gem.MathFunction('conj', evaluation) - # The dual argument contracts over the nodes, so the basis indices are - # reduction indices like the points, not indices of the return value. - evaluation = evaluation * gem_dual[basis_indices] - quadrature_multiindex += tuple(basis_indices) - basis_indices = () - else: + evaluation, quadrature_multiindex, basis_indices = fem.dual_evaluate(expression, to_element, kernel_cfg) + dual_arg, _ = expression.argument_slots() + if not isinstance(dual_arg, ufl.Cofunction): argument_multiindices[dual_arg.number()] = basis_indices argument_multiindices = dict(sorted(argument_multiindices.items())) @@ -387,64 +491,3 @@ def compile_expression_dual_evaluation(expression, ufl_element, *, builder.set_output(return_var) # Build kernel tuple return builder.construct_kernel(impero_c, index_names, needs_external_coords, parameters["add_petsc_events"], name=name) - - -class DualEvaluationCallable(object): - """ - Callable representing a function to dual evaluate. - - When called, this takes in a - :class:`finat.point_set.AbstractPointSet` and returns a GEM - expression for evaluation of the function at those points. - - :param expression: UFL expression for the function to dual evaluate. - :param kernel_cfg: A kernel configuration for creation of a - :class:`GemPointContext` or a :class:`PointSetContext` - - Not intended for use outside of - :func:`compile_expression_dual_evaluation`. - """ - def __init__(self, expression, kernel_cfg): - self.expression = expression - self.kernel_cfg = kernel_cfg - - def __call__(self, ps): - """The function to dual evaluate. - - :param ps: The :class:`finat.point_set.AbstractPointSet` for - evaluating at - :returns: a gem expression representing the evaluation of the - input UFL expression at the given point set ``ps``. - For point set points with some shape ``(*value_shape)`` - (i.e. ``()`` for scalar points ``(x)`` for vector points - ``(x, y)`` for tensor points etc) then the gem expression - has shape ``(*value_shape)`` and free indices corresponding - to the input :class:`finat.point_set.AbstractPointSet`'s - free indices alongside any input UFL expression free - indices. - """ - - if not isinstance(ps, finat.point_set.AbstractPointSet): - raise ValueError("Callable argument not a point set!") - - # Avoid modifying saved kernel config - kernel_cfg = self.kernel_cfg.copy() - - if isinstance(ps, finat.point_set.UnknownPointSet): - # Run time known points - kernel_cfg.update(point_indices=ps.indices, point_expr=ps.expression) - # GemPointContext's aren't allowed to have quadrature rules - kernel_cfg.pop("quadrature_rule", None) - translation_context = fem.GemPointContext(**kernel_cfg) - else: - # Compile time known points - kernel_cfg.update(point_set=ps) - translation_context = fem.PointSetContext(**kernel_cfg) - - gem_expr, = fem.compile_ufl(self.expression, translation_context, point_sum=False) - # In some cases ps.indices may be dropped from expr, but nothing - # new should now appear - argument_multiindices = kernel_cfg["argument_multiindices"].values() - assert set(gem_expr.free_indices) <= set(chain(ps.indices, *argument_multiindices)) - - return gem_expr diff --git a/tsfc/fem.py b/tsfc/fem.py index 943089052e..a6840a7997 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, @@ -171,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. @@ -327,6 +338,106 @@ def needs_coordinate_mapping(element): return isinstance(create_element(element), NeedsCoordinateMappingElement) +def dual_evaluate(expression: ufl.Interpolate, to_element: FiniteElementBase, kernel_cfg: dict) -> tuple: + """Translate an interpolation operand and evaluate its target dual basis. + + Parameters + ---------- + expression : ufl.Interpolate + Preprocessed interpolation expression in the target reference frame. + to_element : finat.FiniteElementBase + Target FInAT element. + kernel_cfg : dict + Configuration for the point-evaluation translation context. + + Returns + ------- + tuple + The GEM expression for the local interpolated values, the multiindex + to contract it on, and the basis indices of the return value. + """ + dual_arg, operand = expression.argument_slots() + fn = DualEvaluationCallable(operand, kernel_cfg) + + if isinstance(to_element, NeedsCoordinateMappingElement): + ctx = PointSetContext(**kernel_cfg) + coefficient = ufl.Coefficient(dual_arg.ufl_function_space().dual()) + coordinate_mapping = CoordinateMapping(analyse_modified_terminal(coefficient), ctx) + else: + coordinate_mapping = None + if isinstance(dual_arg, ufl.Cofunction): + gem_duals = kernel_cfg["interface"].coefficient_components( + dual_arg, None + ) + else: + gem_duals = () + + if len(gem_duals) > 1: + assert to_element.is_mixed + assert len(to_element.elements) == len(gem_duals) + # The summands do not share their points, so each one contracts on its own. + evaluations = [] + for element, gem_dual in zip(to_element.elements, gem_duals): + evaluation, point_indices, basis_indices = element.dual_evaluation( + fn, coordinate_mapping + ) + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + evaluations.append(gem.IndexSum( + evaluation * gem_dual[basis_indices], + tuple(point_indices) + basis_indices + )) + evaluation = gem.optimise.make_sum(evaluations) + quadrature_multiindex = () + basis_indices = () + else: + evaluation, point_indices, basis_indices = to_element.dual_evaluation( + fn, coordinate_mapping + ) + quadrature_multiindex = tuple(point_indices) + if gem_duals: + if is_complex(kernel_cfg["scalar_type"]): + evaluation = gem.MathFunction("conj", evaluation) + # The dual argument contracts over the nodes, so the basis indices + # are reduction indices like the points, not return value indices. + evaluation = evaluation * gem_duals[0][basis_indices] + quadrature_multiindex += tuple(basis_indices) + basis_indices = () + + return evaluation, quadrature_multiindex, basis_indices + + +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 +850,45 @@ 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: + 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( + ctx, domain, mt.restriction), + "ufl_cell": domain.ufl_cell(), + "integration_dim": as_fiat_cell(domain.ufl_cell()).get_dimension(), + "argument_multiindices": ctx.argument_multiindices, + "index_cache": ctx.index_cache, + "scalar_type": ctx.scalar_type, + } + if isinstance(element, finat.QuadratureElement): + kernel_cfg["quadrature_rule"] = element._rule + + evaluation, quadrature_multiindex, basis_indices = dual_evaluate(terminal, element, kernel_cfg) + # The interpolation points are internal to the local solve, so contract + # them here: only the form's own quadrature points stay free. + evaluation = gem.IndexSum(evaluation, quadrature_multiindex) + vec = gem.ComponentTensor(evaluation, basis_indices) + return translate_element(terminal, mt, ctx, vec, element, beta=element.get_indices()) + + @translate.register(Coefficient) def translate_coefficient(terminal, mt, ctx): - domain = extract_unique_domain(terminal) vec = ctx.coefficient(terminal, mt.restriction) element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) + return translate_element(terminal, mt, ctx, vec, element) + + +def translate_element(terminal: ufl.core.expr.Expr, mt: ModifiedTerminal, ctx: ContextBase, + vec: gem.Node, element: FiniteElementBase, beta: tuple | None = None) -> gem.Node: + """Evaluate local finite element values at the current points.""" + domain = extract_unique_domain(terminal) # Collect FInAT tabulation for all entities per_derivative = collections.defaultdict(list) @@ -771,16 +916,31 @@ 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) - 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) @@ -788,7 +948,16 @@ 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)) + 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..128e7643a0 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.domain import MeshSequence, extract_unique_domain +from ufl.algorithms.apply_coefficient_split import CoefficientSplitter +import finat import gem import gem.impero_utils as impero_utils import petsctools @@ -41,6 +44,7 @@ def __init__(self, scalar_type): # Coefficients self.coefficient_map = collections.OrderedDict() + self.coefficient_split = {} # Constants self.constant_map = collections.OrderedDict() @@ -63,6 +67,16 @@ def coefficient(self, ufl_coefficient, restriction): else: return kernel_arg[{'+': 0, '-': 1}[restriction]] + def coefficient_components(self, ufl_coefficient, restriction): + """Return GEM expressions for a coefficient's stored components.""" + coefficients = self.coefficient_split.get( + ufl_coefficient, (ufl_coefficient,) + ) + return tuple( + self.coefficient(coefficient, restriction) + for coefficient in coefficients + ) + def constant(self, const): return self.constant_map[const] @@ -136,6 +150,58 @@ def domain_integral_type_map(self): class KernelBuilderMixin(object): """Mixin for KernelBuilder classes.""" + def compile_interpolate(self, expression, target_element, params, ctx): + """Compile UFL interpolate. + + :arg expression: UFL interpolate. + :arg target_element: UFL element of the interpolation target. This is + not the dual argument's own element when the target is a point + cloud: the points are then only known at run time, so the target + is a quadrature element on the source cell. + :arg params: a dict containing "quadrature_rule". + :arg ctx: context created with :meth:`create_context` method. + + See :meth:`create_context` for typical calling sequence. + """ + expression = CoefficientSplitter(self.coefficient_split)( + expression + ) + target_element = self.create_element(target_element) + config = self.fem_config() + config.update( + argument_multiindices=self.argument_multiindices, + index_cache=ctx["index_cache"], + ) + if isinstance(target_element, finat.QuadratureElement): + config["quadrature_rule"] = target_element._rule + evaluation, quadrature_multiindex, basis_indices = fem.dual_evaluate( + expression, target_element, config + ) + dual_arg, _ = expression.argument_slots() + if not isinstance(dual_arg, Cofunction): + arguments = expression.arguments() + argument_number = arguments.index(dual_arg) + output_indices = self.argument_multiindices[argument_number] + if basis_indices != output_indices: + if tuple(i.extent for i in basis_indices) != tuple( + i.extent for i in output_indices + ): + raise ValueError("Interpolation output index shape mismatch") + mapper = gem.node.MemoizerArg( + gem.optimise.filtered_replace_indices + ) + evaluation = mapper( + evaluation, tuple(zip(basis_indices, output_indices)) + ) + + mode = pick_mode(params["mode"]) + ctx["quadrature_indices"].extend(quadrature_multiindex) + # Argument factorisation does not cancel every Delta here, so lower them. + ctx["finalise_options"]["replace_delta"] = True + return mode.Integrals( + [evaluation], quadrature_multiindex, self.argument_multiindices, params + ) + def compile_integrand(self, integrand, params, ctx): """Compile UFL integrand. @@ -220,6 +286,7 @@ def compile_gem(self, ctx): options = dict(reduce(operator.and_, [mode.finalise_options.items() for mode in mode_irs.keys()])) + options.update(ctx['finalise_options']) expressions = impero_utils.preprocess_gem(expressions, **options) # Let the kernel interface inspect the optimised IR to register @@ -277,6 +344,11 @@ def create_context(self): Dict for mode representations. + *finalise_options* + + Options overriding the modes' own :func:`impero_utils.preprocess_gem` + options. + For each set of integrals to make a kernel for (i,e., `integral_data.integrals`), one must first create a ctx object by calling :meth:`create_context` method. @@ -299,6 +371,7 @@ def create_context(self): """ return {'index_cache': {}, 'quadrature_indices': [], + 'finalise_options': {}, 'mode_irs': collections.OrderedDict()} @@ -577,7 +650,16 @@ 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() + try: + integral_type = domain_integral_type_map[domain] + except KeyError: + # An unsplit argument (e.g. a mixed-space patch argument) reports + # its domain as a MeshSequence rather than a single mesh: every + # mesh it sequences is the same iteration, so they must agree. + if not isinstance(domain, MeshSequence): + raise + integral_type, = {domain_integral_type_map[m] for m in domain.meshes} if integral_type is None: raise RuntimeError(f"Can not determine integral_type on {arg}") if integral_type.startswith("interior_facet"): diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py index cc8fd7a61e..c4c4d91bb8 100644 --- a/tsfc/kernel_interface/firedrake_loopy.py +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -152,18 +152,20 @@ def set_cell_sizes(self, domains): measure of the mesh size around each vertex (hence this lives in P1). - Should the domain have topological dimension 0 this does - nothing. + A domain of topological dimension 0 gets a ``None`` entry: every + domain must keep its slot, since the active domain numbers index + this dict positionally. """ self._cell_sizes = {} for i, domain in enumerate(domains): if domain.ufl_cell().topological_dimension > 0: - # Can't create P1 since only P0 is a valid finite element if - # topological_dimension is 0 and the concept of "cell size" - # is not useful for a vertex. f = Coefficient(FunctionSpace(domain, FiniteElement("P", domain.ufl_cell(), 1))) expr = prepare_coefficient(f, f"cell_sizes_{i}", self._domain_integral_type_map) - self._cell_sizes[domain] = expr + else: + # Only P0 is a valid finite element on a vertex, and the + # concept of "cell size" is not useful there. + expr = None + self._cell_sizes[domain] = expr def create_element(self, element, **kwargs): """Create a FInAT element (suitable for tabulating with) given @@ -293,6 +295,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/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/spectral.py b/tsfc/spectral.py index a521fdb2fd..ac270fafd2 100644 --- a/tsfc/spectral.py +++ b/tsfc/spectral.py @@ -2,7 +2,7 @@ from functools import partial from itertools import chain, zip_longest -from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem.gem import Delta, FlexiblyIndexed, Indexed, Sum, index_sum, one from gem.node import Memoizer, MemoizerArg from gem.optimise import filtered_replace_indices from gem.optimise import delta_elimination as _delta_elimination @@ -123,7 +123,7 @@ def classify(argument_indices, expression, delta_inside): if n == 0: return OTHER elif n == 1: - if isinstance(expression, (Delta, Indexed)) and not delta_inside(expression): + if isinstance(expression, (Delta, FlexiblyIndexed, Indexed)) and not delta_inside(expression): return ATOMIC else: return COMPOUND diff --git a/tsfc/ufl_utils.py b/tsfc/ufl_utils.py index ce25dc3087..eaf8bf4ffd 100644 --- a/tsfc/ufl_utils.py +++ b/tsfc/ufl_utils.py @@ -5,6 +5,7 @@ import numpy import ufl +from ufl.algorithms.map_integrands import map_integrand_dags from ufl import as_tensor, indices, replace from ufl.algorithms import compute_form_data as ufl_compute_form_data from ufl.algorithms import estimate_total_polynomial_degree @@ -23,7 +24,7 @@ from ufl.geometry import QuadratureWeight from ufl.geometry import Jacobian, JacobianDeterminant, JacobianInverse from ufl.classes import (Abs, Argument, CellOrientation, - Expr, FloatValue, Division, + Expr, FloatValue, Division, ReferenceValue, Product, ScalarValue, Sqrt, Zero, CellVolume, FacetArea) from ufl.utils.sorting import sorted_by_count @@ -37,6 +38,39 @@ 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(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) + + +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, @@ -395,6 +429,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)