Merge CliMA#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

This PR also does a replace-all for `state_conservative` -> `state_prognostic`.



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

docs/src/APIs/BalanceLaws/BalanceLaws.md

@@ -10,21 +10,29 @@ 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
-init_state_conservative!
+init_state_prognostic!
 init_state_auxiliary!
 ```
 

docs/src/GettingStarted/Atmos.md

@@ -24,7 +24,7 @@ possible options for each subcomponent.
     source::S = (Gravity(), Coriolis(), GeostrophicForcing{FT}(7.62e-5, 0, 0)),
     tracers::TR = NoTracers(),
     boundarycondition::BC = AtmosBC(),
-    init_state_conservative::IS = nothing,
+    init_state_prognostic::IS = nothing,
     data_config::DC = nothing,
 ```
 
@@ -54,6 +54,6 @@ possible options for each subcomponent.
     source::S = (Gravity(), Coriolis()),
     tracers::TR = NoTracers(),
     boundarycondition::BC = AtmosBC(),
-    init_state_conservative::IS = nothing,
+    init_state_prognostic::IS = nothing,
     data_config::DC = nothing,
 ```

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 |
@@ -71,7 +71,7 @@ the following methods, which are computed locally at each nodal point:
 ## Methods to set initial conditions
 | **Method** | Purpose |
 |:-----|:-----|
-| [`init_state_conservative!`](@ref) | provide initial values for the conservative state as a function of time and space. |
+| [`init_state_prognostic!`](@ref) | provide initial values for the conservative state as a function of time and space. |
 | [`init_state_auxiliary!`](@ref) | provide initial values for the auxiliary variables as a function of the geometry. |
 
 
@@ -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,
@@ -125,7 +122,7 @@ tracers required.  Note that tracer naming is not currently supported,
 i.e. the user must track each tracer variable based on its numerical
 index. Sources can be added to each tracer based on the corresponding
 numerical vector index. Initial profiles must be specified using the
-`init_state_conservative!` hook at the experiment level.
+`init_state_prognostic!` hook at the experiment level.
 
 ```@docs
 ClimateMachine.Atmos.NTracers{N,FT}

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/AtmosGCM/heldsuarez.jl

@@ -100,7 +100,7 @@ function config_heldsuarez(FT, poly_order, resolution)
         hyperdiffusion = DryBiharmonic(τ_hyper),
         moisture = DryModel(),
         source = (Gravity(), Coriolis(), held_suarez_forcing!, sponge),
-        init_state_conservative = init_heldsuarez!,
+        init_state_prognostic = init_heldsuarez!,
         data_config = HeldSuarezDataConfig(T_ref),
         tracers = tracers,
     )

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
 
@@ -444,7 +444,7 @@ function config_bomex(FT, N, resolution, xmax, ymax, zmax)
             ),
             AtmosBC(),
         ),
-        init_state_conservative = ics,
+        init_state_prognostic = ics,
     )
 
     # Assemble configuration

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,
@@ -322,7 +320,7 @@ function config_dycoms(FT, N, resolution, xmax, ymax, zmax)
             ),
             AtmosBC(),
         ),
-        init_state_conservative = ics,
+        init_state_prognostic = ics,
     )
 
     ode_solver = ClimateMachine.ExplicitSolverType(

experiments/AtmosLES/schar_scalar_advection.jl

@@ -182,7 +182,7 @@ function config_schar(FT, N, resolution, xmax, ymax, zmax)
         moisture = DryModel(),
         source = source,
         tracers = NTracers{1, FT}(_δχ),
-        init_state_conservative = init_schar!,
+        init_state_prognostic = init_schar!,
         ref_state = ref_state,
     )
 

experiments/AtmosLES/surfacebubble.jl

@@ -105,7 +105,7 @@ function config_surfacebubble(FT, N, resolution, xmax, ymax, zmax)
             AtmosBC(),
         ),
         moisture = EquilMoist{FT}(),
-        init_state_conservative = init_surfacebubble!,
+        init_state_prognostic = init_surfacebubble!,
     )
     config = ClimateMachine.AtmosLESConfiguration(
         "SurfaceDrivenBubble",

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!,
@@ -136,7 +133,7 @@ function config_greenvortex(
         turbulence = Vreman(_C_smag),       # Turbulence closure model
         moisture = DryModel(),
         source = (),
-        init_state_conservative = init_greenvortex!,             # Apply the initial condition
+        init_state_prognostic = init_greenvortex!,             # Apply the initial condition
     )
 
     # Finally,  we pass a `Problem Name` string, the mesh information, and the model type to  the [`AtmosLESConfiguration`] object.

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