update implicit solver interface

.buildkite/Manifest.toml

@@ -373,7 +373,7 @@ weakdeps = ["CUDA", "Krylov"]
     KrylovExt = "Krylov"
 
 [[deps.ClimaLand]]
-deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
+deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities"]
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.12.1"
@@ -2669,6 +2669,11 @@ git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b"
 uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 version = "0.1.5"
 
+[[deps.UnrolledUtilities]]
+git-tree-sha1 = "b73f7a7c25a2618c5052c80ed32b07e471cc6cb0"
+uuid = "0fe1646c-419e-43be-ac14-22321958931b"
+version = "0.1.2"
+
 [[deps.UnsafeAtomics]]
 git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
 uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"

Project.toml

@@ -21,6 +21,7 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SurfaceFluxes = "49b00bb7-8bd4-4f2b-b78c-51cd0450215f"
 Thermodynamics = "b60c26fb-14c3-4610-9d3e-2d17fe7ff00c"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
@@ -59,5 +60,6 @@ StaticArrays = "1"
 StatsBase = "0.33, 0.34"
 SurfaceFluxes = "0.10, 0.11"
 Thermodynamics = "0.12"
+UnrolledUtilities = "0.1"
 cuDNN = "1"
 julia = "1.9"

docs/Manifest.toml

@@ -370,7 +370,7 @@ weakdeps = ["CUDA", "Krylov"]
     KrylovExt = "Krylov"
 
 [[deps.ClimaLand]]
-deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
+deps = ["Adapt", "ArtifactWrappers", "ClimaComms", "ClimaCore", "ClimaUtilities", "DataFrames", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "IntervalSets", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities"]
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.12.1"
@@ -2692,6 +2692,11 @@ git-tree-sha1 = "6cc9d682755680e0f0be87c56392b7651efc2c7b"
 uuid = "9602ed7d-8fef-5bc8-8597-8f21381861e8"
 version = "0.1.5"
 
+[[deps.UnrolledUtilities]]
+git-tree-sha1 = "b73f7a7c25a2618c5052c80ed32b07e471cc6cb0"
+uuid = "0fe1646c-419e-43be-ac14-22321958931b"
+version = "0.1.2"
+
 [[deps.UnsafeAtomics]]
 git-tree-sha1 = "6331ac3440856ea1988316b46045303bef658278"
 uuid = "013be700-e6cd-48c3-b4a1-df204f14c38f"

docs/src/APIs/Soil.md

@@ -96,5 +96,5 @@ ClimaLand.Soil.RootExtraction
 ## Soil Jacobian Structures
 
 ```@docs
-ClimaLand.Soil.RichardsTridiagonalW
-```
+ClimaLand.Soil.ImplicitEquationJacobian
+```

docs/src/APIs/shared_utilities.md

@@ -61,8 +61,7 @@ ClimaLand.boundary_var_domain_names
 ClimaLand.boundary_var_types
 ClimaLand.make_tendency_jacobian
 ClimaLand.make_update_jacobian
-ClimaLand.∂tendencyBC∂Y
-ClimaLand.AbstractTridiagonalW
+ClimaLand.set_dfluxBCdY!
 ClimaLand.get_drivers
 ```
 

docs/tutorials/standalone/Soil/layered_soil.jl

@@ -29,7 +29,7 @@ t0 = Float64(0)
 tf = Float64(60 * 60)
 dt = Float64(30);
 
-# Define the domain 
+# Define the domain
 zmax = FT(0)
 zmin = FT(-1.1)
 nelems = 75
@@ -109,7 +109,7 @@ soil = Soil.RichardsModel{FT}(;
 Y, p, cds = initialize(soil);
 
 # Initial conditions
-Y.soil.ϑ_l .= 0.0353; # read from Figure 4 of Huang et al. 
+Y.soil.ϑ_l .= 0.0353; # read from Figure 4 of Huang et al.
 
 # We also set the initial conditions of the auxiliary state here:
 set_initial_cache! = make_set_initial_cache(soil)
@@ -131,9 +131,10 @@ ode_algo = CTS.IMEXAlgorithm(
 )
 exp_tendency! = make_exp_tendency(soil)
 imp_tendency! = make_imp_tendency(soil)
-update_jacobian! = make_update_jacobian(soil)
+tendency_jacobian! = make_tendency_jacobian(soil)
+
 jac_kwargs =
-    (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
+    (; jac_prototype = ImplicitEquationJacobian(Y), Wfact = tendency_jacobian!)
 prob = SciMLBase.ODEProblem(
     CTS.ClimaODEFunction(
         T_exp! = exp_tendency!,

docs/tutorials/standalone/Soil/richards_equation.jl

@@ -140,7 +140,7 @@ soil = Soil.RichardsModel{FT}(;
 # We also create the function which is used to update our Jacobian.
 exp_tendency! = make_exp_tendency(soil);
 imp_tendency! = ClimaLand.make_imp_tendency(soil);
-update_jacobian! = ClimaLand.make_update_jacobian(soil);
+tendency_jacobian! = ClimaLand.make_tendency_jacobian(soil);
 
 # # Set up the simulation
 # We can now initialize the prognostic and auxiliary variable vectors, and take
@@ -186,7 +186,7 @@ ode_algo = CTS.IMEXAlgorithm(
 
 # Here we set up the information used for our Jacobian.
 jac_kwargs =
-    (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!);
+    (; jac_prototype = ImplicitEquationJacobian(Y), Wfact = tendency_jacobian!);
 
 # And then we can solve the system of equations, using
 # [SciMLBase.jl](https://github.com/SciML/SciMLBase.jl) and

experiments/standalone/Soil/richards_comparison.jl

@@ -81,7 +81,7 @@ bonan_sand_dataset = ArtifactWrapper(
         Y.soil.ϑ_l .= FT(0.24)
         exp_tendency! = make_exp_tendency(soil)
         imp_tendency! = ClimaLand.make_imp_tendency(soil)
-        update_jacobian! = ClimaLand.make_update_jacobian(soil)
+        tendency_jacobian! = ClimaLand.make_tendency_jacobian(soil)
         set_initial_cache!(p, Y, t0)
 
         stepper = CTS.ARS111()
@@ -98,8 +98,8 @@ bonan_sand_dataset = ArtifactWrapper(
 
         # set up jacobian info
         jac_kwargs = (;
-            jac_prototype = RichardsTridiagonalW(Y),
-            Wfact = update_jacobian!,
+            jac_prototype = ImplicitEquationJacobian(Y),
+            Wfact = tendency_jacobian!,
         )
 
         prob = SciMLBase.ODEProblem(
@@ -181,7 +181,7 @@ end
         Y.soil.ϑ_l .= FT(0.1)
         exp_tendency! = make_exp_tendency(soil)
         imp_tendency! = ClimaLand.make_imp_tendency(soil)
-        update_jacobian! = ClimaLand.make_update_jacobian(soil)
+        tendency_jacobian! = ClimaLand.make_tendency_jacobian(soil)
         set_initial_cache!(p, Y, t0)
 
         stepper = CTS.ARS111()
@@ -197,8 +197,8 @@ end
         )
         # set up jacobian info
         jac_kwargs = (;
-            jac_prototype = RichardsTridiagonalW(Y),
-            Wfact = update_jacobian!,
+            jac_prototype = ImplicitEquationJacobian(Y),
+            Wfact = tendency_jacobian!,
         )
 
         prob = SciMLBase.ODEProblem(

experiments/standalone/Soil/richards_runoff.jl

@@ -218,7 +218,7 @@ Y.soil.ϑ_l .= hydrostatic_profile.(lat, z, ν, θ_r, vg_α, vg_n, S_s, f_max)
 set_initial_cache! = make_set_initial_cache(model)
 exp_tendency! = make_exp_tendency(model);
 imp_tendency! = ClimaLand.make_imp_tendency(model);
-update_jacobian! = ClimaLand.make_update_jacobian(model);
+tendency_jacobian! = ClimaLand.make_tendency_jacobian(model);
 
 set_initial_cache!(p, Y, t0)
 stepper = CTS.ARS111()
@@ -235,7 +235,7 @@ ode_algo = CTS.IMEXAlgorithm(
 
 # set up jacobian info
 jac_kwargs =
-    (; jac_prototype = RichardsTridiagonalW(Y), Wfact = update_jacobian!)
+    (; jac_prototype = ImplicitEquationJacobian(Y), Wfact = tendency_jacobian!)
 
 prob = SciMLBase.ODEProblem(
     CTS.ClimaODEFunction(

experiments/standalone/Soil/water_conservation.jl

@@ -102,7 +102,7 @@ for FT in (Float32, Float64)
     exp_tendency! = make_exp_tendency(soil)
     set_initial_cache! = make_set_initial_cache(soil)
     imp_tendency! = make_imp_tendency(soil)
-    update_jacobian! = make_update_jacobian(soil)
+    tendency_jacobian! = make_tendency_jacobian(soil)
 
     rmses = Array{FT}(undef, length(dts))
     mass_errors = Array{FT}(undef, length(dts))
@@ -114,8 +114,8 @@ for FT in (Float32, Float64)
         set_initial_cache!(p, Y, FT(0.0))
 
         jac_kwargs = (;
-            jac_prototype = RichardsTridiagonalW(Y),
-            Wfact = update_jacobian!,
+            jac_prototype = ImplicitEquationJacobian(Y),
+            Wfact = tendency_jacobian!,
         )
 
         prob = SciMLBase.ODEProblem(
@@ -143,6 +143,7 @@ for FT in (Float32, Float64)
             mass_change_exp = -(flux_in - flux_out) * t_sim
             mass_change_actual = mass_end - mass_start
             relerr = abs(mass_change_actual - mass_change_exp) / mass_change_exp
+            @show relerr
             @assert relerr < FT(1e-9)
             mass_errors[i] = relerr
 
@@ -212,7 +213,7 @@ for FT in (Float32, Float64)
     exp_tendency! = make_exp_tendency(soil_dirichlet)
     set_initial_cache! = make_set_initial_cache(soil_dirichlet)
     imp_tendency! = make_imp_tendency(soil_dirichlet)
-    update_jacobian! = make_update_jacobian(soil_dirichlet)
+    tendency_jacobian! = make_tendency_jacobian(soil_dirichlet)
     update_aux! = make_update_aux(soil_dirichlet)
 
     rmses_dirichlet = Array{FT}(undef, length(dts))
@@ -225,8 +226,8 @@ for FT in (Float32, Float64)
         set_initial_cache!(p, Y, FT(0.0))
 
         jac_kwargs = (;
-            jac_prototype = RichardsTridiagonalW(Y),
-            Wfact = update_jacobian!,
+            jac_prototype = ImplicitEquationJacobian(Y),
+            Wfact = tendency_jacobian!,
         )
 
         prob = SciMLBase.ODEProblem(

src/ClimaLand.jl

@@ -23,7 +23,6 @@ include("shared_utilities/drivers.jl")
 include("shared_utilities/boundary_conditions.jl")
 include("shared_utilities/sources.jl")
 include("shared_utilities/implicit_tendencies.jl")
-include("shared_utilities/implicit_functions.jl")
 include("standalone/Bucket/Bucket.jl")
 
 """

src/shared_utilities/Domains.jl

@@ -658,29 +658,28 @@ extruded center finite difference space, this would return a 2D field
 corresponding to the surface, with values equal to the topmost level.
 """
 function top_center_to_surface(center_field::ClimaCore.Fields.Field)
-    cs = axes(center_field)
-    face_space = obtain_face_space(cs)
-    N = ClimaCore.Spaces.nlevels(face_space)
-    interp_c2f = ClimaCore.Operators.InterpolateC2F(
-        top = ClimaCore.Operators.Extrapolate(),
-        bottom = ClimaCore.Operators.Extrapolate(),
-    )
-    surface_space = obtain_surface_space(cs)
+    center_space = axes(center_field)
+    N_minus_half = ClimaCore.Spaces.nlevels(center_space)
+    surface_space = obtain_surface_space(center_space)
     return ClimaCore.Fields.Field(
         ClimaCore.Fields.field_values(
-            ClimaCore.Fields.level(
-                interp_c2f.(center_field),
-                ClimaCore.Utilities.PlusHalf(N - 1),
-            ),
+            ClimaCore.Fields.level(center_field, N_minus_half),
         ),
         surface_space,
     )
 end
 
+"""
+    top_center_to_surface(val)
+
+When `val` is a scalar (e.g. a single float or struct), returns `val`.
+"""
+top_center_to_surface(val) = val
+
 """
     linear_interpolation_to_surface!(sfc_field, center_field, z)
 
-Linearly interpolate the center field `center_field` to the surface 
+Linearly interpolate the center field `center_field` to the surface
 defined by the top face coordinate of `z`; updates the `sfc_field`
 on the surface (face) space in place.
 """
@@ -705,7 +704,7 @@ end
     get_Δz(z::ClimaCore.Fields.Field)
 
 A function to return a tuple containing the distance between the top boundary
-and its closest center, and the bottom boundary and its closest center, 
+and its closest center, and the bottom boundary and its closest center,
 both as Fields.
 """
 function get_Δz(z::ClimaCore.Fields.Field)