Merge #1280

1280: Add types and unify functions in BalanceLaws r=charleskawczynski a=charleskawczynski

# Description

 - Adds types to balance law (I have no strong opinions about the names here)

```julia
abstract type AbstractStateType end
struct Prognostic <: AbstractStateType end
struct Auxiliary <: AbstractStateType end
struct Gradient <: AbstractStateType end
struct GradientFlux <: AbstractStateType end
struct GradientLaplacian <: AbstractStateType end
struct Hyperdiffusive <: AbstractStateType end
struct UpwardIntegrals <: AbstractStateType end
struct DownwardIntegrals <: AbstractStateType end
```

 - Unifies `vars_state` (BalanceLaws) into a single function
 - Unifies `number_states` (BalanceLaws) into a single function
 - Unifies `create_state` (DGMethods) into a single function
 - Unifies `extract_state` (Diagnostics) into a single function



Co-authored-by: Charles Kawczynski <kawczynski.charles@gmail.com>

docs/src/APIs/BalanceLaws/BalanceLaws.md

@@ -10,15 +10,23 @@ CurrentModule = ClimateMachine.BalanceLaws
 BalanceLaw
 ```
 
+## State variable types
+
+```@docs
+Prognostic
+Auxiliary
+Gradient
+GradientFlux
+GradientLaplacian
+Hyperdiffusive
+UpwardIntegrals
+DownwardIntegrals
+```
+
 ## Variable specification methods
 
 ```@docs
-vars_state_conservative
-vars_state_auxiliary
-vars_state_gradient
-vars_integrals
-vars_reverse_integrals
-vars_state_gradient_flux
+vars_state
 ```
 
 ## Initial condition methods

docs/src/HowToGuides/Numerics/DGMethods/how_to_make_a_balance_law.md

@@ -35,12 +35,12 @@ the following methods, which are computed locally at each nodal point:
 ## Variable name specification methods
 | **Method** | Necessary | Purpose |
 |:-----|:---|:----|
-| [`vars_state_conservative`](@ref)         | **YES** |  specify the names of the variables in the conservative state vector, typically mass, momentum, and various tracers. |
-| [`vars_state_auxiliary`](@ref)           | **YES** |  specify the names of any variables required for the balance law that aren't related to derivatives of the state variables (e.g. spatial coordinates or various integrals) or those needed to solve expensive auxiliary equations (e.g., temperature via a non-linear equation solve)     |
-| [`vars_state_gradient`](@ref)         | **YES** |  specify the names of the gradients of functions of the conservative state variables. used to represent values before **and** after differentiation |
-| [`vars_state_gradient_flux`](@ref)         | **YES** |  specify the names of the gradient fluxes necessary to impose Neumann boundary conditions. typically the product of a diffusivity tensor with a gradient state variable, potentially equivalent to the second-order flux for a conservative state variable |
-| [`vars_integrals`](@ref)         | **NO** |  specify the names of any one-dimensional vertical integrals from **bottom to top** of the domain required for the balance law. used to represent both the integrand **and** the resulting indefinite integral |
-| [`vars_reverse_integrals`](@ref)         | **NO** |  specify the names of any one-dimensional vertical integral from **top to bottom** of the domain required for the balance law. each variable here must also exist in `vars_integrals` since the reverse integral kernels use subtraction to reverse the integral instead of performing a new integral. use to represent the value before **and** after reversing direction |
+| [`vars_state`](@ref)         | **YES** |  specify the names of the variables in the conservative state vector, typically mass, momentum, and various tracers. |
+| [`vars_state`](@ref)           | **YES** |  specify the names of any variables required for the balance law that aren't related to derivatives of the state variables (e.g. spatial coordinates or various integrals) or those needed to solve expensive auxiliary equations (e.g., temperature via a non-linear equation solve)     |
+| [`vars_state`](@ref)         | **YES** |  specify the names of the gradients of functions of the conservative state variables. used to represent values before **and** after differentiation |
+| [`vars_state`](@ref)         | **YES** |  specify the names of the gradient fluxes necessary to impose Neumann boundary conditions. typically the product of a diffusivity tensor with a gradient state variable, potentially equivalent to the second-order flux for a conservative state variable |
+| [`vars_state`](@ref)         | **NO** |  specify the names of any one-dimensional vertical integrals from **bottom to top** of the domain required for the balance law. used to represent both the integrand **and** the resulting indefinite integral |
+| [`vars_state`](@ref)         | **NO** |  specify the names of any one-dimensional vertical integral from **top to bottom** of the domain required for the balance law. each variable here must also exist in `vars_state` since the reverse integral kernels use subtraction to reverse the integral instead of performing a new integral. use to represent the value before **and** after reversing direction |
 
 ## Methods to compute gradients and integrals
 | **Method** |  Purpose |
@@ -84,7 +84,7 @@ argument inside these methods behave as dictionaries, for example:
 ```julia
 struct MyModel <: BalanceLaw end
 
-function vars_state_conservative(m::MyModel, FT)
+function vars_state(m::MyModel, ::Prognostic, FT)
     @vars begin
         ρ::FT
         T::FT

docs/src/Theory/Atmos/Model/tracers.md

@@ -45,10 +45,7 @@ Default stub functions for a generic tracer type are defined here.
 ```julia
 abstract type TracerModel <: BalanceLaw end
 
-vars_state_conservative(::TracerModel, FT) = @vars()
-vars_state_gradient(::TracerModel, FT) = @vars()
-vars_state_gradient_flux(::TracerModel, FT) = @vars()
-vars_state_auxiliary(::TracerModel, FT) = @vars()
+vars_state(::TracerModel, ::AbstractStateType, FT) = @vars()
 
 function atmos_init_aux!(
     ::TracerModel,

docs/src/Theory/Common/Turbulence.md

@@ -34,9 +34,7 @@ be defined for a turbulence model.
 abstract type TurbulenceClosure end
 
 
-vars_state_gradient((::TurbulenceClosure, FT) = @vars()
-vars_state_gradient_flux(::TurbulenceClosure, FT) = @vars()
-vars_state_auxiliary(::TurbulenceClosure, FT) = @vars()
+vars_state(::TurbulenceClosure, ::AbstractStateType, FT) = @vars()
 
 function atmos_init_aux!(
     ::TurbulenceClosure,
@@ -67,7 +65,7 @@ additional state variables or auxiliary variable updates (e.g. TKE based
 models)
 
 ```julia
-vars_state_conservative(::TurbulenceClosure, FT) = @vars()
+vars_state(::TurbulenceClosure, ::Prognostic, FT) = @vars()
 function atmos_nodal_update_auxiliary_state!(
     ::TurbulenceClosure,
     ::AtmosModel,

experiments/AtmosLES/bomex.jl

@@ -77,7 +77,7 @@ using CLIMAParameters.Planet: e_int_v0, grav, day
 struct EarthParameterSet <: AbstractEarthParameterSet end
 const param_set = EarthParameterSet()
 
-import ClimateMachine.BalanceLaws: vars_state_conservative, vars_state_auxiliary
+import ClimateMachine.BalanceLaws: vars_state
 import ClimateMachine.Atmos: source!, atmos_source!, altitude
 import ClimateMachine.Atmos: flux_second_order!, thermo_state
 

experiments/AtmosLES/dycoms.jl

@@ -14,6 +14,8 @@ using ClimateMachine.TemperatureProfiles
 using ClimateMachine.Thermodynamics
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.VariableTemplates
+using ClimateMachine.BalanceLaws:
+    AbstractStateType, Auxiliary, UpwardIntegrals, DownwardIntegrals
 
 using Distributions
 using Random
@@ -28,10 +30,7 @@ struct EarthParameterSet <: AbstractEarthParameterSet end
 const param_set = EarthParameterSet()
 
 import ClimateMachine.BalanceLaws:
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_integrals,
-    vars_reverse_integrals,
+    vars_state,
     indefinite_stack_integral!,
     reverse_indefinite_stack_integral!,
     integral_load_auxiliary_state!,
@@ -43,10 +42,7 @@ import ClimateMachine.BalanceLaws: boundary_state!
 import ClimateMachine.Atmos: flux_second_order!
 
 # -------------------- Radiation Model -------------------------- #
-vars_state_conservative(::RadiationModel, FT) = @vars()
-vars_state_auxiliary(::RadiationModel, FT) = @vars()
-vars_integrals(::RadiationModel, FT) = @vars()
-vars_reverse_integrals(::RadiationModel, FT) = @vars()
+vars_state(::RadiationModel, ::AbstractStateType, FT) = @vars()
 
 function atmos_nodal_update_auxiliary_state!(
     ::RadiationModel,
@@ -108,9 +104,10 @@ struct DYCOMSRadiation{FT} <: RadiationModel
     F_1::FT
 end
 
-vars_state_auxiliary(m::DYCOMSRadiation, FT) = @vars(Rad_flux::FT)
+vars_state(m::DYCOMSRadiation, ::Auxiliary, FT) = @vars(Rad_flux::FT)
 
-vars_integrals(m::DYCOMSRadiation, FT) = @vars(attenuation_coeff::FT)
+vars_state(m::DYCOMSRadiation, ::UpwardIntegrals, FT) =
+    @vars(attenuation_coeff::FT)
 function integral_load_auxiliary_state!(
     m::DYCOMSRadiation,
     integrand::Vars,
@@ -129,7 +126,8 @@ function integral_set_auxiliary_state!(
     aux.∫dz.radiation.attenuation_coeff = integral
 end
 
-vars_reverse_integrals(m::DYCOMSRadiation, FT) = @vars(attenuation_coeff::FT)
+vars_state(m::DYCOMSRadiation, ::DownwardIntegrals, FT) =
+    @vars(attenuation_coeff::FT)
 function reverse_integral_load_auxiliary_state!(
     m::DYCOMSRadiation,
     integrand::Vars,

experiments/AtmosLES/taylor-green.jl

@@ -26,10 +26,7 @@ struct EarthParameterSet <: AbstractEarthParameterSet end
 const param_set = EarthParameterSet()
 
 import ClimateMachine.BalanceLaws:
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_integrals,
-    vars_reverse_integrals,
+    vars_state,
     indefinite_stack_integral!,
     reverse_indefinite_stack_integral!,
     integral_load_auxiliary_state!,

experiments/OceanBoxGCM/homogeneous_box.jl

@@ -7,7 +7,7 @@ using ClimateMachine.ODESolvers
 using ClimateMachine.Mesh.Filters
 using ClimateMachine.VariableTemplates
 using ClimateMachine.Mesh.Grids: polynomialorder
-using ClimateMachine.BalanceLaws: vars_state_conservative
+using ClimateMachine.BalanceLaws: vars_state, Prognostic
 using ClimateMachine.Ocean.HydrostaticBoussinesq
 
 using Test
@@ -91,11 +91,11 @@ function run_homogeneous_box(; imex::Bool = false, BC = nothing)
 
     result = ClimateMachine.invoke!(solver_config)
 
-    maxQ = Vars{vars_state_conservative(driver_config.bl, FT)}(maximum(
+    maxQ = Vars{vars_state(driver_config.bl, Prognostic(), FT)}(maximum(
         solver_config.Q,
         dims = (1, 3),
     ))
-    minQ = Vars{vars_state_conservative(driver_config.bl, FT)}(minimum(
+    minQ = Vars{vars_state(driver_config.bl, Prognostic(), FT)}(minimum(
         solver_config.Q,
         dims = (1, 3),
     ))

src/Atmos/Model/AtmosModel.jl

@@ -36,18 +36,10 @@ using ..Mesh.Grids:
     EveryDirection,
     Direction
 
-using ClimateMachine.BalanceLaws:
-    BalanceLaw, number_state_conservative, num_integrals
+using ClimateMachine.BalanceLaws
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_gradient_laplacian,
-    vars_state_gradient_flux,
-    vars_hyperdiffusive,
-    vars_integrals,
-    vars_reverse_integrals,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -232,109 +224,109 @@ end
 
 
 """
-    vars_state_conservative(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::Prognostic, FT)
 Conserved state variables (Prognostic Variables)
 """
-function vars_state_conservative(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::Prognostic, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
         ρe::FT
-        turbulence::vars_state_conservative(m.turbulence, FT)
-        turbconv::vars_state_conservative(m.turbconv, FT)
-        hyperdiffusion::vars_state_conservative(m.hyperdiffusion, FT)
-        moisture::vars_state_conservative(m.moisture, FT)
-        radiation::vars_state_conservative(m.radiation, FT)
-        tracers::vars_state_conservative(m.tracers, FT)
+        turbulence::vars_state(m.turbulence, st, FT)
+        turbconv::vars_state(m.turbconv, st, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
+        moisture::vars_state(m.moisture, st, FT)
+        radiation::vars_state(m.radiation, st, FT)
+        tracers::vars_state(m.tracers, st, FT)
     end
 end
 
 """
-    vars_state_gradient(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::Gradient, FT)
 Pre-transform gradient variables
 """
-function vars_state_gradient(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::Gradient, FT)
     @vars begin
         u::SVector{3, FT}
         h_tot::FT
-        turbulence::vars_state_gradient(m.turbulence, FT)
-        turbconv::vars_state_gradient(m.turbconv, FT)
-        hyperdiffusion::vars_state_gradient(m.hyperdiffusion, FT)
-        moisture::vars_state_gradient(m.moisture, FT)
-        tracers::vars_state_gradient(m.tracers, FT)
+        turbulence::vars_state(m.turbulence, st, FT)
+        turbconv::vars_state(m.turbconv, st, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
+        moisture::vars_state(m.moisture, st, FT)
+        tracers::vars_state(m.tracers, st, FT)
     end
 end
 """
-    vars_state_gradient_flux(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::GradientFlux, FT)
 Post-transform gradient variables
 """
-function vars_state_gradient_flux(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::GradientFlux, FT)
     @vars begin
         ∇h_tot::SVector{3, FT}
-        turbulence::vars_state_gradient_flux(m.turbulence, FT)
-        turbconv::vars_state_gradient_flux(m.turbconv, FT)
-        hyperdiffusion::vars_state_gradient_flux(m.hyperdiffusion, FT)
-        moisture::vars_state_gradient_flux(m.moisture, FT)
-        tracers::vars_state_gradient_flux(m.tracers, FT)
+        turbulence::vars_state(m.turbulence, st, FT)
+        turbconv::vars_state(m.turbconv, st, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
+        moisture::vars_state(m.moisture, st, FT)
+        tracers::vars_state(m.tracers, st, FT)
     end
 end
 
 """
-    vars_gradient_laplacian(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::GradientLaplacian, FT)
 Pre-transform hyperdiffusive variables
 """
-function vars_gradient_laplacian(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::GradientLaplacian, FT)
     @vars begin
-        hyperdiffusion::vars_gradient_laplacian(m.hyperdiffusion, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
     end
 end
 
 """
-    vars_hyperdiffusive(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::Hyperdiffusive, FT)
 Post-transform hyperdiffusive variables
 """
-function vars_hyperdiffusive(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::Hyperdiffusive, FT)
     @vars begin
-        hyperdiffusion::vars_hyperdiffusive(m.hyperdiffusion, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
     end
 end
 
 """
-    vars_state_auxiliary(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::Auxiliary, FT)
 Auxiliary variables, such as vertical (stack)
 integrals, coordinates, orientation information,
 reference states, subcomponent auxiliary vars,
 debug variables
 """
-function vars_state_auxiliary(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::Auxiliary, FT)
     @vars begin
-        ∫dz::vars_integrals(m, FT)
-        ∫dnz::vars_reverse_integrals(m, FT)
+        ∫dz::vars_state(m, UpwardIntegrals(), FT)
+        ∫dnz::vars_state(m, DownwardIntegrals(), FT)
         coord::SVector{3, FT}
-        orientation::vars_state_auxiliary(m.orientation, FT)
-        ref_state::vars_state_auxiliary(m.ref_state, FT)
-        turbulence::vars_state_auxiliary(m.turbulence, FT)
-        turbconv::vars_state_auxiliary(m.turbconv, FT)
-        hyperdiffusion::vars_state_auxiliary(m.hyperdiffusion, FT)
-        moisture::vars_state_auxiliary(m.moisture, FT)
-        tracers::vars_state_auxiliary(m.tracers, FT)
-        radiation::vars_state_auxiliary(m.radiation, FT)
+        orientation::vars_state(m.orientation, st, FT)
+        ref_state::vars_state(m.ref_state, st, FT)
+        turbulence::vars_state(m.turbulence, st, FT)
+        turbconv::vars_state(m.turbconv, st, FT)
+        hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)
+        moisture::vars_state(m.moisture, st, FT)
+        tracers::vars_state(m.tracers, st, FT)
+        radiation::vars_state(m.radiation, st, FT)
     end
 end
 """
-    vars_integrals(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::UpwardIntegrals, FT)
 """
-function vars_integrals(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::UpwardIntegrals, FT)
     @vars begin
-        radiation::vars_integrals(m.radiation, FT)
+        radiation::vars_state(m.radiation, st, FT)
     end
 end
 """
-    vars_reverse_integrals(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::DownwardIntegrals, FT)
 """
-function vars_reverse_integrals(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::DownwardIntegrals, FT)
     @vars begin
-        radiation::vars_reverse_integrals(m.radiation, FT)
+        radiation::vars_state(m.radiation, st, FT)
     end
 end
 
@@ -558,8 +550,8 @@ end
     ss = soundspeed(m, m.moisture, state, aux)
 
     FT = typeof(state.ρ)
-    ws = fill(uN + ss, MVector{number_state_conservative(m, FT), FT})
-    vars_ws = Vars{vars_state_conservative(m, FT)}(ws)
+    ws = fill(uN + ss, MVector{number_states(m, Prognostic(), FT), FT})
+    vars_ws = Vars{vars_state(m, Prognostic(), FT)}(ws)
 
     wavespeed_tracers!(m.tracers, vars_ws, nM, state, aux, t)
 
@@ -577,7 +569,7 @@ function update_auxiliary_state!(
     FT = eltype(Q)
     state_auxiliary = dg.state_auxiliary
 
-    if num_integrals(m, FT) > 0
+    if number_states(m, UpwardIntegrals(), FT) > 0
         indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
         reverse_indefinite_stack_integral!(dg, m, Q, state_auxiliary, t, elems)
     end

src/Atmos/Model/filters.jl

@@ -11,7 +11,7 @@ struct AtmosFilterPerturbations{M} <: AbstractFilterTarget
 end
 
 vars_state_filtered(target::AtmosFilterPerturbations, FT) =
-    vars_state_conservative(target.atmos, FT)
+    vars_state(target.atmos, Prognostic(), FT)
 
 function compute_filter_argument!(
     ::AtmosFilterPerturbations,