jacobian wip [skip ci]

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()

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

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

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,

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(