Skip to content

Record projection solver parameters on the tape#5175

Open
sghelichkhani wants to merge 2 commits into
mainfrom
sghelichkhani/projectblock-solver-parameters
Open

Record projection solver parameters on the tape#5175
sghelichkhani wants to merge 2 commits into
mainfrom
sghelichkhani/projectblock-solver-parameters

Conversation

@sghelichkhani

Copy link
Copy Markdown
Contributor

Fixes #5172.

When a project call was annotated, the resulting ProjectBlock recorded no solver configuration at all: annotate_project forwarded only the adjoint callback kwargs, the Function.project path forwarded nothing, and the Projector's own cg/icc defaults are applied after annotation has already happened. As a result every taped projection was replayed and adjointed with the global firedrake defaults, i.e. a direct solve with mumps, regardless of what the forward projection actually did. The issue has a small example where the replayed functional differs from the taped one by three orders of magnitude.

The first commit extracts the Projector's parameter defaulting into resolve_projection_solver_parameters and has the annotation record the effective parameters of the forward solve on the block. The forward kwargs carry the solver and form compiler parameters, while the adjoint kwargs carry only the solver parameters, since the adjoint solve goes through the assembled-matrix path which accepts no form compiler parameters; the mass matrix is self-adjoint so reusing the forward parameters is the right thing. The duplicated annotation wrapper on Function.project is removed in favour of the annotation on projection.project, to which it delegates, so both entry points now record the same block. Supermesh projections pass the parameters through to the Projector the block already holds.

Note this deliberately changes replay behaviour: an un-parameterised taped projection now replays with the cg/icc mass solve it actually performed forward (to ksp_rtol 1e-8) instead of an exact direct solve. That is the point of the fix, but it may shift replayed values within solver tolerance for tests that relied on the old behaviour.

The second commit addresses the second inconsistency noted in the issue: projecting a Function onto its own space shortcuts to an assignment in the forward model, with no solve at all, while the tape recorded a ProjectBlock that performed a mass solve on every replay. The tape now records a FunctionAssignBlock, matching the forward model.

This touches firedrake/adjoint_utils/blocks/solving.py territory adjacent to the CachedSolverBlock work in #4638; I have left a note there. A more structural follow-up, where the block holds the Projector itself so that replay reuses its cached solver (as SupermeshProjectBlock already does), is left for after #4638 settles.

Previously annotate_project forwarded only the adjoint callback kwargs to
the ProjectBlock, so the solver parameters used by the forward projection
(and the Projector's cg/icc defaults, which are applied after annotation)
never reached the tape. Replay and adjoint of any taped projection
silently fell back to the global firedrake defaults (preonly + lu).

The Projector's parameter defaulting is extracted into
resolve_projection_solver_parameters so the annotation can record the
effective parameters of the forward solve. The forward kwargs carry the
solver and form compiler parameters; the adjoint kwargs carry only the
solver parameters, since the adjoint solve takes an assembled matrix. The
mass matrix is self-adjoint so the forward parameters are reused.

Function.project was annotated separately with a duplicated wrapper that
forwarded no kwargs at all; it now relies on the annotation of
projection.project, to which it delegates.

Fixes #5172
Projecting a Function onto its own space (without boundary conditions)
shortcuts to an Assigner in the forward model, so no solve happens.
The tape nevertheless recorded a ProjectBlock, which performed a mass
solve on every replay and adjoint evaluation. Record a
FunctionAssignBlock instead, matching what the forward model does.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Solver parameters passed to project are never recorded on the ProjectBlock

1 participant