From d60e2b27f0b87a37792d984ea08c7f6d810fb81d Mon Sep 17 00:00:00 2001 From: Julia Sloan Date: Wed, 13 Mar 2024 16:41:34 -0700 Subject: [PATCH] handle MoistureStateBC [skip ci] --- docs/src/APIs/shared_utilities.md | 2 +- src/shared_utilities/implicit_tendencies.jl | 6 +- src/shared_utilities/models.jl | 17 +++- src/standalone/Soil/boundary_conditions.jl | 78 +++++++++++---- src/standalone/Soil/rre.jl | 100 +++++++++++++++++--- 5 files changed, 164 insertions(+), 39 deletions(-) diff --git a/docs/src/APIs/shared_utilities.md b/docs/src/APIs/shared_utilities.md index 5b763c9dbb..545fa97881 100644 --- a/docs/src/APIs/shared_utilities.md +++ b/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 ``` diff --git a/src/shared_utilities/implicit_tendencies.jl b/src/shared_utilities/implicit_tendencies.jl index 9b22348a0e..4ec9a2c163 100644 --- a/src/shared_utilities/implicit_tendencies.jl +++ b/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, diff --git a/src/shared_utilities/models.jl b/src/shared_utilities/models.jl index f9ccbcf9c6..3a81d6bc7f 100644 --- a/src/shared_utilities/models.jl +++ b/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 diff --git a/src/standalone/Soil/boundary_conditions.jl b/src/standalone/Soil/boundary_conditions.jl index 4f53313b46..afd967e663 100644 --- a/src/standalone/Soil/boundary_conditions.jl +++ b/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}) diff --git a/src/standalone/Soil/rre.jl b/src/standalone/Soil/rre.jl index 1a91128b89..e77654c8e7 100644 --- a/src/standalone/Soil/rre.jl +++ b/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