Mass-conservative update for non-stiffly-accurate stage_value

[sghelichkhani]

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_new from the stages.

Diagnosis

For inner(Dt(g(u)), v)*dx + spatial(u) with nonlinear g (Richards moisture content theta(h) is the canonical case), PR #204 already builds the per-stage equations as the two-evaluation form g(U_i) - g(u_0) inside getFormStage. 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 forms u_new = (1 - sum bAinv) u_0 + sum bAinv U_i, a linear combination of u_0 and the stage values. For linear Dt(u) that is algebraically equivalent to the conservative time integration. For Dt(theta(h)) it is not: theta of a linear combination is not the linear combination of per-stage theta increments, so the conservation property the stage equations establish is destroyed at the update step. Stiffly accurate methods avoid this entirely because u_new = U_s is identically the last stage value, so the conservative finite difference sits in that stage's equation and gets copied straight to u_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:

replace(F_dtless, {u_0: u_new}) - replace(F_dtless, {u_0: u_0})
    + dt * sum_i b_i * F_rem(stage_i) = 0

solved for u_new. For g = identity it collapses to the original inner(u_new - u_0, v)*dx head, so nothing changes for the linear case. For nonlinear g it is a discrete statement of the global conservation invariant and u_new satisfies it by construction. Order is preserved: the stages are still order p, the RK quadrature on the spatial term is unchanged, and inverting the conservation equation for u_new preserves order provided g is sufficiently smooth.

Linear Dt(u) deliberately keeps _update_Ainv. The new path is algebraically identical to the linear-combination update for g = identity (verified by direct calculation), and _update_Ainv handles 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 in irksome/ufl/manipulation.py as has_composite_time_derivative, with documented exemptions for structural views, linear differential operators, DG facet operators, and arithmetic that is linear in u_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_value and let accuracy buy back the runtime — this PR does not touch DIRKTimeStepper. 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 through stage_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.py already runs the exponential Richards problem with a flux BC and asserts mass error near machine precision. I extended the parametrize list to include GaussLegendre(1), GaussLegendre(2) and QinZhang alongside the existing BackwardEuler / RadauIIA(2) / BDF / AdamsMoulton entries.

Master Irksome with the three new entries:

Scheme	Mass error
BackwardEuler, RadauIIA2, BDF1, AM0..2	machine precision
ImplicitMidpoint	8.88e-08
GaussLegendre(2)	4.78e-09
QinZhang	3.94e-09

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-b Alexander case as a sign regression on the SA path.

tests/test_has_composite_time_derivative.py pins 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-SA stage_value silently 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_curl all 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_balance test, and a new test_mass_balance_tableaux parametrising 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

[pbrubeck]

There's a potential issue here. What if the solver options for the stage problem use fieldsplit?

[sghelichkhani]

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:

1. Is there an Irksome-wide convention for default solver_parameters? I would rather match an existing pattern than introduce a one-off.
2. CG + bjacobi/icc bakes in monotone g (g' > 0 a.e.). Non-monotone g breaks 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.

[pbrubeck]

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.

[sghelichkhani]

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.