O2.3.6 Implicit solver for full soil model

[kmdeck]

Purpose

We need an implicit solver for the soil model (both for Richards equation alone, and the full soil system), and possibly in the future for other models as well. Currently, the implicit solver for Richards equation is implemented and in use. The full soil implicit solver is still in progress.

Components

• A way to specify separate explicit and implicit tendencies
  • Richards equation: imp, exp tendency functions instead of ode_function/rhs #191
• A way to set tolerances and other convergence requirements (max iters, etc)
  • Included by using ClimaTimeSteppers: https://github.com/CliMA/ClimaLand.jl/pull/218/files
• A way to specify Jacobians, and how often to update them; Function to compute Richards model Jacobian (derivation here)
  • Richards equation: Jacobian update defaults #190

Inputs

• For Richards equation: https://www.overleaf.com/9252289794ykktrwjxqvcj. We will be implementing modified Picard iterations.
• For full soil equation: TODO need to write out jacobian derivation

Results and deliverables/Validation

• For Richards equation: plots of mass conservation and solution error vs timestep for modified Picard iterations, for constant flux and state BC test runs. Comparison to an explicit solver of the same order - obtain max dt the explicit solver runs at. - in our buildkite
• Richards equation: unit test of Jacobian computed by the code with the analytic expression
• Make sure the above validation results are preserved with new CTS interface
• For full soil model : plots of mass conservation and solution error vs timestep for modified Picard iterations, for constant flux, moisture state BC, and for AtmosDrivenBC. Comparison to an explicit solver of the same order - obtain max dt the explicit solver runs at.
• Unit tests of jacobian
• flexible implementation that allows for LSM models - with many component models - to also have Jacobians specified. Start with soil/canopy model

Initial task breakdown

A preliminary list of PRs and a preliminary timeline of PRs, milestones, and key results.

• Have reviewers review equations (Tapio)
• Add method computing Jacobian for approximate Newton (modified Picard). Make sure A and b update every X timesteps, where X is chosen at the top-most level.
• Simulations with the implicit solver using ClimaTimesteppers. comparison of approx Newton to explicit runs
• State varying boundary conditions treated implicitly
• Implement in general fashion Jacobian update defaults #190 imp, exp tendency functions instead of ode_function/rhs #191 Use ClimaTimeSteppers, add implicit helper functions #218
• update to new ClimaTimesteppers interface use new interface for implicit RRE #528
• Add Jacobian for full soil model; repeat tests step full soil model implicitly #538
• Add Jacobian for an LSM model; repeat tests step full soil model implicitly #538
• Time dependent BC TBD

Producers

@juliasloan25 (design and implementation) @kmdeck @dennisYatunin (design, implementation support)

Reviewers

Katherine, Dennis - PRs
Tapio - SDI review (@tapios)

SDI Changelog

• 4 March 2024: SDI updated for clarity and to more accurately reflect current status (by @juliasloan25)

step RichardsModel implicitly

Beta Give feedback

1. Step RRE implicitly #224
   enhancement +1
   
2. Remove general implicit solver functions #268
   enhancement +1
   
3. use new interface for implicit RRE #528
   enhancement +1
   
4. test GPU compatibility after interface change #541
   GPU enhancement +2
   
5. Add global GPU RichardsModel to benchmarks pipeline #647
   enhancement +1
   
6. smooth derivative dPsi/dtheta at theta = \nu (if needed)

step EnergyHydrology implicitly

Beta Give feedback

1. step full soil model implicitly #538
   enhancement +1
   
2. Reduce code duplication in jacobian functions #659
   enhancement +1
3. Unify jacobian/tendency naming convention #660
   enhancement +1
   
4. Generalize variable names in Jacobian initialization #661
   enhancement +1
5. Update ozark conservation test for IMEX stepping #662
   enhancement +1
6. add implicit tendency for EnergyHydrology energy #663
   enhancement +1
   
7. include boundary condition contributions (if needed)
8. include contribution from drhoe_int / dtheta_l (if needed)
9. Understand how to treat soil phase change implicitly

QA

Beta Give feedback

1. write tutorial demonstrating how to use IMEX model

[tapios]

LGTM.

We should keep the tridiagonal solver in one central place that we do not have to duplicate it (so that we update it in one place, e.g., for GPU performance, this percolates through).

[kmdeck]

@dennisYatunin could we move the thomas algorithm function to ClimaTimesteppers?

also, if the matrix passed is the identity, will the algorithm spend time inverting it, or will it just return b?

[dennisYatunin]

Actually, ClimaCore might be a slightly better location than ClimaTimeSteppers. We're currently computing Jacobian blocks as ClimaCore Fields, then copying them into arrays to run the Thomas Algorithm. It would probably be better to have a function in ClimaCore that can directly invert a Jacobian block Field (I left a note in ClimaCore last year to add a function like this).

If the Jacobian matrix is I (or some constant times I, which is encoded by the same struct), the matrix inversion will be ldiv!(x, I, b), which is defined in LinearAlgebra to just set x .= b (scaled by a constant, if applicable).

[kmdeck]

If the Jacobian matrix is I (or some constant times I, which is encoded by the same struct), the matrix inversion will be ldiv!(x, I, b), which is defined in LinearAlgebra to just set x .= b (scaled by a constant, if applicable).

excellent!