diff --git a/AGENTS.md b/AGENTS.md index de62cc1d29..f6fb08f064 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,8 +1,10 @@ # Firedrake Firedrake is an automated system for the portable solution of partial differential equations using -the finite element method (FEM). The codebase is primarily Python, relying heavily on code generation -and high-performance C backends to achieve scalability and speed. +the finite element method. The codebase is primarily Python, relying heavily on code generation +and high-performance C backends to achieve scalability and speed. Firedrake is highly composable and +fully differentiable, enabling the automatic generation of tangent linear and adjoint models for +PDE-constrained optimization. Firedrake's full contribution process is documented at [Contributing to Firedrake](https://firedrakeproject.org/contribute.html). In short, for AI-assisted @@ -30,6 +32,12 @@ toolchain: 2. **Lowering to Loopy:** The GEM expressions are then lowered into **loopy** kernels. * **PyOP2:** Finally, the generated loopy kernels are wrapped and executed by PyOP2, which handles the parallel execution of loops over mesh cells and facets. +* **pyadjoint:** Firedrake integrates with `pyadjoint` for algorithmic differentiation. Annotated + operations (assembly, interpolation, variational solves, boundary condition application, ...) are + recorded on a `Tape` as a DAG of `Block`s; a `ReducedFunctional` composed with one or more `Control`s + can then be evaluated, differentiated (adjoint or tangent-linear), and checked with a `taylor_test`. + Taping happens at the level of Firedrake's own operations, not the underlying numerics, so any new + feature assembled purely from already-annotated building blocks is differentiable automatically. ## Core Working Rules @@ -56,6 +64,13 @@ toolchain: prose explaining what the removed, incorrect approach used to do or why it was wrong. Keep comments and documentation focused on the current, correct code; a reader should never need the history of what used to be there to understand why the present code is right. +* **Composability And Differentiability:** New features are expected to compose with existing ones + without special-casing, and, via `pyadjoint`, to remain differentiable when built from already-taped + operations. Prefer the annotated, top-level API (e.g. `firedrake.assemble`) over a lower-level + equivalent that bypasses it (e.g. calling an `Interpolator`'s own `.assemble()` directly) even when + both give the same forward numbers — the lower-level call silently drops out of the tape, and no test + will notice unless it specifically exercises pyadjoint. When a change could plausibly sit on the tape, + verify differentiability explicitly with a `taylor_test`, not just a forward-value check. ## Coding Style And Conventions @@ -77,6 +92,29 @@ toolchain: * **Docstrings:** All public-facing APIs must include properly formatted `numpydoc`-style docstrings. * **Type Hints:** New code should include type hints on function/method signatures. +## Design And Debugging Method + +How to spend effort on any feature or bug: the first three shape a design before writing code, +the last two localize a failure before reading code. + +* **Design by nearest working neighbor.** Some existing feature already solves a structurally + identical problem. Grep for the invariant yours must satisfy (the type handled, the hook fired, + the kwarg accepted), read how the neighbor earns it, and implement only the delta. +* **Write the state contract before the code.** For anything flowing through a cached, replayed, + or lazily-refreshed system, answer up front: who owns it, when is it refreshed, what happens on + reuse or in-place mutation, how does a replay recover it. Stale-state bugs fail far from their + cause and only under reuse. +* **Classify the mathematical structure before choosing machinery.** Affine and linear + dependencies have closed-form contributions (identities, fixed operators, exact zeros for higher + derivatives); recognizing an exact zero lets you skip a code path instead of fixing it. +* **Attribute before you analyze.** Shrink *where* with one-delta experiments — each ingredient + toggled in isolation against a known-good baseline, consecutive CI failure sets diffed, a single + hunk reverted, the merge base rerun — then read code for *why*. +* **Census the consumers before changing a lifecycle.** Changing *when* or *how often* something + is computed (rather than its value) is safe only after grepping every access site: some consumer + calls it in a context you did not design for (per nonlinear iteration, inside a PETSc callback, + on every attribute read). + ## Testing Requirements * **Pull Requests:** All PRs must include comprehensive tests demonstrating that the new feature works @@ -104,6 +142,12 @@ toolchain: in the install docs to get a component installed in editable mode so source edits take effect without reinstalling, and check which branch/commit of each component is actually active before assuming a fix belongs in Firedrake itself. +* **Branch pairing across the stack:** Firedrake's `main` and `release` branches go hand in hand with + the `main` and `release` branches of its components (FIAT, UFL, ...). A CI failure may be + unreproducible locally simply because a component checkout in the venv sits on some other branch — + check `git -C $VIRTUAL_ENV/src/ branch`, switch to the branch matching the Firedrake branch + under test, run `firedrake-clean`, and reproduce again before hunting for the bug in Firedrake + itself. * **`petsc4py`/PETSc version skew:** `petsc4py` is a compiled extension built against one specific PETSc checkout. If you switch the PETSc branch/commit underneath an existing venv (e.g. to bisect a PETSc-side issue) without rebuilding `petsc4py` against it, `import firedrake` fails with a confusing @@ -140,12 +184,25 @@ toolchain: conclude a parallel code path is untested just because a plain, unmarked `pytest` run was green. * **Splitting for CI:** `firedrake-run-split-tests` shards the suite by process count for CI; look at it (and `.github/workflows/pr.yml`/`core.yml`) if a failure only reproduces in CI and not locally. +* **CI triage:** `gh pr checks ` lists job statuses. When `gh run view --log` returns nothing + (it does for very large logs), download the log with + `gh api repos///actions/jobs//logs` and grep for `FAILED`. Before debugging + anything, fetch the failure list of the *previous* run of the same PR: the difference between the + two failure sets attributes each failure to the commits pushed in between. * **Narrow reproduction first:** Run the single failing test node (`pytest path::test_name -k ...`) before the full module; the suite is large and full-module reruns are slow to iterate against. +* **Test mathematical correctness, not just that it runs or looks structurally right.** Neither "no + exception was raised" nor `==` agreement between two expressions proves the result is + correct — two independently-built expressions can match structurally while sharing the same wrong + derivative or simplification rule. Verify the actual mathematical claim: evaluate numerically and + compare against a hand-computed or finite-difference value, or use a Taylor test for anything + claiming to be a derivative. +* **Taylor-test-everything is the immune system:** Taylor-test a `ReducedFunctional`, ensuring that any + new feature built from existing, annotated Firedrake operations is automatically differentiable. ### Debugging -* **Generated kernels (niche, rarely needed):** By default, generated C is compiled optimized and +* **Generated kernels:** By default, generated C is compiled optimized and without debug symbols, so a debugger attached to the Python process cannot meaningfully step through it. Set `PYOP2_DEBUG=1` to compile with `-O0 -g` instead, which is the prerequisite for using `gdb`/`cgdb` on the compiled kernel at all. @@ -156,7 +213,7 @@ toolchain: computed differently per rank and fed into code generation (e.g. a rank-local decision that should be a collective/global one) — make that decision the same on every rank, rather than patching the generated source or the difference itself. -* **Parallel deadlocks (niche, rarely needed):** `PYOP2_SPMD_STRICT=1` adds barriers around calls +* **Parallel deadlocks:** `PYOP2_SPMD_STRICT=1` adds barriers around calls marked `@collective` and around cache access, trading overhead for a much narrower failure point when ranks disagree about control flow. * **Logging:** `firedrake.logging.set_log_level()` (or the `PYOP2_LOG_LEVEL` environment variable) @@ -165,6 +222,23 @@ toolchain: standard PETSc options (`-ksp_view`, `-snes_view`, `-ksp_monitor`, `-log_view`, `-start_in_debugger`) can be passed through Firedrake's `solver_parameters` or the command line exactly as in a plain PETSc application. +* **Errors inside PETSc callbacks do not surface as their own traceback:** under pytest they often + appear as a bare `Segmentation fault` with no Firedrake frames; standalone they appear as + `petsc4py.PETSc.Error: error code 101` whose *first* chained traceback (e.g. a `TypeError` about + the callback context in `petscsnes.pxi`) describes the corrupted callback state, not the cause. + Rerun the failing test as a standalone script to expose the chained tracebacks, and read the PETSc + call stack inside the error (`PCSetUp_MG` → `SNESComputeFunction`, ...) to identify *which* + callback was executing. +* **Construction-time code re-runs inside solver callbacks:** geometric multigrid coarsens the + entire problem lazily inside `PCSetUp`, through the `coarsen` singledispatch in + `firedrake/mg/ufl_utils.py` — whatever a feature does at construction time (e.g. `DirichletBC` + interpolating or projecting its boundary value into the space) is re-executed per level inside + that PETSc callback. Objects that carry solvers or attach DM hooks (a `Projector`, a variational + solver) must be built once and cached, never rebuilt on each call of an accessor that may fire + there: repeatedly constructing and garbage-collecting a solver stack inside `PCSetUp_MG` corrupts + the DM callback state and segfaults far from the allocation site. Extruded (hexahedral) + hierarchies are the stress test: tensor-product elements (NCE/NCF) have no dual-basis + interpolation, so paths that interpolate on simplices take the projection fallback there. ### Reproducible Environments diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000000..43c994c2d3 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1 @@ +@AGENTS.md diff --git a/firedrake/adjoint/AGENTS.md b/firedrake/adjoint/AGENTS.md new file mode 100644 index 0000000000..d3c957bab8 --- /dev/null +++ b/firedrake/adjoint/AGENTS.md @@ -0,0 +1,71 @@ +# Firedrake Adjoint + +Two properties are treated as defining features of Firedrake, not optional extras: **composability** +(function spaces, forms, boundary conditions, solvers, and preconditioners are expected to combine +freely, without special-casing particular combinations) and **differentiability** (via `pyadjoint`, +essentially any computation built from Firedrake's own operations — assembly, interpolation, +variational solves, boundary condition application — can be taped and differentiated end-to-end). +Differentiability is expected to fall out of composability: a new feature built from already-annotated +Firedrake operations should be differentiable for free, with no extra work. A feature that instead +reaches past those operations into a lower-level, unannotated API silently breaks this guarantee — +forward runs stay numerically correct, but the adjoint quietly goes wrong, and nothing fails until +pyadjoint is specifically exercised. + +## Coding Style And Conventions + +* **Keep Annotation Out Of Plain Modules:** pyadjoint bookkeeping never appears inside a method body in + `firedrake/*.py`. Differentiable types instead have a `*Mixin` in `firedrake/adjoint_utils/` exposing + one `_ad_annotate_` decorator per method that needs taping, applied where the method is defined + (`@SomeMixin._ad_annotate_foo` on `def foo`). The decorator wraps the whole method, so the real + implementation stays pyadjoint-agnostic. Add a new decorator to the `Mixin` rather than calling `_ad_*` + from `firedrake/*.py` directly. + +## Design And Debugging Method + +Guiding principles for building and debugging taped operations with `firedrake.adjoint`, in the +order they should be applied: + +* **Adjoints come from composition, not re-derivation.** If implementing a block's + `evaluate_adj_component` has you calling `derivative`/`adjoint`/`action` on an operation + Firedrake already tapes (assembly, interpolation, projection, a solve), the tape structure is + wrong, not incomplete: make the block depend on that operation's *output*, and the operation's + own block supplies the adjoint, TLM, Hessian, and recompute. A block's dependency is the value + its operation actually consumes, never the raw user input that value was derived from. +* **Tape derived state when its consumer is taped, not when it is computed.** A value lazily + re-derived from a mutable input must be re-annotated at the moment the consuming block is + recorded, so the dependency edge points at the input's *current* block variable; for a + `FloatingType`, override `_ad_will_add_as_dependency` to refresh (and thereby tape) the value + before `super()` tapes the block. Object reuse in a time loop then records one correct chain per + step with no extra bookkeeping. Conversely, run internal updates at construction/setter time + under `stop_annotating()`: taping work nothing depends on leaves dangling blocks that recompute + pays for. +* **Prove the linchpin primitive in isolation before restructuring around it.** Check in a few + lines that the symbolic machinery you plan to rely on (e.g. + `assemble(action(adjoint(derivative(form, c)), cof))`) accepts every input class you must + handle. This surfaces hard limits early, and often shows that existing special-case code is + subsumed by the general path — or was already dead. +* **Fix the block class that actually runs, not just the base.** Solves taped through solver + objects execute `NonlinearVariationalSolveBlock`, which overrides several `GenericSolveBlock` + methods (`prepare_evaluate_adj`, `evaluate_adj_component` with its own `_dFdm_cache`, and + `_ad_assign_map`, which refreshes its cached cloned solvers by matching coefficients across + clones via `.count()`); a case added only to the generic method silently never executes there. + Print `type(block)` from the tape and grep the subclass for overrides before editing the base. +* **Verify the structure before the numbers.** Print each block in + `get_working_tape().get_blocks()` with the identities of its `.get_dependencies()` and + `.get_outputs()`, and check the DAG is exactly the chain you designed — no stale or dangling + block variables — for both freshly-created and reused objects. Only then run `taylor_test` with + a genuinely nonlinear control: rate ≈ 2 is a pass; residuals all ~1e-16 mean the functional is + accidentally linear in the control (square it); rate ≈ 1 means a stale value reached the tape. + A zero gradient is named by the warning `Adjoint value is None, is the functional independent of + the control variable?` — any later `ZeroDivisionError`/`nan` is just its echo. In a + `taylor_to_dict` check, rates that start correct then collapse with residuals stuck at a small + floor mean a small absolute error in the highest-order term supplied (a nearly-right Hessian or + gradient); attribute it by re-running the same Taylor test over a matrix of feature toggles + (one new argument at a time) against a known-good baseline. +* **Dependency discovery is structural, not conceptual.** `form.coefficients()` finds `Function`s + (including on `"R"` spaces) but not `firedrake.Constant`, so a bare `Constant` silently records + no dependency — represent differentiable scalars as `Function`s on an `"R"` space. +* **A cache reused across recomputes must resync every axis that varies between reuses.** A cached + solver refreshes some of its inputs automatically (form coefficients) but not others (its + problem's boundary conditions); each missed axis silently reuses stale data under a perturbed + control.