handle MoistureStateBC [skip ci]

docs/src/APIs/shared_utilities.md

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

src/shared_utilities/implicit_tendencies.jl

@@ -1,5 +1,5 @@
 export make_tendency_jacobian,
-    make_update_jacobian, AbstractTridiagonalW, ∂tendencyBC∂Y
+    make_update_jacobian, AbstractTridiagonalW, set_dfluxBCdY!
 
 
 """
@@ -48,7 +48,7 @@ function make_update_jacobian(model::AbstractModel)
 end
 
 """
-    ∂tendencyBC∂Y(::AbstractModel,
+    set_dfluxBCdY!(::AbstractModel,
                   ::AbstractBC,
                   ::AbstractBoundary,
                   _...)::Union{ClimaCore.Fields.FieldVector, Nothing}
@@ -57,7 +57,7 @@ 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(
+function set_dfluxBCdY!(
     ::AbstractModel,
     ::AbstractBC,
     ::AbstractBoundary,

src/shared_utilities/models.jl

@@ -164,7 +164,7 @@ end
      make_update_cache(model::AbstractModel)
 
 A helper function which updates all cache variables of a model;
-currently only used in `set_initial_cache` since not all 
+currently only used in `set_initial_cache` since not all
 cache variables are updated at the same time.
 """
 function make_update_cache(model::AbstractModel)
@@ -404,6 +404,15 @@ function get_drivers(model::AbstractModel)
     return (nothing, nothing)
 end
 
+"""
+    add_dfluxBCdY_to_cache(p::NamedTuple, model::AbstractModel, _, _)
+
+We only need to add this field in the case of a `RichardsModel`, so we
+provide a fallback method for all other models which does nothing.
+"""
+function add_dfluxBCdY_to_cache(p::NamedTuple, model::AbstractModel, _, _)
+    return p
+end
 
 """
     initialize(model::AbstractModel)
@@ -418,5 +427,11 @@ function initialize(model::AbstractModel{FT}) where {FT}
     Y = initialize_prognostic(model, coords)
     p = initialize_auxiliary(model, coords)
     p = add_drivers_to_cache(p, model, coords)
+    p = add_dfluxBCdY_to_cache(
+        p,
+        model,
+        model.boundary_conditions.top,
+        ClimaLand.TopBoundary(),
+    )
     return Y, p, coords
 end

src/standalone/Soil/boundary_conditions.jl

@@ -161,7 +161,7 @@ end
                                               <:Runoff.TOPMODELRunoff,
                                               }, ::ClimaLand.TopBoundary)
 
-An extension of the `boundary_vars` method for RichardsAtmosDrivenFluxBC with 
+An extension of the `boundary_vars` method for RichardsAtmosDrivenFluxBC with
 TOPMODELRunoff.
 
 These variables are updated in place in `boundary_flux`.
@@ -181,7 +181,7 @@ boundary_vars(
                               ::ClimaLand.TopBoundary)
 
 An extension of the `boundary_var_domain_names` method for RichardsAtmosDrivenFluxBC
-with TOPMODELRunoff. 
+with TOPMODELRunoff.
 """
 boundary_var_domain_names(
     bc::RichardsAtmosDrivenFluxBC{
@@ -191,11 +191,12 @@ boundary_var_domain_names(
     ::ClimaLand.TopBoundary,
 ) = (:surface, :surface, :surface, :surface, :surface)
 """
-    boundary_var_types(::RichardsAtmosDrivenFluxBC{<:PrescribedPrecipitation,
+    boundary_var_types(::RichardsModel{FT},
+                        ::RichardsAtmosDrivenFluxBC{<:PrescribedPrecipitation,
                                                    <:Runoff.TOPMODELRunoff{FT},
                                                   },
-                       ::ClimaLand.TopBoundary,
-                       ) where {FT}
+                        ::ClimaLand.TopBoundary,
+                        ) where {FT}
 
 An extension of the `boundary_var_types` method for RichardsAtmosDrivenFluxBC
 with TOPMODELRunoff.
@@ -214,7 +215,7 @@ boundary_var_types(
                                               <:Runoff.AbstractRunoffModel,
                                               }, ::ClimaLand.TopBoundary)
 
-An extension of the `boundary_vars` method for RichardsAtmosDrivenFluxBC with 
+An extension of the `boundary_vars` method for RichardsAtmosDrivenFluxBC with
 no runoff modeled.
 
 These variables are updated in place in `boundary_flux`.
@@ -234,7 +235,7 @@ boundary_vars(
                               ::ClimaLand.TopBoundary)
 
 An extension of the `boundary_var_domain_names` method for RichardsAtmosDrivenFluxBC
-with no runoff modeled. 
+with no runoff modeled.
 """
 boundary_var_domain_names(
     bc::RichardsAtmosDrivenFluxBC{
@@ -244,11 +245,12 @@ boundary_var_domain_names(
     ::ClimaLand.TopBoundary,
 ) = (:surface, :surface)
 """
-    boundary_var_types(::RichardsAtmosDrivenFluxBC{<:PrescribedPrecipitation,
+    boundary_var_types(::RichardsModel{FT}
+                        ::RichardsAtmosDrivenFluxBC{<:PrescribedPrecipitation,
                                                    <:Runoff.AbstractRunoffModel,
                                                   },
-                       ::ClimaLand.TopBoundary,
-                       ) where {FT}
+                        ::ClimaLand.TopBoundary,
+                        ) where {FT}
 
 An extension of the `boundary_var_types` method for RichardsAtmosDrivenFluxBC
 with no runoff modeled.
@@ -447,7 +449,7 @@ function boundary_flux(
 end
 
 """
-    ClimaLand.∂tendencyBC∂Y(
+    ClimaLand.set_dfluxBCdY!(
         model::RichardsModel,
         ::MoistureStateBC,
         boundary::ClimaLand.TopBoundary,
@@ -466,7 +468,7 @@ variable, this is given by
 `∂T_N∂Y_N = [-∂/∂z(∂F_bc/∂Y_N)]_N`, where `N` indicates the top
 layer cell index.
 """
-function ClimaLand.∂tendencyBC∂Y(
+function ClimaLand.set_dfluxBCdY!(
     model::RichardsModel,
     ::MoistureStateBC,
     boundary::ClimaLand.TopBoundary,
@@ -487,13 +489,14 @@ function ClimaLand.∂tendencyBC∂Y(
         interpc2f_op.(dψdϑ.(hydrology_cm, Y.soil.ϑ_l, ν, θ_r, S_s)),
         face_len,
     )
-    return ClimaCore.Fields.FieldVector(;
-        soil = (; ϑ_l = @. -K / Δz * dψ / (2 * Δz)),
-    )
+    # Update `dfluxBCdY` in place
+    p.dfluxBCdY .=
+        ClimaCore.Fields.FieldVector(; soil = (; ϑ_l = @. -K * dψ / Δz))
+    return nothing
 end
 
 """
-    ClimaLand.∂tendencyBC∂Y(
+    ClimaLand.set_dfluxBCdY!(
         ::RichardsModel,
         ::AbstractWaterBC,
         boundary::ClimaLand.TopBoundary,
@@ -516,7 +519,8 @@ layer cell index.
 If `F_bc` can be approximated as independent of `Y_N`, the derivative
 is zero.
 """
-function ClimaLand.∂tendencyBC∂Y(
+# TODO do we even need to define this method? won't be called if not using MoistureStateBC at TopBoundary
+function ClimaLand.set_dfluxBCdY!(
     ::RichardsModel,
     ::AbstractWaterBC,
     boundary::ClimaLand.TopBoundary,
@@ -525,7 +529,7 @@ function ClimaLand.∂tendencyBC∂Y(
     p,
     t,
 )
-    return ClimaCore.Fields.FieldVector(; soil = (; ϑ_l = zeros(axes(Δz))))
+    return nothing
 end
 
 # BC type for the soil heat equation
@@ -808,6 +812,7 @@ boundary_var_domain_names(
 ) = (:surface, :surface, :surface, :surface)
 """
     boundary_var_types(
+        ::EnergyHydrology{FT},
         ::AtmosDrivenFluxBC{
             <:PrescribedAtmosphere{FT},
             <:AbstractRadiativeDrivers{FT},
@@ -886,7 +891,7 @@ end
                                     <:Runoff.TOPMODELRunoff,
                                     }, ::ClimaLand.TopBoundary)
 
-An extension of the `boundary_vars` method for AtmosDrivenFluxBC with 
+An extension of the `boundary_vars` method for AtmosDrivenFluxBC with
 TOPMODELRunoff. This
 adds the surface conditions (SHF, LHF, evaporation, and resistance) and the
 net radiation to the auxiliary variables.
@@ -924,6 +929,7 @@ boundary_var_domain_names(
 ) = (:surface, :surface, :surface, :surface, :surface, :surface, :surface)
 """
     boundary_var_types(
+        model::EnergyHydrology{FT},
         ::AtmosDrivenFluxBC{
             <:PrescribedAtmosphere{FT},
             <:AbstractRadiativeDrivers{FT},
@@ -952,3 +958,37 @@ boundary_var_types(
     FT,
     FT,
 )
+
+"""
+    boundary_vars(::MoistureStateBC, ::ClimaLand.TopBoundary)
+
+An extension of the `boundary_vars` method for MoistureStateBC at the
+top boundary.
+
+These variables are updated in place in `boundary_flux`.
+"""
+boundary_vars(bc::MoistureStateBC, ::ClimaLand.TopBoundary) =
+    (:top_bc, :dfluxBCdY)
+
+"""
+    boundary_var_domain_names(::MoistureStateBC, ::ClimaLand.TopBoundary)
+
+An extension of the `boundary_var_domain_names` method for MoistureStateBC at the
+top boundary.
+"""
+boundary_var_domain_names(bc::MoistureStateBC, ::ClimaLand.TopBoundary) =
+    (:surface, :surface)
+"""
+    boundary_var_types(::RichardsModel{FT},
+                        ::MoistureStateBC,
+                        ::ClimaLand.TopBoundary,
+                        ) where {FT}
+
+An extension of the `boundary_var_types` method for MoistureStateBC at the
+    top boundary.
+"""
+boundary_var_types(
+    model::RichardsModel{FT},
+    bc::MoistureStateBC,
+    ::ClimaLand.TopBoundary,
+) where {FT} = (FT, ClimaCore.MatrixFields.UpperBidiagonalMatrixRow{FT})

src/standalone/Soil/rre.jl

@@ -326,6 +326,19 @@ function ClimaLand.make_update_aux(model::RichardsModel)
             effective_saturation(ν, Y.soil.ϑ_l, θ_r),
         )
         @. p.soil.ψ = pressure_head(hydrology_cm, θ_r, Y.soil.ϑ_l, ν, S_s)
+
+        # If p.soil.dfluxBCdY is present, update it
+        if haskey(p, :dfluxBCdY)
+            set_dfluxBCdY!(
+                model,
+                model.boundary_conditions.top,
+                TopBoundary(),
+                Δz_top,
+                Y,
+                p,
+                t,
+            )
+        end
     end
     return update_aux!
 end
@@ -425,26 +438,48 @@ function ClimaLand.make_update_jacobian(model::RichardsModel{FT}) where {FT}
         # The derivative of the residual with respect to the prognostic variable
         ∂ϑres∂ϑ =
             matrix[MatrixFields.@name(soil.ϑ_l), MatrixFields.@name(soil.ϑ_l)]
-        @. ∂ϑres∂ϑ =
-            divf2c_matrix() ⋅
-            MatrixFields.DiagonalMatrixRow(interpc2f_op(p.soil.K)) ⋅
-            gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
-                ClimaLand.Soil.dψdϑ(hydrology_cm, Y.soil.ϑ_l, ν, θ_r, S_s),
+        # If the top BC is a `MoistureStateBC`, add the term from the top BC
+        #  flux before applying divergence
+        if haskey(p, :dfluxBCdY)
+            dfluxBCdY = p.dfluxBCdY
+            # TODO not sure if this usage is correct
+            topBC_op = Operators.SetBoundaryOperator(
+                top = Operators.SetValue(Geometry.WVector(dfluxBCdY)),
+                bottom = Operators.SetValue(Geometry.WVector(zero(FT))),
             )
+            # Add term from top boundary condition before applying divergence
+            @. ∂ϑres∂ϑ =
+                divf2c_matrix() ⋅ (
+                    MatrixFields.DiagonalMatrixRow(interpc2f_op(p.soil.K)) ⋅
+                    gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                        ClimaLand.Soil.dψdϑ(
+                            hydrology_cm,
+                            Y.soil.ϑ_l,
+                            ν,
+                            θ_r,
+                            S_s,
+                        ),
+                    ) + MatrixFields.BidiagonalMatrixRow(topBC_op)
+                )
+        else
+            @. ∂ϑres∂ϑ =
+                divf2c_matrix() ⋅
+                MatrixFields.DiagonalMatrixRow(interpc2f_op(p.soil.K)) ⋅
+                gradc2f_matrix() ⋅ MatrixFields.DiagonalMatrixRow(
+                    ClimaLand.Soil.dψdϑ(hydrology_cm, Y.soil.ϑ_l, ν, θ_r, S_s),
+                )
+
+        end
+
 
         # Hardcoded for single column: FIX!
         # The derivative of the tendency may eventually live in boundary vars and be updated there. but TBD
         z = Fields.coordinate_field(axes(Y.soil.ϑ_l)).z
         Δz_top, Δz_bottom = get_Δz(z)
-        ∂T_bc∂YN = ClimaLand.∂tendencyBC∂Y(
-            model,
-            model.boundary_conditions.top,
-            ClimaLand.TopBoundary(),
-            Δz_top,
-            Y,
-            p,
-            t,
-        )
+
+
+        ∂T_bc∂YN = p.dfluxBCdY
+
         #TODO: allocate space in W? See how final implementation of stencils with boundaries works out
         N = ClimaCore.Spaces.nlevels(axes(Y.soil.ϑ_l))
         parent(ClimaCore.Fields.level(∂ϑₜ∂ϑ.coefs.:2, N)) .=
@@ -478,9 +513,44 @@ end
 """
     is_saturated(Y, model::RichardsModel)
 
-A helper function which can be used to indicate whether a layer of soil is 
+A helper function which can be used to indicate whether a layer of soil is
 saturated.
 """
 function is_saturated(Y, model::RichardsModel)
     return @. ClimaLand.heaviside(Y.soil.ϑ_l - model.parameters.ν)
 end
+
+
+"""
+    add_dfluxBCdY_to_cache(p::NamedTuple,
+                            model::RichardsModel,
+                            bc::MoistureStateBC,
+                            ::ClimaLand.TopBoundary,
+                        )
+
+Allocate space in `p` for the derivative of the top boundary condition with
+respect to the state variables, and merges it into `p` under the key `dfluxBCdY`.
+Specifically, this stores the derivative of the top face flux with respect to
+the state variable at the top cell center.
+This is used in the implicit solver for
+`RichardsModel` when using a `MoistureStateBC` at the top boundary, in which
+case we need to handle the boundary specially in computing the Jacobian.
+
+"""
+# TODO should this go in p.dfluxBCdY or p.soil.dfluxBCdY?
+function ClimaLand.add_dfluxBCdY_to_cache(
+    p::NamedTuple,
+    model::RichardsModel,
+    bc::MoistureStateBC,
+    ::ClimaLand.TopBoundary,
+)
+    surface_space = model.domain.space.surface
+    FT = ClimaCore.Spaces.undertype(surface_space)
+
+    # Allocate a field containing a `BidiagonalMatrixRow` at each point
+    dfluxBCdY_field = ClimaCore.Fields.Field(
+        ClimaCore.MatrixFields.UpperBidiagonalMatrixRow{FT},
+        surface_space,
+    )
+    return merge(p, (; dfluxBCdY = dfluxBCdY_field))
+end