Merge #190

190: Jacobian update defaults r=kmdeck a=kmdeck

## Purpose 
Part of SDI #135 
Adds functions which compute the Jacobian and Jacobian due to the boundary terms in the tendency.
Adds default functions + extends those for Richards equation.

## Content
- adds three functions: `make_update_jacobian`, `make_tendency_jacobian`, and `\partialtendencyBC\partialY`. 
- adds the methods of these for RichardsModel.
- adds unit tests which demonstrate that the Jacobian is computed correctly.

Review checklist

I have:
- followed the codebase contribution guide: https://clima.github.io/ClimateMachine.jl/latest/Contributing/
- followed the style guide: https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/
- followed the documentation policy: https://github.com/CliMA/policies/wiki/Documentation-Policy
- checked that this PR does not duplicate an open PR.

In the Content, I have included 
- relevant unit tests, and integration tests, 
- appropriate docstrings on all functions, structs, and modules, and included relevant documentation.

----
- [x] I have read and checked the items on the review checklist.


Co-authored-by: kmdeck <kdeck@caltech.edu>

Project.toml

@@ -14,6 +14,7 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"

docs/src/APIs/SharedUtilities.md

@@ -48,6 +48,10 @@ ClimaLSM.BottomBoundary
 ClimaLSM.boundary_flux
 ClimaLSM.diffusive_flux
 ClimaLSM.get_Δz
+ClimaLSM.make_tendency_jacobian
+ClimaLSM.make_update_jacobian
+ClimaLSM.∂tendencyBC∂Y
+ClimaLSM.AbstractTridiagonalW
 ```
 
 ## Drivers

docs/src/APIs/Soil.md

@@ -36,6 +36,7 @@ ClimaLSM.Soil.temperature_from_ρe_int
 ClimaLSM.Soil.thermal_conductivity
 ClimaLSM.Soil.phase_change_source
 ClimaLSM.Soil.thermal_time
+ClimaLSM.Soil.dψdϑ
 ```
 
 ## Soil Parameters
@@ -56,4 +57,10 @@ ClimaLSM.Soil.FreeDrainage
 ClimaLSM.Soil.AtmosDrivenFluxBC
 ClimaLSM.Soil.AbstractSoilSource
 ClimaLSM.Soil.PhaseChange
+```
+
+## Soil Jacobian Structures
+
+```@docs
+ClimaLSM.Soil.RichardsTridiagonalW
 ```

experiments/Manifest.toml

@@ -193,10 +193,10 @@ uuid = "d934ef94-cdd4-4710-83d6-720549644b70"
 version = "0.3.7"
 
 [[deps.ClimaLSM]]
-deps = ["ArtifactWrappers", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaCoreTempestRemap", "Dates", "DocStringExtensions", "Insolation", "IntervalSets", "JLD2", "NCDatasets", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnPack"]
+deps = ["ArtifactWrappers", "CLIMAParameters", "ClimaComms", "ClimaCore", "ClimaCoreTempestRemap", "Dates", "DocStringExtensions", "Insolation", "IntervalSets", "JLD2", "LinearAlgebra", "NCDatasets", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnPack"]
 path = ".."
 uuid = "7884a58f-fab6-4fd0-82bb-ecfedb2d8430"
-version = "0.2.2"
+version = "0.2.3"
 
 [[deps.CloseOpenIntervals]]
 deps = ["ArrayInterface", "Static"]

src/ClimaLSM.jl

@@ -16,6 +16,7 @@ include("SharedUtilities/drivers.jl")
 include("SharedUtilities/utils.jl")
 include("SharedUtilities/boundary_conditions.jl")
 include("SharedUtilities/sources.jl")
+include("SharedUtilities/implicit_tendencies.jl")
 include("Bucket/Bucket.jl")
 export make_interactions_update_aux
 export domain

src/SharedUtilities/boundary_conditions.jl

@@ -17,13 +17,13 @@ convenient conditions.
 abstract type AbstractBC end
 
 """
-    AbstractBoundary{}
+    AbstractBoundary
 
 An abstract type to indicate which boundary we are doing calculations for.
 Currently, we support the top boundary (TopBoundary)
 and bottom boundary (BottomBoundary).
 """
-abstract type AbstractBoundary{} end
+abstract type AbstractBoundary end
 
 
 """

src/SharedUtilities/implicit_tendencies.jl

@@ -0,0 +1,72 @@
+export make_tendency_jacobian,
+    make_update_jacobian, AbstractTridiagonalW, ∂tendencyBC∂Y
+
+
+"""
+   make_tendency_jacobian(model::AbstractModel)
+
+Creates and returns a function which updates the auxiliary
+variables `p` in place and then updates the entries of the 
+Jacobian matrix `W` for the `model` in place.
+
+The default is that no updates are required, no implicit tendency is
+present, and hence the timestepping is entirely explicit.
+
+Note that the returned function `tendency_jacobian!` should be
+used as `Wfact!` in `ClimaTimeSteppers.jl` and `OrdinaryDiffEq.jl`.
+"""
+function make_tendency_jacobian(model::AbstractModel)
+    update_aux! = make_update_aux(model)
+    update_jacobian! = make_update_jacobian(model)
+    function tendency_jacobian!(W, Y, p, dtγ, t)
+        update_aux!(p, Y, t)
+        update_jacobian!(W, Y, p, dtγ, t)
+    end
+    return tendency_jacobian!
+end
+
+"""
+    make_update_jacobian(model::AbstractModel)
+
+Creates and returns a function which updates the entries
+of the Jacobian matrix `W` in place.
+
+If the implicit tendency function is given by 
+`T!(dY, Y, p, t) = make_implicit_tendency(model)`, the Jacobian 
+should be given by `W_{i,j}! = ∂T!_i/∂Y_j`, where `Y_j` is the 
+`j-th` state variable
+and `T!_i` is the implicit tendency of the `i-th` state variable.
+
+The default is that no updates are required, no implicit tendency is
+present, and hence the timestepping is entirely explicit.
+"""
+function make_update_jacobian(model::AbstractModel)
+    function update_jacobian!(W, Y, p, dtγ, t) end
+    return update_jacobian
+end
+
+"""
+    ∂tendencyBC∂Y(::AbstractModel,
+                  ::AbstractBC,
+                  ::AbstractBoundary,
+                  _...)::Union{ClimaCore.Fields.FieldVector, Nothing}
+
+A function stub which returns the derivative of the implicit tendency
+term of the `model` arising from the boundary condition,
+with respect to the state Y.
+"""
+function ∂tendencyBC∂Y(
+    ::AbstractModel,
+    ::AbstractBC,
+    ::AbstractBoundary,
+    _...,
+)::Union{ClimaCore.Fields.FieldVector, Nothing} end
+
+"""
+    AbstractTridiagonalW
+
+An abstract type for tridiagonal Jacobian matrices.
+"""
+abstract type AbstractTridiagonalW end
+
+Base.similar(w::AbstractTridiagonalW) = w

src/SharedUtilities/utils.jl

@@ -10,3 +10,10 @@ function heaviside(x::FT)::FT where {FT}
         return FT(0.0)
     end
 end
+
+"""
+    to_scalar_coefs(vector_coefs)
+
+Helper function used in computing tendencies of vertical diffusion terms.
+"""
+to_scalar_coefs(vector_coefs) = map(vector_coef -> vector_coef.u₃, vector_coefs)

src/Soil/Soil.jl

@@ -56,18 +56,20 @@ for interactions between components.
 =#
 
 using ClimaLSM
-using UnPack
 using DocStringExtensions
+using LinearAlgebra
 using ClimaCore
 import ..Parameters as LSMP
 import ClimaCore: Fields, Operators, Geometry, Spaces
 using Thermodynamics
 
 import ClimaLSM.Domains: Column, HybridBox, SphericalShell
+using ClimaLSM: AbstractTridiagonalW
 import ClimaLSM:
     AbstractModel,
     make_update_aux,
     make_rhs,
+    make_update_jacobian,
     prognostic_vars,
     auxiliary_vars,
     name,
@@ -84,6 +86,7 @@ import ClimaLSM:
     surface_height
 export RichardsModel,
     RichardsParameters,
+    RichardsTridiagonalW,
     EnergyHydrology,
     EnergyHydrologyParameters,
     AbstractSoilModel,

src/Soil/boundary_conditions.jl

@@ -159,7 +159,7 @@ function ClimaLSM.boundary_flux(
     ψ_c = Fields.level(p.soil.ψ, p_len)
 
     # Calculate pressure head using boundary condition
-    @unpack vg_α, vg_n, vg_m, θ_r, ν, S_s = params
+    (; vg_α, vg_n, vg_m, θ_r, ν, S_s) = params
     θ_bc = rre_bc.bc(p, t)
     ψ_bc = @. pressure_head(vg_α, vg_n, vg_m, θ_r, θ_bc, ν, S_s)
 
@@ -186,7 +186,7 @@ function ClimaLSM.boundary_flux(
     ψ_c = Fields.level(p.soil.ψ, 1)
 
     # Calculate pressure head using boundary condition
-    @unpack vg_α, vg_n, vg_m, θ_r, ν, S_s = params
+    (; vg_α, vg_n, vg_m, θ_r, ν, S_s) = params
     θ_bc = rre_bc.bc(p, t)
     ψ_bc = @. pressure_head(vg_α, vg_n, vg_m, θ_r, θ_bc, ν, S_s)
 
@@ -270,3 +270,85 @@ function soil_boundary_fluxes(bc::NamedTuple, boundary, model, Y, Δz, p, t)
     return ClimaLSM.boundary_flux(bc.water, boundary, Δz, p, t, params),
     ClimaLSM.boundary_flux(bc.heat, boundary, Δz, p, t, params)
 end
+
+"""
+    ClimaLSM.∂tendencyBC∂Y(
+        model::RichardsModel,
+        ::MoistureStateBC,
+        boundary::ClimaLSM.TopBoundary,
+        Δz,
+        Y,
+        p,
+        t,
+)
+
+Computes and returns the derivative of the part of the
+implicit tendency in the top layer, due to the boundary
+condition, with respect to the state variable in the top layer.
+
+For a diffusion equation like Richards equation with a single state
+variable, this is given by
+`∂T_N∂Y_N = [-∂/∂z(∂F_bc/∂Y_N)]_N`, where `N` indicates the top 
+layer cell index.
+"""
+function ClimaLSM.∂tendencyBC∂Y(
+    model::RichardsModel,
+    ::MoistureStateBC,
+    boundary::ClimaLSM.TopBoundary,
+    Δz,
+    Y,
+    p,
+    t,
+)
+    (; ν, vg_α, vg_n, vg_m, S_s, θ_r) = model.parameters
+    fs = ClimaLSM.Domains.obtain_face_space(model.domain.space)
+    face_len = ClimaCore.Utilities.PlusHalf(ClimaCore.Spaces.nlevels(fs) - 1)
+    interpc2f_op = Operators.InterpolateC2F(
+        bottom = Operators.Extrapolate(),
+        top = Operators.Extrapolate(),
+    )
+    K = Fields.level(interpc2f_op.(p.soil.K), face_len)
+    ϑ_l = Fields.level(interpc2f_op.(Y.soil.ϑ_l), face_len)
+    return ClimaCore.Fields.FieldVector(;
+        :soil => (;
+            :ϑ_l => @. -K / Δz * dψdϑ(ϑ_l, ν, θ_r, vg_α, vg_n, vg_m, S_s) /
+               (2 * Δz)
+        ),
+    )
+end
+
+"""
+    ClimaLSM.∂tendencyBC∂Y(
+        ::AbstractSoilModel,
+        ::AbstractSoilBC,
+        boundary::ClimaLSM.TopBoundary,
+        Δz,
+        Y,
+        p,
+        t,
+)
+
+A default method which computes and returns the zero for the 
+derivative of the part of the
+implicit tendency in the top layer, due to the boundary
+condition, with respect to the state variable in the top layer.
+
+For a diffusion equation like Richards equation with a single state
+variable, this is given by
+`∂T_N∂Y_N = [-∂/∂z(∂F_bc/∂Y_N)]_N`, where `N` indicates the top 
+layer cell index.
+
+If `F_bc` can be approximated as independent of `Y_N`, the derivative 
+is zero.
+"""
+function ClimaLSM.∂tendencyBC∂Y(
+    ::AbstractSoilModel,
+    ::AbstractSoilBC,
+    boundary::ClimaLSM.TopBoundary,
+    Δz,
+    Y,
+    p,
+    t,
+)
+    return ClimaCore.Fields.FieldVector(; :soil => (; :ϑ_l => zeros(axes(Δz))))
+end

src/Soil/energy_hydrology.jl

@@ -318,18 +318,20 @@ This has been written so as to work with Differential Equations.jl.
 """
 function ClimaLSM.make_update_aux(model::EnergyHydrology)
     function update_aux!(p, Y, t)
-        @unpack ν,
-        vg_α,
-        vg_n,
-        vg_m,
-        K_sat,
-        S_s,
-        θ_r,
-        Ω,
-        γ,
-        γT_ref,
-        κ_sat_frozen,
-        κ_sat_unfrozen = model.parameters
+        (;
+            ν,
+            vg_α,
+            vg_n,
+            vg_m,
+            K_sat,
+            S_s,
+            θ_r,
+            Ω,
+            γ,
+            γT_ref,
+            κ_sat_frozen,
+            κ_sat_unfrozen,
+        ) = model.parameters
 
         p.soil.θ_l = volumetric_liquid_fraction.(Y.soil.ϑ_l, ν .- Y.soil.θ_i)
         p.soil.κ =
@@ -406,7 +408,7 @@ function ClimaLSM.source!(
     model,
 ) where {FT}
     params = model.parameters
-    @unpack ν, earth_param_set = params
+    (; ν, earth_param_set) = params
     _ρ_l = FT(LSMP.ρ_cloud_liq(earth_param_set))
     _ρ_i = FT(LSMP.ρ_cloud_ice(earth_param_set))
     ρc = volumetric_heat_capacity.(p.soil.θ_l, Y.soil.θ_i, Ref(params))