Mass-conservative update for non-stiffly-accurate stage_value#226
Mass-conservative update for non-stiffly-accurate stage_value#226sghelichkhani wants to merge 12 commits into
Conversation
For Dt(g(u)) with a nonlinear g (Richards moisture content theta(h) is the canonical case), PR firedrakeproject#204 already builds the conservative two-evaluation form g(U_i) - g(u0) inside getFormStage, so the per-stage equations are correct for any RK method. Mass conservation still failed for non-stiffly-accurate methods like ImplicitMidpoint and GaussLegendre(2) because _update_Ainv computes u_new as a linear combination of u0 and the stage values. For linear Dt(u) that linear combination is algebraically equivalent to the conservative time integration, but for theta nonlinear it is not, and the conservation property the stage equations establish is destroyed at the update step. This routes the non-SA non-Bernstein path through a new conservative update solve when the form has a composite Dt: solve replace(F_dtless, {u0: unew}) - replace(F_dtless, {u0: u0}) + dt * sum_i b_i * F_remainder(stage_i) = 0 for u_new. For g = identity it collapses to the old inner(unew - u0, v)*dx head, so the linear case is unchanged. For nonlinear g it is a discrete statement of the global conservation invariant and u_new satisfies it by construction. Linear Dt(u) keeps _update_Ainv since that path also handles DAE structure correctly, which the conservative variational head does not (the algebraic block is left unconstrained, leading to a singular linear solve). The composite-Dt detection lives in irksome/ufl/manipulation.py as the public helper has_composite_time_derivative(F, u0). It walks TimeDerivative nodes, exempting structural views (Indexed, ListTensor, ComponentTensor), linear differential operators (Grad, Div, Curl), DG facet operators (CellAvg, FacetAvg, restrictions), and arithmetic that is linear in u0 (Sum, Product/Division with at most one operand depending on u0). Pure-time forcing terms Dt(f(t,x)) are passed straight through to the chain rule. The docstring carries an explicit warning that detection is syntactic: pre-interpolating an expression in u0 into a Function and writing Dt(that_function) loses the dependence and silently produces a non-conservative discretisation. The DIRK stage-derivative path folds the chain rule into the K-substitution and has no clean conservative analogue without restructuring. Rather than silently mis-discretising, DIRKTimeStepper now raises NotImplementedError when the form has a composite Dt and points users at stage_type=value. use_collocation_update=True similarly refuses composite Dt because the collocation polynomials terminal value is also a linear combination of stages. New tests in tests/test_conservative_dt.py cover the seven implicit tableaux Irksome ships (BackwardEuler, RadauIIA(2), LobattoIIIC(2), Alexander, ImplicitMidpoint, GaussLegendre(2), QinZhang) on the existing Richards problem with mass error below 1e-10 over twenty steps; the DIRK and use_collocation_update refusals; the linear-arithmetic classification matrix (Dt(2u), Dt(u/2), Dt(u + sin(t)) classified linear; Dt(u*u), Dt(theta(u)), Dt(1/u) classified composite); and identity-runs-and-decays sanity across stage_value and DIRK. Targeted regression on test_dirk, test_dae, test_explicit, test_galerkin, test_disc_galerkin, test_split, test_time_form_splitting, test_mass_conservation, test_pc, test_bern, test_butcher, test_differentiation, test_curl all pass: 243 passed, 2 skipped, 0 failed.
# Conflicts: # irksome/stage_value.py
The auto-detection PR initially refused composite Dt(g(u)) in the DIRK stage-derivative path and in stage_value with use_collocation_update=True, since both produce a u_new that is a linear combination of stages and is therefore not mass-conservative for nonlinear g. That refusal sat outside the scope of this PR (which is about making non-stiffly-accurate stage_value mass-conservative) and was a behaviour change for users who knew about the non-conservation in DIRK and were prepared to live with it. Dropping the guards keeps this PR focused on the one diagnosis it is making and lets the DIRK question be handled separately if and when we choose to fix or warn there. The corresponding two refusal tests are removed; the linear-arithmetic classification matrix and the mass-conservation/identity regressions remain.
Drop tests/test_conservative_dt.py and route its content into the appropriate existing places. * tests/test_mass_conservation.py already runs the exponential Richards problem with a flux BC and asserts mass error near machine precision. Add GaussLegendre(1) (= ImplicitMidpoint), GaussLegendre(2) and QinZhang to its parametrize list. These three tableaux are exactly the ones that exercise the new conservative variational update path for non-stiffly-accurate stage_value -- they fail on master with mass errors of 4e-9 to 9e-8 (the linear-combination update destroying the conservation property the stage equations build for nonlinear theta) and sit at machine precision on this branch. * tests/test_has_composite_time_derivative.py is a new standalone test for the public helper introduced in irksome/ufl/manipulation.py. Two cases: linear-arithmetic forms (Dt(c*u), Dt(u/c), Dt(u + sin(t))) must not be flagged, otherwise the conservative variational head would replace the linear-combination update and break DAEs; truly nonlinear forms (Dt(u*u), Dt(theta(u)), Dt(1/u)) must be flagged, otherwise non-SA stage_value would silently produce a non-conservative discretisation. The identity-regression heat tests previously in test_conservative_dt are dropped as redundant: the existing test_dirk, test_galerkin, test_disc_galerkin and similar suites already exercise linear forms across every tableau the project ships.
Per Pablo's review:
* Rename has_composite_time_derivative -> has_nonlinear_time_derivative.
"Composite" was internal jargon; "nonlinear" is what we are actually
detecting and reads more directly at the call sites. Test file
renamed in lockstep.
* Simplify the conservative update head from
replace(F_dtless, {u0: unew}) - replace(F_dtless, {u0: u0})
to replace(F_dtless, {u0: unew}) - F_dtless: the second replacement
is an identity transformation, so the two forms are byte-identical.
Pablo's suggested phrasing.
* Move the explanation of the two-evaluation conservative head from
inline comments into the docstring of get_update_solver, and state
the math in user-facing terms (inner(g(unew) - g(u0), v) * dx)
rather than via the internal F_dtless symbol. Drop the redundant
"subtracting the two replaced forms keeps v unchanged" comment.
* Compress the warm-start comment in _update_general to a single line
and drop the application-specific (Haverkamp / moisture capacity)
language. The failure mode is generic: g prime of u0 can vanish,
in which case a zero-initialised iterate gives a singular initial
Jacobian.
* Move the update_solver_parameters = solver_parameters default into
the conservative-update branch that actually consumes it. Pablo
pointed out the previous placement implied the inheritance always
applies; in fact the linear _update_Ainv path does not use it, so
scoping the default makes the intent clearer.
| if update_solver_parameters is None: | ||
| update_solver_parameters = solver_parameters |
There was a problem hiding this comment.
There's a potential issue here. What if the solver options for the stage problem use fieldsplit?
There was a problem hiding this comment.
I see! Agreed! Dropped the inheritance. When update_solver_parameters is None, the update solve now uses a fixed default tuned to its own structure rather than copying the stage parameters.
The default is CG + bjacobi/icc rather than LU:
_DEFAULT_UPDATE_SOLVER_PARAMETERS = {
'snes_type': 'newtonls',
'ksp_type': 'cg',
'pc_type': 'bjacobi',
'sub_pc_type': 'icc',
}Reasoning for this: the update Jacobian is g'(u_new) * v * phi * dx so the RK quadrature contributes spatial terms at known stage values, constant in u_new, so they drop from the linearisation. The whole operator is a weighted mass matrix: SPD when g' > 0, condition number bounded independent of mesh size, so CG converges in O(1) iterations with any cheap preconditioner. LU works at small sizes but does not scale; the iterative default is strictly better across the size range and loses nothing at the small end. bjacobi + icc picks up block structure for vector / mixed V automatically and degenerates gracefully on scalar V.
Two things I would value your view on:
- Is there an Irksome-wide convention for default
solver_parameters? I would rather match an existing pattern than introduce a one-off. - CG + bjacobi/icc bakes in monotone
g(g' > 0a.e.). Non-monotonegbreaks SPD and needs an explicit override (GMRES + a more general preconditioner). Comfortable with that assumption as the default, or would you prefer GMRES + bjacobi/ilu as the conservative choice at the cost of slightly slower convergence on the canonical case?
Docstring on get_update_solver now states the no-inheritance asymmetry explicitly so it does not surprise anyone.
There was a problem hiding this comment.
Irksome does not change the firedrake default solver parameters. I don't see why we should be changing them in any particular instance.
Firedrake by default uses a sparse direct solver, depending on which packages are installed. If the user wants an iterative solver they should set it explicitly.
There was a problem hiding this comment.
All sensible. Now reverted in 05e0d8e -- update_solver_parameters defaults to None, which lets NonlinearVariationalSolver pick up Firedrake's default. Test updated to pass tight tolerances through explicitly.
Per Pablo's review on manipulation.py: the bespoke walker with its
_PASS_THROUGH_OPS exemption list and _is_pass_through / _depends_on
helpers is a maintenance liability -- every linear UFL operator added
upstream silently misclassifies until Irksome's exemption list is
updated to match.
ufl.derivative already knows the linearity of Grad, Div, Indexed,
ListTensor, ComponentTensor, restrictions, CellAvg, FacetAvg, and any
future linear node UFL adds. Reformulating the classifier as
D = expand_derivatives(derivative(f, u0, TrialFunction(V)))
if u0 in extract_coefficients(D): nonlinear
inherits that classification for free and collapses the walker to its
extract_type(F, TimeDerivative) loop plus three lines.
Verified against tests/test_has_nonlinear_time_derivative.py (all
seven cases agree with the walker) and tests/test_mass_conservation.py
(all nine tableaux still conserve to machine precision).
Per Pablo's review on stage_value.py: silently inheriting
solver_parameters into the conservative update solve is wrong-answer
unsafe, not just suboptimal. Stage-tuned options can fail badly when
applied to the update problem -- fieldsplit indices keyed to V^s do
not match V, snes_type='ksponly' returns a wrong answer when the
update head is nonlinear, lagged Jacobians and custom MG transfers
are tuned to operators that have nothing to do with the update.
Drop the inheritance. When update_solver_parameters is None, use a
fixed default tuned to the update problem itself. The Jacobian is
g'(u_new) * v * phi * dx
(the RK quadrature contributes spatial terms at known stage values,
constant in u_new, so they drop from the linearisation) -- a weighted
mass matrix, SPD when g' > 0, with condition number bounded
independent of mesh size. For that operator the right default is
Newton + CG + bjacobi/icc: scales to any mesh size, single-digit
iteration counts in the canonical case, no failure modes that LU
would handle better.
Non-monotone g breaks SPD and needs an explicit override (GMRES +
a more general preconditioner). Points where g' vanishes (plateau
soils with theta' = 0 in the saturated zone) trip a zero pivot
regardless of preconditioner choice -- a separate structural issue.
Docstring on get_update_solver now states the no-inheritance
asymmetry explicitly so users do not assume copy-through behaviour.
Co-authored-by: Pablo Brubeck <brubeck@protonmail.com>
Per Pablo's follow-up on the update solver default: Irksome does not override Firedrake's default solver_parameters anywhere else, and should not introduce a one-off here. When update_solver_parameters is None, NonlinearVariationalSolver picks up Firedrake's default (typically a sparse direct solve). Users who want CG + bjacobi/icc or anything else configure it explicitly. tests/test_mass_conservation.py was previously inheriting tight snes_rtol/snes_atol from the stage solver into the update solve; with the inheritance dropped, the conservation assertion (10 * eps) fails by a factor of ~1 because the default SNES rtol of 1e-8 is the binding constraint on the update solve precision rather than machine arithmetic. Pass update_solver_parameters explicitly for stage_value runs so the assertion measures conservation rather than SNES tolerance. The DG path does not accept update_solver_parameters (it has no update solve in the same sense) so the kwarg is gated on stage_type.
These were exemption-list classes for the bespoke walker that 68de837 replaced with the ufl.derivative probe. Walker is gone, classes are unused.
Summary
@rckirby — this is the follow-up to our Slack conversation about implicit midpoint and mass conservation. Your hand-derivation was right and there was a bug, just subtler than "the discretisation is wrong". The per-stage equations are fine; the bug lives in the update step that produces
u_newfrom the stages.Diagnosis
For
inner(Dt(g(u)), v)*dx + spatial(u)with nonlinearg(Richards moisture contenttheta(h)is the canonical case), PR #204 already builds the per-stage equations as the two-evaluation formg(U_i) - g(u_0)insidegetFormStage. So the stage-level conservation property is correct for any RK method — exactly what your hand-derivation predicts at the PDE level.Conservation was still being lost for non-stiffly-accurate methods (
GaussLegendre(1)= implicit midpoint,GaussLegendre(2),QinZhang). The cause is in_update_Ainv: after the stages are computed, Irksome formsu_new = (1 - sum bAinv) u_0 + sum bAinv U_i, a linear combination ofu_0and the stage values. For linearDt(u)that is algebraically equivalent to the conservative time integration. ForDt(theta(h))it is not:thetaof a linear combination is not the linear combination of per-stagethetaincrements, so the conservation property the stage equations establish is destroyed at the update step. Stiffly accurate methods avoid this entirely becauseu_new = U_sis identically the last stage value, so the conservative finite difference sits in that stage's equation and gets copied straight tou_new. That is why RadauIIA, LobattoIIIA/C, BackwardEuler and Alexander all look fine and only the non-SA methods flagged.Fix
Route the non-SA path through a conservative variational solve when the form has a composite
Dt:solved for
u_new. Forg = identityit collapses to the originalinner(u_new - u_0, v)*dxhead, so nothing changes for the linear case. For nonlineargit is a discrete statement of the global conservation invariant andu_newsatisfies it by construction. Order is preserved: the stages are still orderp, the RK quadrature on the spatial term is unchanged, and inverting the conservation equation foru_newpreserves order providedgis sufficiently smooth.Linear
Dt(u)deliberately keeps_update_Ainv. The new path is algebraically identical to the linear-combination update forg = identity(verified by direct calculation), and_update_Ainvhandles DAE structure correctly while the conservative variational head does not (the algebraic block is unconstrained and the linear solve is singular). Composite-Dt detection is the right syntactic gate for routing — it lives inirksome/ufl/manipulation.pyashas_composite_time_derivative, with documented exemptions for structural views, linear differential operators, DG facet operators, and arithmetic that is linear inu_0.What is intentionally not in this PR
In line with your earlier comment about DIRKs — that we run DIRK in stage-derivative form today but could plausibly do a stage-value form for them at some point, and that the right play is to push fully-implicit solvers in
stage_valueand let accuracy buy back the runtime — this PR does not touchDIRKTimeStepper. Stage-derivative DIRK folds the chain rule into the K-substitution and has no clean conservative analogue without restructuring; users who need conservation can run any DIRK tableau throughstage_type='value'instead. There is a telescoping W-ladder fix on a separate branch but it belongs to its own conversation.Test
The existing
tests/test_mass_conservation.pyalready runs the exponential Richards problem with a flux BC and asserts mass error near machine precision. I extended the parametrize list to includeGaussLegendre(1),GaussLegendre(2)andQinZhangalongside the existing BackwardEuler / RadauIIA(2) / BDF / AdamsMoulton entries.Master Irksome with the three new entries:
This branch: all nine sit at machine precision.
The failure on master is the bug; the pass on this branch is the fix. The seven-tableau spread covers single-stage and multi-stage non-SA methods, order 2 and order 4, plus the negative-
bAlexander case as a sign regression on the SA path.tests/test_has_composite_time_derivative.pypins both directions of the classifier contract: linear-arithmetic forms (Dt(c*u),Dt(u/c),Dt(u + sin(t))) must not be flagged (else DAEs break), nonlinear forms (Dt(u*u),Dt(theta(u)),Dt(1/u)) must be flagged (else non-SAstage_valuesilently mis-discretises).Wider regression:
tests/test_dirk,test_dae,test_explicit,test_galerkin,test_disc_galerkin,test_split,test_time_form_splitting,test_pc,test_bern,test_butcher,test_differentiation,test_curlall pass.Cross-validation outside Irksome
The same test runs through G-ADOPT's Richards solver wrapper on a Haverkamp soil model: a function-space sweep on BackwardEuler is the existing
test_mass_balancetest, and a newtest_mass_balance_tableauxparametrising over BE, ImplicitMidpoint and GaussLegendre(2) at DQ1 produces the same master-fail / branch-pass pattern. File:https://github.com/g-adopt/g-adopt/blob/48a8cf44/tests/richards/test_mass_balance.py#L145-L168