Skip to content

Conversation

@connorjward
Copy link
Contributor

@connorjward connorjward commented Jan 11, 2024

Replace PyOP2 with pyop3.

Thesis link

Summary

Data structures for unstructured meshes are much more complicated than N-dimensional arrays. Firedrake therefore does an awful lot of very involved bookkeeping to make sure that things work.

The key contribution of pyop3 is a new abstraction for describing data layouts, called axis trees. Axis trees are able to describe mesh data layouts completely from a high-level specification. This has the dual set of benefits of: (a) being able to delete a lot of complex Firedrake code, and (b) we can now express a whole suite of algorithms that would otherwise be very difficult to write.

You can think of pyop3 as 'a JIT-compiled numpy for unstructured mesh data'.

pyop3

Core concepts

Expected benefits

API changes visible in Firedrake

function.dat.data now always returns a flat array

This means that coordinates.dat.data gives back an array with shape (nverts*gdim,) instead of (nverts, gdim).

Explanation

Now that data layouts are much more flexible than before it no longer makes sense to apply multi-dimensional shape to dats. For example, consider a mixed function space with both scalar and vector subspaces. mixed_function.dat.data cannot straightforwardly have multi-dimensional shape.

Similarly, since we can now permute data layouts freely, it is valid for coordinates to be stored with effective layout (gdim, nverts).

To avoid all this confusion dat.data always returns something flat. The axis tree (dat.axes) provides the information needed to interpret it.

Drop support for variable layer extruded meshes

Users should use a cell submesh instead to get the right behaviour.

Explanation

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.

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

function.dat.vec is now a function not a property.

This allows one to pass block size to the generated Vec.

@connorjward connorjward changed the title pyop3 (WIP) pyop3 Mar 21, 2024
connorjward and others added 27 commits July 28, 2025 17:25
WIP
Currently failing because ragged temporaries crash. This is currently
hitting a breakpoint in layouts.py so we can progress from there.
WIP
Fixed a lot of intermediate bugs. Now getting a segfault... Probably
should start looking where we have those ragged temporaries.
Now a failure much much later on in the process.
What a momentous occasion.
Seems like some of the traversals (esp. replace) can be quite expensive.
We seem to gain with eager evals of some conditions...
--global and others added 30 commits December 12, 2025 17:14
Works for axis trees too!

Next is to keep up the disk caching work and get the right mapping for
the buffer indices.

Then can maybe start translating more traversals into this form.
I.e. it's actually still broken.
Appears to work in parallel, now need to apply to more cases.
Not currently working because need to be clever about cache keys for
iteration_spec.
Immutable objects should be reusable.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants