refactor interpolation #4595
refactor interpolation #4595
Conversation
|
@leo-collins this looks like it includes changes from the other ongoing interpolation PRs. This will be much easier to review if you point this PR at that branch instead of |
leo-collins
force-pushed
the
leo/refactor_interpolate
branch
from
deb0a8c to
4097207
Compare
pbrubeck
changed the base branch from
main
to
pbrubeck/interp-adjoint-explicit
This is fine, but we might want to preserve a user-facing interface for reusable interpolators. |
pbrubeck
force-pushed
the
pbrubeck/interp-adjoint-explicit
branch
7 times, most recently
from
210bae0 to
8d631f8
Compare
firedrake/interpolation.py
Outdated
| Parameters | ||
| ---------- | ||
| expr : ufl.Interpolate | ufl.ZeroBaseForm | ||
| The symbolic interpolation expression, or a zero form. Zero forms |
There was a problem hiding this comment.
"zero form" is too close to 0-form, which is not necesarily equal to 0.
| The symbolic interpolation expression, or a zero form. Zero forms | |
| The symbolic interpolation expression, or a ZeroBaseForm. ZeroBaseForms |
|
|
||
| # NOTE: The par_loop is always over the target mesh cells. | ||
| target_mesh = as_domain(V) | ||
| target_mesh = V.mesh() |
There was a problem hiding this comment.
Now we might want to reconsider leaving as_domain here...
There was a problem hiding this comment.
Are we sure this has been properly resolved?
There was a problem hiding this comment.
It'll be useful to have some assertions here
firedrake/interpolation.py
Outdated
| if access == op2.INC: | ||
| copyin += (tensor.zero,) |
There was a problem hiding this comment.
For the 1-form adjoint on a mixed space, we might be adding different sub-interpolates onto the same tensor. The zeroing should be handled outside the for loop of line 633
There was a problem hiding this comment.
tensor is already the indexed sub-tensor in this case. Can you write a test for this?
leo-collins
force-pushed
the
leo/refactor_interpolate
branch
from
89550eb to
9fd2344
Compare
connorjward
left a comment
connorjward
left a comment
There was a problem hiding this comment.
I haven't reviewed this in great detail but this seems really good!
| """ | ||
| from firedrake.assemble import assemble | ||
| # If assembling the operator, we need the concrete permutation matrix | ||
| matfree = False if self.rank == 2 else self.matfree |
There was a problem hiding this comment.
This seems bizarre. A user could pass a matfree option that gets ignored.
There was a problem hiding this comment.
We can't assemble the cross-mesh interpolation matrix if matfree=True. By doing this the user can do assemble(interpolate(TrialFunction(V), U)) and get the operator without having to pass an additional matfree=False.
There was a problem hiding this comment.
In which case I think we should change (and document) the interface to
def interpolate(expr, ..., *, matfree=None):
if matfree is None:
if some_condition:
matfree = False
else:
matfree = Truebecause currently the user could run
assemble(interpolate(TrialFunction(V), U, matfree=True))and the option would be silently ignored.
There was a problem hiding this comment.
The same could be said about the allow_missing_dofs and default_missing_val parameters if the user is doing same-mesh interpolation. I think it's fine if the docstring says something like "this parameter only applies if interpolating between a VertexOnlyMesh and its input_ordering".
There was a problem hiding this comment.
Well the right answer there is to disallow those options for same-mesh interpolation. I don't want to force you to make big changes here though so I'd be happy provided that this behaviour is made clear in the docstring and you open an issue about it.
| for indices, form in firedrake.formmanipulation.split_form(expr): | ||
| if isinstance(form, ufl.ZeroBaseForm): | ||
| # Get sub-interpolators and sub-bcs for each block | ||
| Isub: dict[tuple[int, int], tuple[Interpolator, list[DirichletBC]]] = {} |
There was a problem hiding this comment.
This is a great idea. Big fan.
There was a problem hiding this comment.
I'm not a fan. Why can we not have the BCs passed on to the Interpolator?
There was a problem hiding this comment.
I guess that's just a matter of taste. We don't generate code for the boundary conditions so stashing things on the Interpolator is less important.
I was specifically praising the use of type hints to make an extremely complex data structure clear.
There was a problem hiding this comment.
We can't change the BCs on matrices on repeated assembly calls, or can we?
There was a problem hiding this comment.
Also the tuple[int, int] could in some cases just be tuple[int]
There was a problem hiding this comment.
We can't change the BCs on matrices on repeated assembly calls, or can we?
Good point. I'm not sure. We switch the lgmaps out at runtime so feasibly it's OK to do.
There was a problem hiding this comment.
My thought was that BCs should be passed to assemble, like we do for all other forms
dham
left a comment
dham
left a comment
There was a problem hiding this comment.
subject to checking that adjoint interpolation on mixed spaces is really correct, this is great. thanks.
| sub_tensor.assign(sum(self[i]._interpolate(*function, adjoint=adjoint, **kwargs) | ||
| for i in self if i[0] == k)) | ||
| return output | ||
| Isub[indices] = (get_interpolator(form), sub_bcs) |
There was a problem hiding this comment.
Can we not do this instead?
| Isub[indices] = (get_interpolator(form), sub_bcs) | |
| Isub[indices] = get_interpolator(form, bcs=sub_bcs) |
There was a problem hiding this comment.
Or if that is not your intended design, I suggest we enable something like this
| Isub[indices] = (get_interpolator(form), sub_bcs) | |
| Isub[indices] = get_assembler(form, bcs=sub_bcs) |
There was a problem hiding this comment.
A third option is to do something like
| Isub[indices] = (get_interpolator(form), sub_bcs) | |
| Isub[indices] = partial(get_interpolator(form).assemble, bcs=sub_bcs) |
firedrake/interpolation.py
Outdated
|
|
||
| operand, = expr.ufl_operands | ||
| target_mesh = expr.target_space.mesh() | ||
| source_mesh = extract_unique_domain(operand) or target_mesh |
There was a problem hiding this comment.
We should also consider the possibility of the operand having a MeshSequence. Right now extract_unique_domain will raise an Exception, but MixedMeshInterpolator should be able to handle the case where operand is an Argument of a MixedFunctionSpace. The case where operand is a more general Expr on multiple meshes might involve more complicated code generation.
This should be done in a separate PR.
| coarse_space_bcs = tuple(coarse_space_bcs) | ||
| if G_callback is None: | ||
| interp_petscmat = chop(Interpolator(dminus(trial), V, bcs=bcs + coarse_space_bcs).callable().handle) | ||
| interp_petscmat = chop(assemble(interpolate(dminus(trial), V), bcs=bcs + coarse_space_bcs).mat()) |
There was a problem hiding this comment.
Why not assemble(...).petscmat?
fixes
| self.target_space = dual_arg.function_space().dual() | ||
| """The primal space we are interpolating into.""" | ||
| self.target_mesh = self.target_space.mesh() | ||
| self.target_mesh = self.target_space.mesh().unique() |
There was a problem hiding this comment.
I think MixedInterpolator will be fine with MeshSequence, so this call should go in the subclasses
| dual_arg, operand = expr.argument_slots() | ||
| V = dual_arg.arguments()[0].function_space() | ||
|
|
||
| V = dual_arg.function_space().dual() |
There was a problem hiding this comment.
| V = dual_arg.function_space().dual() | |
| assert isinstance(dual_arg, Cofunction | Coargument) | |
| V = dual_arg.function_space().dual() |
Simplifies interpolation code and introduces new features. So far:
InterpolateOptionsdataclass. We get type hinting, better IDE support, single source of truth for these options.__new__method and dispatch instead with aget_interpolatorfunction.