Fieldsplit: replace empty Forms with ZeroBaseForm #3947

pbrubeck · 2024-12-30T04:03:51Z

Description

This PR addresses the consequences of FEniCS/ufl#336, which introduces more agressive simplifications that cause empty forms to be returned when extracting a subblock that is supposed to be zero (e.g. the pressure block in Stokes).

Here empty Jacobians are replaced with ZeroBaseForm, as we require two Arguments in order for the adjoint to make sense.

The simplifications should also speed up code generation and even reduce flop counts in some cases.

github-actions · 2024-12-30T04:24:36Z

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake real	8162 ran	7487 passed	675 skipped	0 failed

github-actions · 2024-12-30T04:25:15Z

	Tests	Passed ✅	Skipped ⏭️	Failed ❌
Firedrake complex	8156 ran	6682 passed	1474 skipped	0 failed

firedrake/adjoint_utils/variational_solver.py

firedrake/formmanipulation.py

firedrake/adjoint_utils/variational_solver.py

firedrake/formmanipulation.py

JHopeCollins · 2025-01-09T10:51:53Z

firedrake/formmanipulation.py

+        if len(indices) == 1:
+            i = indices[0]
+            W = V[i]
+            W = DualSpace(W.mesh(), W.ufl_element())


Can you use the subspace function here too? It will need a switch adding for primal/dual but I think it would be nicer to use subspace for both.

I think it would also be nice to lift subspace out, possibly into one of the functionspace files. It's useful more generally - see for example line 678 in slate.py where we have the same if/else block to construct the subspace.

I was thinking to extend FunctionSpace.__getitem__ to allow a list index return the subspace, since this is how sub-arrays are extracted in numpy.

The if statement can be avoided by using FunctionSpace.collapse()

I was actually thinking exactly the same thing the other day.

It would be great to have it for Function too so that you can get a Mixed subfunction viewing the relevant subcomponents. ExtractSubBlock.cofunction would then basically just be calling cofunction[idxs].
There are definitely other places where that would be useful too, for example in the EnsembleFunction I've been working on here.

What does FunctionSpace.collapse() do? It just looks like it would return a copy of the FunctionSpace.

@ksagiyam what do you think about this __getitem__ idea? Currenty, indexing by a scalar returns IndexedProxyFunctionSpace, but indexing by a list of a single item would return the collapsed version of this space.

__getitem__ for Function is implemented in UFL so the changes would be needed there I think.

I'm not too worried for the Function case, that one seems more involved, as they are Expressions in the UFL sense, and extracting a component naturally means a scalar compoment, but here it seems that we want a MixedFunctionSpace (possibly vector-valued) component.

It could also be added to sub or subfunctions, although they both also have issues.
sub explicitly says it is for bcs, and subfunctions is already a tuple not a method so it would probably have to become a local class with __getitem__ overridden so we could deal with single or multiple indices.

That sounds slightly odd to me. Can you instead rewrite collapse() so that it would take a list of indices as an optional argument?

@JHopeCollins collapse() basically forgets parents. They do something nontrivial in ProxyFunctionSpace and in MixedFunctionSpace.

Ok thanks.

It does strip the parent information, I don't think the name collapse would be where I would look for this functionality - it would be more intuitive in subfunctions or sub.

JHopeCollins · 2025-01-09T10:53:26Z

firedrake/matrix_free/operators.py

@@ -135,6 +137,11 @@ def __init__(self, a, row_bcs=[], col_bcs=[],

        self.action = action(self.a, self._x)
        self.actionT = action(self.aT, self._y)
+        # TODO prevent action from returning empty Forms


Is this to do in this PR or at a later date? If it won't be part of this PR please can you open an issue for it to keep a record?

firedrake/slate/slate.py

JHopeCollins · 2025-01-09T14:49:52Z

firedrake/functionspaceimpl.py

@@ -321,6 +321,14 @@ def __iter__(self):
        return iter(self.subfunctions)

    def __getitem__(self, i):
+        from firedrake.functionspace import MixedFunctionSpace
+        if isinstance(i, list):


We should allow any iterable, not just a list.

Suggested change

if isinstance(i, list):

if not isinstance(i, int):

I wanted to mimic numpy.array indexing. tuple indices are used for multi-dimensional arrays, lists give you sub-arrays.

Does that matter here? This is to produce mixed subspaces so it will always be 1D. I don't think it introduces any confusion, and it makes the interface cleaner rather than having a superfluous list at many call sites.

I think the behaviour of __getitem__ is a bit subtle here. Calling obj[1, 2, 3] is exactly the same as calling obj[(1, 2, 3)] which is not quite the same as obj[[1, 2, 3]].

I think having something like space[1, 2] would suggest to someone that space is somehow nested, whereas space[[1, 2]] suggests a subset of the indices.

JHopeCollins · 2025-01-09T14:57:06Z

firedrake/functionspaceimpl.py

@@ -321,6 +321,14 @@ def __iter__(self):
        return iter(self.subfunctions)

    def __getitem__(self, i):
+        from firedrake.functionspace import MixedFunctionSpace


I think we should cache the mixed subspaces like we do for the non-mixed subspaces. Otherwise we get a "different" space every time we call with the same indices.

ksagiyam · 2025-01-09T15:37:17Z

firedrake/functionspaceimpl.py

+            # Return a collapsed subspace if the index is a list
+            if len(i) == 1:
+                return self[i[0]].collapse()
+            else:
+                return MixedFunctionSpace([self[isub] for isub in i])
+


As I said elsewhere, I think this is odd. I feel we are abusing __getitem__(). If we want this feature, we should do it somewhere else (in a new method or in collapse()), not in __getitem__().

I removed the call to collapse(), and will leave V[(i,)] == V[i]

So now V[(i,)] is a non-collapsed indexed space and V[(i, j)] is a new collapsed mixed space?

Not quite. V[(i, )] is no longer collapsed, which means that it is a IndexedProxyFunctionSpace. But collapsing a MixedSpace does nothing, as the new MixedSpace subspace has lost any reference to the parent MixedSpace. I only required collapsing the subspaces because Fieldsplit builds a new NLVP, and we cannot use proxy spaces for BCs.

Please read my comment carefully. I understand why we need this, but I have been claiming that we should not have __getitem__() have multiple meanings depending on if we pass int or tuple/list. If the former returns something which remembers the parent and the latter returns something which does not, it indicates that this feature should be implemented somewhere else.

That's a good point, I hadn't quite appreciated whether the subspace remembers the parent or not. If I understand correctly then we currently have:

subfunction (and by extension __getitem__) returns a FunctionSpace/Function for a single subcomponent that remembers the parent.

collapse returns a FunctionSpace/Function with all components that forgets the parent.

What we want is a utility for returning a FunctionSpace/Function with some subcomponents that forgets the parent. The fact that it returns a subspace matches with subfunction, but the fact that it forgets the parent matches with collapse. This means that if we extend either of the existing methods then we will introduce some ambiguity into that method's purpose.

If we are to avoid this ambiguity then Koki is right, we need a new method given how important clarity is with symbolic information (e.g. the split->subfunction change).

This information should also be added to the docstrings for subfunctions/sub/__getitem__/collapse/the new method for all the FunctionSpace classes plus the Function classes and Cofunction.

DO NOT MERGE

2286596

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from cf743a1 to be6a7b4 Compare January 2, 2025 00:43

pbrubeck changed the title ~~DO NOT MERGE~~ Fieldsplit: replace empty Forms with ZeroBaseForm Jan 2, 2025

pbrubeck marked this pull request as ready for review January 2, 2025 01:07

pbrubeck force-pushed the pbrubeck/simplify-indexed branch 2 times, most recently from e0c7ba1 to 00b80e4 Compare January 2, 2025 02:07

Replace empty Jacobians with ZeroBaseForm

bb04bb0

pbrubeck force-pushed the pbrubeck/simplify-indexed branch 2 times, most recently from 3f399e5 to bf10f91 Compare January 2, 2025 15:46

Split Cofunction

d82039d

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from bf10f91 to d82039d Compare January 2, 2025 16:00

pbrubeck added 2 commits January 2, 2025 10:07

Do not split off-diagonal blocks if we only want the diagonal

af53302

Zero-simplify slate Tensors

7f40504

pbrubeck force-pushed the pbrubeck/simplify-indexed branch 2 times, most recently from 730e299 to 35638df Compare January 3, 2025 19:54

ImplicitMatrixContext: handle empty action

3d06fc5

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from 35638df to 3d06fc5 Compare January 3, 2025 20:32

pbrubeck added 4 commits January 3, 2025 20:27

Only extract constants referenced in the kernel

6078f93

Adjoint: only skip expand_derivatives if necessary

5894b49

style

d99ba50

EquationBC: do not reconstruct empty Forms

d6bb7dd

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from 897cd3c to d6bb7dd Compare January 5, 2025 16:22

pbrubeck mentioned this pull request Jan 6, 2025

Simplify indexed ListTensor objects FEniCS/ufl#336

Open

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from 9c1c516 to d6bb7dd Compare January 6, 2025 19:56

pbrubeck added 2 commits January 6, 2025 17:34

lower degree for EquationBC tests

ed58467

style

2a0c03b

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from d7f0a1b to 2a0c03b Compare January 7, 2025 01:33

ksagiyam reviewed Jan 8, 2025

View reviewed changes

firedrake/adjoint_utils/variational_solver.py Show resolved Hide resolved

firedrake/formmanipulation.py Show resolved Hide resolved

pbrubeck commented Jan 8, 2025

View reviewed changes

firedrake/formmanipulation.py Outdated Show resolved Hide resolved

firedrake/adjoint_utils/variational_solver.py Show resolved Hide resolved

firedrake/formmanipulation.py Show resolved Hide resolved

firedrake/formmanipulation.py Show resolved Hide resolved

JHopeCollins reviewed Jan 9, 2025

View reviewed changes

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from d46a06e to 2eb6a56 Compare January 9, 2025 14:47

JHopeCollins requested changes Jan 9, 2025

View reviewed changes

ksagiyam requested changes Jan 9, 2025

View reviewed changes

FunctionSpace: multiindex returns subspace

934ff6f

pbrubeck force-pushed the pbrubeck/simplify-indexed branch from 2eb6a56 to 934ff6f Compare January 9, 2025 17:48

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fieldsplit: replace empty Forms with ZeroBaseForm #3947

Fieldsplit: replace empty Forms with ZeroBaseForm #3947

pbrubeck commented Dec 30, 2024 •

edited

Loading

github-actions bot commented Dec 30, 2024 •

edited

Loading

github-actions bot commented Dec 30, 2024 •

edited

Loading

JHopeCollins Jan 9, 2025

pbrubeck Jan 9, 2025

pbrubeck Jan 9, 2025 •

edited

Loading

JHopeCollins Jan 9, 2025

pbrubeck Jan 9, 2025 •

edited

Loading

connorjward Jan 9, 2025

pbrubeck Jan 9, 2025 •

edited

Loading

JHopeCollins Jan 9, 2025

ksagiyam Jan 9, 2025

JHopeCollins Jan 9, 2025

JHopeCollins Jan 9, 2025

JHopeCollins Jan 9, 2025

pbrubeck Jan 9, 2025 •

edited

Loading

JHopeCollins Jan 9, 2025

connorjward Jan 9, 2025

JHopeCollins Jan 9, 2025

ksagiyam Jan 9, 2025

pbrubeck Jan 9, 2025

ksagiyam Jan 9, 2025

pbrubeck Jan 9, 2025

ksagiyam Jan 9, 2025

JHopeCollins Jan 10, 2025

Fieldsplit: replace empty Forms with ZeroBaseForm #3947

Are you sure you want to change the base?

Fieldsplit: replace empty Forms with ZeroBaseForm #3947

Conversation

pbrubeck commented Dec 30, 2024 • edited Loading

Description

github-actions bot commented Dec 30, 2024 • edited Loading

github-actions bot commented Dec 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbrubeck Jan 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbrubeck Jan 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbrubeck Jan 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbrubeck Jan 9, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

pbrubeck commented Dec 30, 2024 •

edited

Loading

github-actions bot commented Dec 30, 2024 •

edited

Loading

github-actions bot commented Dec 30, 2024 •

edited

Loading

pbrubeck Jan 9, 2025 •

edited

Loading

pbrubeck Jan 9, 2025 •

edited

Loading

pbrubeck Jan 9, 2025 •

edited

Loading

pbrubeck Jan 9, 2025 •

edited

Loading