Implement implicit timestepping for soil energy and hydrology

[bgroenks96]

WIP: Currently just the initial draft of the plan proposed by Claude Fable.

CC @olivierbonte and @maximilian-gelbrecht: would be good to start a discussion already before actually implementing anything.

[olivierbonte]

In general, I think my preference would be an approach somewhere between the two that are proposed now in terms of numerical correctness:

• Approach 1 (JFNK) is very ambitious I think and I am doubtful if it will work at global scale. See my comments for background info
• Approach 2: Going for a 1 sweep tridiagonal solve with or picard iteration can be valuable, but it is quite a specific methodology that already does operator splitting between hydrology and energy

I think the intermediate route is closer to what is done in the ClimaLand paper (https://doi.org/10.1029/2025MS005118, Appendix F):

• Implicit Euler (for some variables) where the nonlinear system is solved by Newonts method
• No global nonlinear solve, but 1 per spatial column
• To get the Jacobian of the implicit tendencies, I would propose 2 systems:
  1. Analytical jacobians for robustness and speed
  2. AD calculated Jacobians for generality. Ideally, you would leverage sparsity here to reduce the amount of JVPs needed, but I think for an initial implementation I would go for dense.
• At each Newton iteration, a linear system has to be solved. The system is sparse (because of the sparse Jacobian, see F9 climaland paper). ClimaCore implements many of the typical sparse matrices (https://clima.github.io/ClimaCore.jl/dev/reference/matrix_fields/#ClimaCore.MatrixFields.BandMatrixRow) and related solvers (https://clima.github.io/ClimaCore.jl/dev/reference/matrix_fields/#Linear-Solvers). Here, we would need to think of what the equivalance in Oceananigans is. I am afraid that a simple A \ b for each column won't work inside of a kernel : /

Edit: the only nonlinear solve inside of a GPU kernel I am aware of is using StaticArrays.jl and combintation with (Simple)NonLinearSolve.jl: https://docs.sciml.ai/NonlinearSolve/stable/tutorials/nonlinear_solve_gpus/#GPU-Acceleration-over-Large-Parameter-Searches-using-KernelAbstractions.jl

[olivierbonte]

All prognostic variables in the model can be seen as elements of the full state vector, which is then the variable (see appendix F of ClimaLand paper). So I don't think this criticism is really valid.

[bgroenks96]

I think you might be misunderstanding the criticism. The point is that IMEX is typically applied per term on the RHS: you split $u' = f(u)$ into $u' = f_1(u) + f_2(u)$ where $f_1$ is the implicit part and $f_2$ is the explicit part. Thus, for a single variable, a typical IMEX scheme would compute two tendencies, one explicitly and one implicitly. The current IMEX solver assumes that variables $u_i$ can be grouped such that they only have either $f_1$ or $f_2$.

[olivierbonte]

No I understand, I just wanted to point out that if you treat all variables as one big $\mathbf{u}$ vector, then its vector RHS is split. I suppose it would be interesting if you could split out per variable even what is treated explicilty and what implicitly, but does seems like quite a big change.

[bgroenks96]

I just wanted to point out that if you treat all variables as one big u vector, then its vector RHS is split.

Yes but its split with a very specific sparsity structure. Putting variables into implicit or explicit groups implies that one of $f_1$ or $f_2$ is always zero.

[bgroenks96]

so in other words, the supervector structure does not make it fully general.

[olivierbonte]

with this linearisation path, the goal is to get the implicit step in 1 tridiagonal sweep right? But as we discussed @bgroenks96 this requires quite some simplifications.

A more general nonlinear solve to the implicit statevariables based on the residual (the Path 1) makes more sense to me

[olivierbonte]

I would have to check the Tubini paper (https://doi.org/10.5194/tc-15-2541-2021) in more detail again, but if I remember correctly they have a part on the limitations of the apparent heat capacity approach and why it is more consistent to solve in the enthalpy/internal energy form.

[bgroenks96]

We already solve the enthalpy form and that would not change under this approach. But you still solve for temperature as the unknown, also in Tubini's method.

[olivierbonte]

I see, I just checked Tubini's paper again briefly, and indeed they are solving for temperature. They do use a (modified) Newton algorithm though (the NCZ of their paper).

[olivierbonte]

Rolling out custom numerical linear algebra for Terrarium only does not seem ideal to me. Preferably via Oceananigans

[bgroenks96]

I am not sure whether or not Krylov.jl will be GPU and AD compatible though, might be a huge can of worms.

[olivierbonte]

This would indeed be the ideal case, but I am doubtful of this will allow to include implicit state variables beyond those in the soil. Parfllow e.g. does this for 3D modelling of soil.
Again referring to the climaland paper, there they also implicitly timestep the canopy leaf temperature. Or the canopy air space in the formulation of #132 (comment)

[olivierbonte]

The only source I can find on tyring to do full coupling of processes + all spatial cells at once is the Amanzi Advanced Terrestrial Simulator code: https://amanzi.github.io/ats/stable/index.html.

illustrates that potentially all processes can be strongly coupled in their framework (https://doi.org/10.1002/2015WR018427). https://doi.org/10.1016/j.envsoft.2015.12.017 describes the framework.

They also seem to apply operator splitting in https://doi.org/10.1007/s10596-017-9679-3, where they argue for the split in:

1. later transport of water and energy
2. per colunm solving of surface energy balance + subsurface hydrology and energy. In this way, if one land column hits a difficult solve due to e.g. freezing-thawing, this doens't affect the full global solve.

In this paper the sey that they couldn't run a 468 spatial cell domain in full 3D mode computationally, so makes you wonder how hard it would be for global.

[olivierbonte]

Note that this is the approach implemented in https://numericalmathematics.github.io/Ariadne.jl/stable/#Ariadne.newton_krylov!. Ardiane.jl is only for normal arrays though :/

[bgroenks96]

No global nonlinear solve, but 1 per spatial column

We can consider adding an alternative scheme later that does this for CPU-based models with a relatively small number of columns, but this is not going to be the main route forward. It is fundamentally at odds with the way our timestepping scheme is structured and is not GPU friendly. We definitely do not want (by default) a scheme where different grid cells run different numbers of iterations.

When faced with the trade-off of numerical exactness vs. scalable performance, we have to favor the latter by default.

[olivierbonte]

Apparently, you can actually use A \ b notation to solve wihtin a kernel but you have to use StaticArrays.jl and you are limited to matrices A of 14 x 14 or smaller. The example script below:

using KernelAbstractions, StaticArrays, LinearAlgebra
using CUDA   # only needed so that CUDABackend() exists; swap in AMDGPU/Metal/oneAPI as needed

@kernel function solve_k!(x, @Const(A), @Const(b))
    i = @index(Global)
    @inbounds x[i] = A[i] \ b[i]
end

to_device(backend, x) = copyto!(KernelAbstractions.allocate(backend, eltype(x), size(x)), x)

n_parallel_solves = 1024
workgroup_size = 128
for N in (2, 3, 4, 8, 14, 15, 25), backend in (CPU(), CUDABackend())
    A = [SMatrix{N, N}(rand(N, N) + N * I) for _ in 1:n_parallel_solves]
    b = [SVector{N}(rand(N)) for _ in 1:n_parallel_solves]
    dA, db = to_device(backend, A), to_device(backend, b)
    dx = similar(db)
    try
        solve_k!(get_backend(dx), workgroup_size)(dx, dA, db; ndrange = length(b))
        println((N, backend, maximum(norm.(Array(dx) .- A .\ b))))
    catch e
        println((N, backend, "FAILED: ", typeof(e)))
    end
end

fails on GPU from N = 15 onwards.
I don't think this is something that can be solved easily. In the DiffEqGPU.jl paper (https://doi.org/10.1016/j.cma.2023.116591) for example it also mentioned as a limitation in the discussion:
For example, while using stack-allocated arrays provides a workaround to using arrays inside GPU kernels, they are not suitable for higher-dimensional problems due to the limited memory of static allocations. The model might compile, but there might not be any realizable speedups.

[olivierbonte]

No global nonlinear solve, but 1 per spatial column

We can consider adding an alternative scheme later that does this for CPU-based models with a relatively small number of columns, but this is not going to be the main route forward. It is fundamentally at odds with the way our timestepping scheme is structured and is not GPU friendly. We definitely do not want (by default) a scheme where different grid cells run different numbers of iterations.

For global simulations, you could set your newton iterations fixed no? So then you can still solve per column (which is an easier problem than the global solve) without having the problem that the GPU is idle while one cell is finishing. ClimaLand also does it this way, see
https://github.com/CliMA/ClimaLand.jl/blob/5dc19d08cfc6bec65c2b6746ef598263f74298ee/src/simulations/Simulations.jl#L127-L135

For a small number of columns, you could then choose to use the same algorithm, but this time prescribe a convergence criterion.

[maximilian-gelbrecht]

I'm not so negative on nested AD. It might still work, it would at least try it

[maximilian-gelbrecht]

We had this discussion already in one of the Reactant PRs, in principle this is traceable. It might not be ideal from a performance point of view to have a non-fixed-iteration solver with Reactant at the moment, but it should be possible

[bgroenks96]

I would say from my attempts to make it work on RootSolvers.jl that it's only theoretically possible. It was fraught with problems. Especially unclear is how (or whether) Reactant can handle cases where you have a nested dynamic condition within a @traced condition. @trace doesn't nest, AFAIK.

[maximilian-gelbrecht]

@trace just don't nest within a single function scope

function foo(x)
     @trace while x < 10 
                 bar(x) 
      end
end 

function bar(x)
      @trace for i=1:10
                  x += 0.1  
       end 
end 

should work. Afaik there was a PR open some time ago that all make nesting @trace easier.

[maximilian-gelbrecht]

In Speedy I also already had nested @traces like this (that I got rid of later again because there weren't necessary)

[maximilian-gelbrecht]

Could we also have an LLM written analytic Jacobian as a fallback for common options?

[bgroenks96]

I would rather not because it introduces a maintenance burden.

[bgroenks96]

For global simulations, you could set your newton iterations fixed no? So then you can still solve per column (which is an easier problem than the global solve) without having the problem that the GPU is idle while one cell is finishing.

The problem is that this requires all of the model logic, i.e. compute_auxiliary! and compute_tendencies! to run within a single kernel, and effectively commits us to exclusively column-based modeling per the timestepping design. That's a diversion that I don't think we should follow.