off diagonal term in jacobian (#952)

Author: juliasloan25

.buildkite/pipeline.yml

@@ -96,6 +96,9 @@ steps:
       - label: "Soilbiogeochem"
         command: "julia --color=yes --project=.buildkite experiments/standalone/Biogeochemistry/experiment.jl"
 
+      - label: "Implicit stepper full soil; saturated soil"
+        command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/energy_hydrology_saturated.jl"
+
       - label: "Water conservation"
         command: "julia --color=yes --project=.buildkite experiments/standalone/Soil/water_conservation.jl"
         artifact_paths: "experiments/standalone/Soil/water_conservation/output_active/water_conservation*png"

docs/tutorials/standalone/Soil/freezing_front.jl

@@ -221,7 +221,7 @@ ode_algo = CTS.IMEXAlgorithm(
     timestepper,
     CTS.NewtonsMethod(
         max_iters = 3,
-        update_j = CTS.UpdateEvery(CTS.NewTimeStep),
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
     ),
 );
 

experiments/integrated/fluxnet/US-Var/US-Var_simulation.jl

@@ -20,4 +20,4 @@ h_stem = FT(0) # m
 t0 = Float64(0)
 
 # Time step size:
-dt = Float64(900)
+dt = Float64(1800)

experiments/long_runs/land.jl

@@ -380,7 +380,7 @@ end
 function setup_and_solve_problem(; greet = false)
 
     t0 = 0.0
-    tf = 60 * 60.0 * 24 * 365
+    tf = 60 * 60.0 * 24 * 365 * 2
     Δt = 450.0
     nelements = (101, 15)
     if greet
@@ -413,7 +413,12 @@ short_names = ["gpp", "swc", "et", "ct"]
 mktempdir(root_path) do tmpdir
     for short_name in short_names
         var = get(simdir; short_name)
-        times = [ClimaAnalysis.times(var)[1], ClimaAnalysis.times(var)[end]]
+        N = length(ClimaAnalysis.times(var))
+        times = [
+            ClimaAnalysis.times(var)[1],
+            ClimaAnalysis.times(var)[div(N, 2, RoundNearest)],
+            ClimaAnalysis.times(var)[N],
+        ]
         for t in times
             fig = CairoMakie.Figure(size = (600, 400))
             kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()

experiments/long_runs/snowy_land.jl

@@ -482,7 +482,12 @@ short_names = ["gpp", "swc", "et", "ct", "swe", "si"]
 mktempdir(root_path) do tmpdir
     for short_name in short_names
         var = get(simdir; short_name)
-        times = [ClimaAnalysis.times(var)[end]]
+        N = length(ClimaAnalysis.times(var))
+        times = [
+            ClimaAnalysis.times(var)[1],
+            ClimaAnalysis.times(var)[div(N, 2, RoundNearest)],
+            ClimaAnalysis.times(var)[N],
+        ]
         for t in times
             fig = CairoMakie.Figure(size = (600, 400))
             kwargs = ClimaAnalysis.has_altitude(var) ? Dict(:z => 1) : Dict()
@@ -500,13 +505,7 @@ mktempdir(root_path) do tmpdir
     end
     figures = readdir(tmpdir, join = true)
     pdfunite() do unite
-        run(
-            Cmd([
-                unite,
-                figures...,
-                joinpath(root_path, "last_year_figures.pdf"),
-            ]),
-        )
+        run(Cmd([unite, figures..., joinpath(root_path, "figures.pdf")]))
     end
 end
 

experiments/long_runs/soil.jl

@@ -221,7 +221,7 @@ end
 function setup_and_solve_problem(; greet = false)
 
     t0 = 0.0
-    tf = 60 * 60.0 * 24 * 365
+    tf = 60 * 60.0 * 24 * 365 * 2
     Δt = 450.0
     nelements = (101, 15)
     if greet
@@ -254,7 +254,12 @@ short_names = ["swc", "si", "sie"]
 mktempdir(root_path) do tmpdir
     for short_name in short_names
         var = get(simdir; short_name)
-        times = [ClimaAnalysis.times(var)[1], ClimaAnalysis.times(var)[end]]
+        N = length(ClimaAnalysis.times(var))
+        times = [
+            ClimaAnalysis.times(var)[1],
+            ClimaAnalysis.times(var)[div(N, 2, RoundNearest)],
+            ClimaAnalysis.times(var)[N],
+        ]
         for t in times
             var = get(simdir; short_name)
             fig = CairoMakie.Figure(size = (600, 400))

experiments/standalone/Soil/energy_hydrology_saturated.jl

@@ -0,0 +1,133 @@
+import SciMLBase
+using Statistics
+import ClimaTimeSteppers as CTS
+using ClimaCore
+import ClimaParams as CP
+using ClimaLand
+using ClimaLand.Domains: Column
+using ClimaLand.Soil
+import ClimaLand
+import ClimaLand.Parameters as LP
+
+FT = Float32;
+earth_param_set = LP.LandParameters(FT)
+
+vg_n = FT(2.9)
+vg_α = FT(6)
+K_sat = FT(4.42 / 3600 / 100) # m/s
+hcm = vanGenuchten{FT}(; α = vg_α, n = vg_n)
+ν = FT(0.43)
+θ_r = FT(0.045)
+S_s = FT(1e-3)
+ν_ss_om = FT(0.0)
+ν_ss_quartz = FT(1.0)
+ν_ss_gravel = FT(0.0)
+
+params = ClimaLand.Soil.EnergyHydrologyParameters(
+    FT;
+    ν,
+    ν_ss_om,
+    ν_ss_quartz,
+    ν_ss_gravel,
+    hydrology_cm = hcm,
+    K_sat,
+    S_s,
+    θ_r,
+    earth_param_set,
+);
+
+surface_water_flux = WaterFluxBC((p, t) -> -K_sat / 10)
+bottom_water_flux = WaterFluxBC((p, t) -> 0.0);
+
+surface_heat_flux = HeatFluxBC((p, t) -> 0.0)
+bottom_heat_flux = HeatFluxBC((p, t) -> 0.0);
+
+boundary_fluxes = (;
+    top = WaterHeatBC(; water = surface_water_flux, heat = surface_heat_flux),
+    bottom = WaterHeatBC(; water = bottom_water_flux, heat = bottom_heat_flux),
+);
+
+# Domain - single column
+zmax = FT(0)
+zmin = FT(-0.5)
+nelems = 25
+soil_domain = Column(; zlim = (zmin, zmax), nelements = nelems)
+z = ClimaCore.Fields.coordinate_field(soil_domain.space.subsurface).z;
+
+# Soil model, and create the prognostic vector Y and cache p:
+soil = Soil.EnergyHydrology{FT}(;
+    parameters = params,
+    domain = soil_domain,
+    boundary_conditions = boundary_fluxes,
+    sources = (),
+)
+Y, p, cds = initialize(soil);
+
+set_initial_cache! = make_set_initial_cache(soil)
+exp_tendency! = make_exp_tendency(soil)
+imp_tendency! = make_imp_tendency(soil);
+jacobian! = ClimaLand.make_jacobian(soil);
+jac_kwargs = (; jac_prototype = ImplicitEquationJacobian(Y), Wfact = jacobian!);
+
+function hydrostatic_equilibrium(z, z_interface, params)
+    (; ν, S_s, hydrology_cm) = params
+    (; α, n, m) = hydrology_cm
+    if z < z_interface
+        return -S_s * (z - z_interface) + ν
+    else
+        return ν * (1 + (α * (z - z_interface))^n)^(-m)
+    end
+end
+function init_soil!(Y, z, params)
+    FT = eltype(Y.soil.ϑ_l)
+    Y.soil.ϑ_l .= hydrostatic_equilibrium.(z, FT(-0.45), params)
+    Y.soil.θ_i .= 0
+    T = FT(296.15)
+    ρc_s = @. Soil.volumetric_heat_capacity(
+        Y.soil.ϑ_l,
+        FT(0),
+        params.ρc_ds,
+        params.earth_param_set,
+    )
+    Y.soil.ρe_int =
+        Soil.volumetric_internal_energy.(FT(0), ρc_s, T, params.earth_param_set)
+end
+
+# Timestepping:
+t0 = Float64(0)
+tf = Float64(24 * 3600)
+timestepper = CTS.ARS111();
+ode_algo = CTS.IMEXAlgorithm(
+    timestepper,
+    CTS.NewtonsMethod(
+        max_iters = 6,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+    ),
+);
+
+# Define the problem and callbacks:
+prob = SciMLBase.ODEProblem(
+    CTS.ClimaODEFunction(
+        T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+        dss! = ClimaLand.dss!,
+    ),
+    Y,
+    (t0, tf),
+    p,
+);
+# Solve at small dt
+init_soil!(Y, z, soil.parameters);
+set_initial_cache!(p, Y, t0);
+sol_dt_small = SciMLBase.solve(prob, ode_algo; dt = 100.0);
+init_soil!(Y, z, soil.parameters);
+set_initial_cache!(p, Y, t0);
+sol_dt_large = SciMLBase.solve(prob, ode_algo; dt = 900.0);
+@assert sum(isnan.(sol_dt_large.u[end])) == 0
+
+norm(x) = sqrt(mean(parent(x .^ 2)))
+rmse(x, y) = sqrt(mean(parent((x .- y) .^ 2)))
+@assert rmse(sol_dt_small[end].soil.ϑ_l, sol_dt_large[end].soil.ϑ_l) /
+        norm(sol_dt_small[end].soil.ϑ_l) < sqrt(eps(FT))
+@assert rmse(sol_dt_small[end].soil.ρe_int, sol_dt_large[end].soil.ρe_int) /
+        norm(sol_dt_small[end].soil.ρe_int) < sqrt(eps(FT))

src/shared_utilities/implicit_timestepping.jl

@@ -148,6 +148,21 @@ function ImplicitEquationJacobian(Y::ClimaCore.Fields.FieldVector)
                 get_j_field(axes(MatrixFields.get_field(Y, var)), FT),
         available_implicit_vars,
     )
+
+    # We include some terms ∂T_x/∂y where x ≠ y
+    # These are the off-diagonal terms in the Jacobian matrix
+    # Here, we take the convention that each pair has order (T_x, y) to produce ∂T_x/∂y as above
+    off_diagonal_pairs = ((@name(soil.ρe_int), @name(soil.ϑ_l)),)
+    available_off_diagonal_pairs = MatrixFields.unrolled_filter(
+        pair -> all(is_in_Y.(pair)),
+        off_diagonal_pairs,
+    )
+    implicit_off_diagonals = MatrixFields.unrolled_map(
+        pair ->
+            (pair[1], pair[2]) =>
+                get_j_field(axes(MatrixFields.get_field(Y, pair[1])), FT),
+        available_off_diagonal_pairs,
+    )
     # For explicitly-stepped variables, use the negative identity matrix
     # Note: We have to use FT(-1) * I instead of -I because inv(-1) == -1.0,
     # which means that multiplying inv(-1) by a Float32 will yield a Float64.
@@ -156,12 +171,21 @@ function ImplicitEquationJacobian(Y::ClimaCore.Fields.FieldVector)
         available_explicit_vars,
     )
 
+    matrix = MatrixFields.FieldMatrix(
+        implicit_blocks...,
+        implicit_off_diagonals...,
+        explicit_blocks...,
+    )
 
-
-    matrix = MatrixFields.FieldMatrix(implicit_blocks..., explicit_blocks...)
-
-    # Set up block diagonal solver for block Jacobian
-    alg = MatrixFields.BlockDiagonalSolve()
+    # Choose algorithm based on whether off-diagonal blocks are present
+    if is_in_Y(@name(soil.ρe_int)) && is_in_Y(@name(soil.ϑ_l))
+        # Set up lower triangular solver for block Jacobian with off-diagonal blocks
+        # Specify which variable to compute ∂T_x/∂x for first
+        alg = MatrixFields.BlockLowerTriangularSolve(@name(soil.ϑ_l))
+    else
+        # Set up block diagonal solver for block Jacobian with no off-diagonal blocks
+        alg = MatrixFields.BlockDiagonalSolve()
+    end
     solver = MatrixFields.FieldMatrixSolver(alg, matrix, Y)
 
     return ImplicitEquationJacobian(matrix, solver)

src/standalone/Soil/energy_hydrology.jl

@@ -322,6 +322,7 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
         # The derivative of the residual with respect to the prognostic variable
         ∂ϑres∂ϑ = matrix[@name(soil.ϑ_l), @name(soil.ϑ_l)]
         ∂ρeres∂ρe = matrix[@name(soil.ρe_int), @name(soil.ρe_int)]
+        ∂ρeres∂ϑ = matrix[@name(soil.ρe_int), @name(soil.ϑ_l)]
         # If the top BC is a `MoistureStateBC`, add the term from the top BC
         #  flux before applying divergence
         if haskey(p.soil, :dfluxBCdY)
@@ -373,6 +374,26 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
                     )
                 ) - (I,)
         end
+        @. ∂ρeres∂ϑ =
+            -dtγ * (
+                divf2c_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                    -interpc2f_op(
+                        volumetric_internal_energy_liq(
+                            p.soil.T,
+                            model.parameters.earth_param_set,
+                        ) * p.soil.K,
+                    ),
+                ) ⋅ gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                    ClimaLand.Soil.dψdϑ(
+                        hydrology_cm,
+                        Y.soil.ϑ_l,
+                        ν - Y.soil.θ_i, #ν_eff
+                        θ_r,
+                        S_s,
+                    ),
+                )
+            ) - (I,)
+
         @. ∂ρeres∂ρe =
             -dtγ * (
                 divf2c_matrix() ⋅

test/shared_utilities/implicit_timestepping/energy_hydrology_model.jl

@@ -145,5 +145,27 @@ for FT in (Float32, Float64)
         # Check the main diagonal, entry corresponding to top of column
         @test Array(parent(jac_ρe.entries.:2))[end] .≈
               dtγ * (-κ_ic / dz^2 * dTdρ_ic) - I
+
+        # Off diagonal blocks, ∂T_ρe_int/∂ϑ_l
+        jac_ρeϑ = jacobian.matrix[
+            MatrixFields.@name(soil.ρe_int),
+            MatrixFields.@name(soil.ϑ_l)
+        ]
+
+        ρe_liq_ic = ClimaLand.Soil.volumetric_internal_energy_liq(
+            T,
+            params.earth_param_set,
+        )
+        # Check the main diagonal, entry corresponding to bottom of column
+        @test Array(parent(jac_ρeϑ.entries.:2))[1] .≈
+              dtγ * (-ρe_liq_ic * K_ic / dz^2 * dψdϑ_ic) - I
+        # Check the main diagonal, entries corresponding to interior of domain
+        @test all(
+            Array(parent(jac_ρeϑ.entries.:2))[2:(end - 1)] .≈
+            dtγ * (-2 * ρe_liq_ic * K_ic / dz^2 * dψdϑ_ic) - I,
+        )
+        # Check the main diagonal, entry corresponding to top of column
+        @test Array(parent(jac_ρeϑ.entries.:2))[end] .≈
+              dtγ * (-ρe_liq_ic * K_ic / dz^2 * dψdϑ_ic) - I
     end
 end