diff --git a/src/shared_utilities/implicit_timestepping.jl b/src/shared_utilities/implicit_timestepping.jl index ccfe9128e0..8b2623d219 100644 --- a/src/shared_utilities/implicit_timestepping.jl +++ b/src/shared_utilities/implicit_timestepping.jl @@ -1,5 +1,6 @@ using ClimaCore.MatrixFields import ClimaCore.MatrixFields: @name, ⋅ +using ClimaCore: Spaces import UnrolledUtilities import LinearAlgebra import LinearAlgebra: I @@ -108,52 +109,64 @@ to the `explicit_names` tuple. """ function ImplicitEquationJacobian(Y::ClimaCore.Fields.FieldVector) FT = eltype(Y) - center_space = axes(Y.soil.ϑ_l) - - # Construct a tridiagonal matrix that will be used as the Jacobian - tridiag_type = MatrixFields.TridiagonalMatrixRow{FT} - # Create a field containing a `TridiagonalMatrixRow` at each point - tridiag_field = Fields.Field(tridiag_type, center_space) - fill!(parent(tridiag_field), NaN) - # Only add jacobian blocks for fields that are in Y for this model - is_in_Y(name) = MatrixFields.has_field(Y, name) + is_in_Y(var) = MatrixFields.has_field(Y, var) # Define the implicit and explicit variables of any model we use - implicit_names = (@name(soil.ϑ_l),) - explicit_names = ( + implicit_vars = (@name(soil.ϑ_l), @name(canopy.energy.T)) + explicit_vars = ( @name(soilco2.C), @name(soil.θ_i), @name(soil.ρe_int), @name(canopy.hydraulics.ϑ_l), - @name(canopy.energy.T), ) # Filter out the variables that are not in this model's state, `Y` - available_implicit_names = - MatrixFields.unrolled_filter(is_in_Y, implicit_names) - available_explicit_names = - MatrixFields.unrolled_filter(is_in_Y, explicit_names) + available_implicit_vars = + MatrixFields.unrolled_filter(is_in_Y, implicit_vars) + available_explicit_vars = + MatrixFields.unrolled_filter(is_in_Y, explicit_vars) + + + function get_j_field( + space::Union{ + Spaces.FiniteDifferenceSpace, + Spaces.ExtrudedFiniteDifferenceSpace, + }, + ) + # Construct a tridiagonal matrix that will be used as the Jacobian + tridiag_type = MatrixFields.TridiagonalMatrixRow{FT} + # Create a field containing a `TridiagonalMatrixRow` at each point + tridiag_field = ClimaCore.Fields.Field(tridiag_type, space) + fill!(parent(tridiag_field), 0) + return tridiag_field + end + function get_j_field( + space::Union{Spaces.PointSpace, Spaces.SpectralElementSpace2D}, + ) + diag_type = (I,) + diag_field = ClimaCore.Fields.Field(diag_type, space) + fill!(parent(diag_field), 0) + return diag_field + end + + implicit_blocks = MatrixFields.unrolled_map( + var -> + (var, var) => get_j_field(axes(MatrixFields.get_field(Y, var))), + available_implicit_vars, + ) # 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. - identity_blocks = MatrixFields.unrolled_map( - name -> (name, name) => FT(-1) * I, - available_explicit_names, + explicit_blocks = MatrixFields.unrolled_map( + var -> (var, var) => FT(-1) * I, + available_explicit_vars, ) - # For implicitly-stepped variables, use a tridiagonal matrix - tridiagonal_blocks = MatrixFields.unrolled_map( - name -> - (name, name) => Fields.Field( - tridiag_type, - axes(MatrixFields.get_field(Y, name)), - ), - available_implicit_names, - ) - matrix = MatrixFields.FieldMatrix(identity_blocks..., tridiagonal_blocks...) + + matrix = MatrixFields.FieldMatrix(implicit_blocks..., explicit_blocks...) # Set up block diagonal solver for block Jacobian alg = MatrixFields.BlockDiagonalSolve() diff --git a/src/standalone/Vegetation/Canopy.jl b/src/standalone/Vegetation/Canopy.jl index 457a1b1fed..a0c2f7941c 100644 --- a/src/standalone/Vegetation/Canopy.jl +++ b/src/standalone/Vegetation/Canopy.jl @@ -3,6 +3,9 @@ using DocStringExtensions using Thermodynamics using ClimaLand using ClimaCore +using ClimaCore.MatrixFields +import ClimaCore.MatrixFields: @name, ⋅ +import LinearAlgebra: I using ClimaLand: AbstractRadiativeDrivers, AbstractAtmosphericDrivers import ..Parameters as LP @@ -20,6 +23,8 @@ import ClimaLand: make_update_boundary_fluxes, make_update_aux, make_compute_exp_tendency, + make_compute_imp_tendency, + make_update_jacobian, get_drivers using ClimaLand.Domains: Point, Plane, SphericalSurface @@ -617,6 +622,68 @@ function make_compute_exp_tendency( end return compute_exp_tendency! end + +""" + make_compute_imp_tendency(canopy::CanopyModel) + +Creates and returns the compute_imp_tendency! for the `CanopyModel`. +""" +function make_compute_imp_tendency( + canopy::CanopyModel{ + FT, + <:AutotrophicRespirationModel, + <:Union{BeerLambertModel, TwoStreamModel}, + <:Union{FarquharModel, OptimalityFarquharModel}, + <:MedlynConductanceModel, + <:PlantHydraulicsModel, + <:Union{PrescribedCanopyTempModel, BigLeafEnergyModel}, + }, +) where {FT} + components = canopy_components(canopy) + compute_imp_tendency_list = map( + x -> make_compute_imp_tendency(getproperty(canopy, x), canopy), + components, + ) + function compute_imp_tendency!(dY, Y, p, t) + for f! in compute_imp_tendency_list + f!(dY, Y, p, t) + end + + end + return compute_imp_tendency! +end + +""" + make_update_jacobian(canopy::CanopyModel) + +Creates and returns the update_jacobian! for the `CanopyModel`. +""" +function make_update_jacobian( + canopy::CanopyModel{ + FT, + <:AutotrophicRespirationModel, + <:Union{BeerLambertModel, TwoStreamModel}, + <:Union{FarquharModel, OptimalityFarquharModel}, + <:MedlynConductanceModel, + <:PlantHydraulicsModel, + <:Union{PrescribedCanopyTempModel, BigLeafEnergyModel}, + }, +) where {FT} + components = canopy_components(canopy) + update_jacobian_list = map( + x -> make_update_jacobian(getproperty(canopy, x), canopy), + components, + ) + function update_jacobian!(W, Y, p, dtγ, t) + for f! in update_jacobian_list + f!(W, Y, p, dtγ, t) + end + + end + return update_jacobian! +end + + function ClimaLand.get_drivers(model::CanopyModel) return (model.atmos, model.radiation) end diff --git a/src/standalone/Vegetation/canopy_energy.jl b/src/standalone/Vegetation/canopy_energy.jl index 94dc934d10..389991cf86 100644 --- a/src/standalone/Vegetation/canopy_energy.jl +++ b/src/standalone/Vegetation/canopy_energy.jl @@ -87,11 +87,11 @@ where the canopy temperature is modeled prognostically. canopy_temperature(model::BigLeafEnergyModel, canopy, Y, p, t) = Y.canopy.energy.T -function make_compute_exp_tendency( +function make_compute_imp_tendency( model::BigLeafEnergyModel{FT}, canopy, ) where {FT} - function compute_exp_tendency!(dY, Y, p, t) + function compute_imp_tendency!(dY, Y, p, t) area_index = p.canopy.hydraulics.area_index ac_canopy = model.parameters.ac_canopy # Energy Equation: @@ -114,7 +114,7 @@ function make_compute_exp_tendency( @. ac_canopy * max(area_index.leaf + area_index.stem, eps(FT)) @. dY.canopy.energy.T = -net_energy_flux / c_per_ground_area end - return compute_exp_tendency! + return compute_imp_tendency! end """ @@ -151,3 +151,17 @@ function root_energy_flux_per_ground_area!( ) where {FT} fa_energy .= FT(0) end + +function ClimaLand.make_update_jacobian( + model::BigLeafEnergyModel{FT}, + canopy, +) where {FT} + function update_jacobian!(jacobian::ImplicitEquationJacobian, Y, p, dtγ, t) + (; matrix) = jacobian + + # The derivative of the residual with respect to the prognostic variable + ∂Tres∂T = matrix[@name(canopy.energy.T), @name(canopy.energy.T)] + @. ∂Tres∂T = dtγ * 1 - (I,) + end + return update_jacobian! +end diff --git a/src/standalone/Vegetation/component_models.jl b/src/standalone/Vegetation/component_models.jl index 3697e0ce7c..2b95ec0c79 100644 --- a/src/standalone/Vegetation/component_models.jl +++ b/src/standalone/Vegetation/component_models.jl @@ -123,11 +123,48 @@ may be needed and passed in (via the `canopy` model itself). The right hand side for the entire canopy model can make use of these functions for the individual components. """ -function ClimaLand.make_compute_exp_tendency(::AbstractCanopyComponent, canopy) - function compute_exp_tendency!(dY, Y, p, t) end +function ClimaLand.make_compute_exp_tendency( + component::AbstractCanopyComponent, + canopy, +) + function compute_exp_tendency!(dY, Y, p, t) + vars = prognostic_vars(component) + if vars != () + getproperty(getproperty(dY, name(canopy)), name(component)) .= 0 + end + end return compute_exp_tendency! end +""" + ClimaLand.make_compute_exp_tendency(component::AbstractCanopyComponent, canopy) + +Creates the compute_imp_tendency!(dY,Y,p,t) function for the canopy `component`. + +Since component models are not standalone models, other information +may be needed and passed in (via the `canopy` model itself). +The right hand side for the entire canopy model can make use of +these functions for the individual components. +""" +function ClimaLand.make_compute_imp_tendency( + component::AbstractCanopyComponent, + canopy, +) + function compute_imp_tendency!(dY, Y, p, t) + vars = prognostic_vars(component) + if vars != () + getproperty(getproperty(dY, name(canopy)), name(component)) .= 0 + end + + end + return compute_imp_tendency! +end + + +function make_update_jacobian(component::AbstractCanopyComponent, canopy) + function update_jacobian!(W, Y, p, dtγ, t) end + return update_jacobian! +end """ set_canopy_prescribed_field!(component::AbstractCanopyComponent, diff --git a/test/standalone/Vegetation/canopy_model.jl b/test/standalone/Vegetation/canopy_model.jl index 8df8895f5f..2c4a9f876a 100644 --- a/test/standalone/Vegetation/canopy_model.jl +++ b/test/standalone/Vegetation/canopy_model.jl @@ -719,6 +719,14 @@ for FT in (Float32, Float64) t0 = FT(0.0) set_initial_cache!(p, Y, t0) exp_tendency! = make_exp_tendency(canopy) + imp_tendency! = ClimaLand.make_imp_tendency(canopy) + tendency_jacobian! = ClimaLand.make_tendency_jacobian(canopy) + # set up jacobian info + jac_kwargs = (; + jac_prototype = ImplicitEquationJacobian(Y), + Wfact = tendency_jacobian!, + ) + dY = similar(Y) exp_tendency!(dY, Y, p, t0) turb_fluxes = ClimaLand.Canopy.canopy_turbulent_fluxes(