[connorjward]

This PR entirely replaces PyOP2 with pyop3. It represents the culmination of 5 years of work.

Motivation

The Firedrake team is in the business of designing software abstractions. For many years now it has been clear that the design of PyOP2 isn't quite right. In PyOP2's view of the world there is no such thing as a mesh (or anything that looks remotely like one) and whole bunch of assumptions are baked right into the abstraction: no loop nesting allowed, no (proper) map composition allowed, only specific types of indirection maps are permitted, that are near impossible to work around once encountered. pyop3 is a fundamental and complete rewrite of PyOP2 to address all of these issues.

pyop3 can be thought of as 2 major pieces:

1. A new abstraction for data layouts that is 'mesh aware'. It becomes possible to ask questions such as 'where are all of the DoFs that lie in the closure of this cell'.
2. A DSL for expressing computations over these data structures. A much richer set of operations are supported including nested loops and arbitrary map composition.

This is very abstract and so I like to describe pyop3 in an alternative way: pyop3 is like numpy but for unstructured mesh data.

To learn more about pyop3 my thesis is the canonical reference.

Benefits

Replacing PyOP2 with pyop3 will enable a broad swathe of new research to occur because some operations should now become so much easier to express and experiment with. A non-exhaustive list may be found here.

In addition to enabling lots of new technologies, another very significant benefit to integrating pyop3 is the reduction in technical debt and code complexity. Many complex pieces of preexisting Firedrake code have been simplified and/or centralised into very few functions. For example, extruded meshes are now just the same as regular meshes. These changes reduce special-casing throughout Firedrake, making everything more composable!

Further, the current work on GPU support in Firedrake and FUSE both depend on pyop3.

Key changes

PyOP2 is gone

Obviously. Anywhere that you explicitly use the PyOP2 API (i.e. from pyop2 import op2) will have to be replaced with pyop3 (import pyop3 as op3). If you use PyOP2 indirectly (e.g. function.dat.data, mesh.cell_set, V.dof_dset) then some things will be expected to work (e.g. function.dat.data) but others have no equivalent pyop3 object and thus no longer exist.

Extruded meshes are now DMPlexs

Extruded meshes are now fully fledged PETSc DMPlex objects, instead of hybrid objects where only Firedrake knew about the extrusion. This simplifies an awful lot of code and will enable things like submeshing extruded meshes.

Breaking API changes

mesh.interior_facets and mesh.exterior_facets are now AxisTrees, not _Facets

Instead of mesh.exterior_facets.unique_markers and mesh.interior_facets.unique_markers you should use mesh.facet_markers.

Other attributes of _Facets are also available via the mesh with new names.

function_space.cell_node_map is now a pyop3 Map instead of a PyOP2 one

This means that it will not work as expected. In particular V.cell_node_map() will now raise an error because pyop3 maps expect arguments.

It is common practice for downstream packages to inspect function_space.cell_node_map().values. Instead you should simply do function_space.cell_node_list.

Support for variable layer extruded meshes is dropped

The functionality appeared not to be used by the community (e.g. thetisproject/thetis#405) and porting the functionality would have been a large amount of work. Eventually the functionality will be recoverable via Submesh.

API deprecations

.sub() expects a tuple instead of integer for tensor spaces.

For example, to get the (1, 1) component of a 2x2 tensor function space one should now pass f.sub((1, 1)) instead of f.sub(3). The latter will continue to work for now but a warning is raised.

ASMExtrudedStarPC is being deprecated

Now that extruded meshes are just meshes, you should just use ASMStarPC. To manage the fact that until now ASMStarPC on an extruded mesh would give column patches, a new column option has been added.

Known issues

There are a number of things that are known to be currently broken with pyop3 and we hope to fix them in due course.

1. Performance is currently a fair bit slower than previous. This is expected to (a) be predominantly in the setup parts of the code, and hence amortize over long running simulations, and (b) will be addressed over time.
2. 4D extrusion is not currently supported
3. ASMExtrudedStarPC does not work in 3D (see https://gitlab.com/petsc/petsc/-/work_items/1946)
4. MATIS and FDM and BDDC preconditioners have bugs that will be fixed in later PRs

Disabled tests

The tests that are disabled in this PR and that will need subsequent fixing are:

• regression/test_moore_spence.py (in complex mode)
• regression/test_bddc.py
• regression/test_fdm.py (only some)
• regression/test_assemble.py::test_matis (only some)
• regression/test_mass_lumping.py (4D extrusion test case)
• regression/test_adjoint_operators.py::test_init_constant_diff_mesh
• external_operators/test_jax_operator.py::test_mixed_space_bcs
• external_operators/test_pytorch_operator.py::test_mixed_space_bcs

[sghelichkhani]

@connorjward this is what Claude Code came up with for the pyop3 memory growth in #5534: #5538