diff --git a/docs/src/APIs/BalanceLaws/BalanceLaws.md b/docs/src/APIs/BalanceLaws/BalanceLaws.md
index e4fd778d231..1c8a4e9bce9 100644
--- a/docs/src/APIs/BalanceLaws/BalanceLaws.md
+++ b/docs/src/APIs/BalanceLaws/BalanceLaws.md
@@ -13,12 +13,7 @@ BalanceLaw
 ## 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
diff --git a/docs/src/HowToGuides/Numerics/DGMethods/how_to_make_a_balance_law.md b/docs/src/HowToGuides/Numerics/DGMethods/how_to_make_a_balance_law.md
index 222d2927747..6dc14b240d6 100644
--- a/docs/src/HowToGuides/Numerics/DGMethods/how_to_make_a_balance_law.md
+++ b/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, ::Conservative, FT)
     @vars begin
         ρ::FT
         T::FT
diff --git a/docs/src/Theory/Atmos/Model/tracers.md b/docs/src/Theory/Atmos/Model/tracers.md
index eed892954ed..81f27797565 100644
--- a/docs/src/Theory/Atmos/Model/tracers.md
+++ b/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,
diff --git a/docs/src/Theory/Common/Turbulence.md b/docs/src/Theory/Common/Turbulence.md
index 9ae74f820dc..216b867a77a 100644
--- a/docs/src/Theory/Common/Turbulence.md
+++ b/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, ::Conservative, FT) = @vars()
 function atmos_nodal_update_auxiliary_state!(
     ::TurbulenceClosure,
     ::AtmosModel,
diff --git a/experiments/AtmosLES/bomex.jl b/experiments/AtmosLES/bomex.jl
index 21e4b07309d..942ef236bc3 100755
--- a/experiments/AtmosLES/bomex.jl
+++ b/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
 
diff --git a/experiments/AtmosLES/dycoms.jl b/experiments/AtmosLES/dycoms.jl
index a97be5e4cb6..6d0eaa95fea 100755
--- a/experiments/AtmosLES/dycoms.jl
+++ b/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,
diff --git a/experiments/AtmosLES/taylor-green.jl b/experiments/AtmosLES/taylor-green.jl
index 92a5c1ebe91..d6fcd64e50c 100644
--- a/experiments/AtmosLES/taylor-green.jl
+++ b/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!,
diff --git a/experiments/OceanBoxGCM/homogeneous_box.jl b/experiments/OceanBoxGCM/homogeneous_box.jl
index 494832830e7..e93d1a42dfd 100755
--- a/experiments/OceanBoxGCM/homogeneous_box.jl
+++ b/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, Conservative
 using ClimateMachine.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, Conservative(), 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, Conservative(), FT)}(minimum(
         solver_config.Q,
         dims = (1, 3),
     ))
diff --git a/src/Atmos/Model/AtmosModel.jl b/src/Atmos/Model/AtmosModel.jl
index 70656fe0040..8d551dabc5b 100644
--- a/src/Atmos/Model/AtmosModel.jl
+++ b/src/Atmos/Model/AtmosModel.jl
@@ -34,18 +34,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!,
@@ -219,105 +211,105 @@ end
 
 
 """
-    vars_state_conservative(m::AtmosModel, FT)
+    vars_state(m::AtmosModel, ::Conservative, FT)
 Conserved state variables (Prognostic Variables)
 """
-function vars_state_conservative(m::AtmosModel, FT)
+function vars_state(m::AtmosModel, st::Conservative, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
         ρe::FT
-        turbulence::vars_state_conservative(m.turbulence, 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)
+        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)
-        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)
+        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)
-        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)
+        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)
-        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)
+        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
 
@@ -518,8 +510,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, Conservative(), FT), FT})
+    vars_ws = Vars{vars_state(m, Conservative(), FT)}(ws)
 
     wavespeed_tracers!(m.tracers, vars_ws, nM, state, aux, t)
 
@@ -537,7 +529,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
diff --git a/src/Atmos/Model/filters.jl b/src/Atmos/Model/filters.jl
index 74ef2104aef..d89ace0d9ca 100644
--- a/src/Atmos/Model/filters.jl
+++ b/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, Conservative(), FT)
 
 function compute_filter_argument!(
     ::AtmosFilterPerturbations,
diff --git a/src/Atmos/Model/hyperdiffusion.jl b/src/Atmos/Model/hyperdiffusion.jl
index 1418999bd50..52ffc89813f 100644
--- a/src/Atmos/Model/hyperdiffusion.jl
+++ b/src/Atmos/Model/hyperdiffusion.jl
@@ -4,12 +4,8 @@ using LinearAlgebra
 export HyperDiffusion, NoHyperDiffusion, StandardHyperDiffusion
 
 abstract type HyperDiffusion end
-vars_state_conservative(::HyperDiffusion, FT) = @vars()
-vars_state_auxiliary(::HyperDiffusion, FT) = @vars()
-vars_state_gradient(::HyperDiffusion, FT) = @vars()
-vars_gradient_laplacian(::HyperDiffusion, FT) = @vars()
-vars_state_gradient_flux(::HyperDiffusion, FT) = @vars()
-vars_hyperdiffusive(::HyperDiffusion, FT) = @vars()
+vars_state(::HyperDiffusion, ::AbstractStateType, FT) = @vars()
+
 function atmos_init_aux!(
     ::HyperDiffusion,
     ::AtmosModel,
@@ -72,12 +68,12 @@ struct StandardHyperDiffusion{FT} <: HyperDiffusion
     τ_timescale::FT
 end
 
-vars_state_auxiliary(::StandardHyperDiffusion, FT) = @vars(Δ::FT)
-vars_state_gradient(::StandardHyperDiffusion, FT) =
+vars_state(::StandardHyperDiffusion, ::Auxiliary, FT) = @vars(Δ::FT)
+vars_state(::StandardHyperDiffusion, ::Gradient, FT) =
     @vars(u::SVector{3, FT}, h_tot::FT)
-vars_gradient_laplacian(::StandardHyperDiffusion, FT) =
+vars_state(::StandardHyperDiffusion, ::GradientLaplacian, FT) =
     @vars(u::SVector{3, FT}, h_tot::FT)
-vars_hyperdiffusive(::StandardHyperDiffusion, FT) =
+vars_state(::StandardHyperDiffusion, ::Hyperdiffusive, FT) =
     @vars(ν∇³u::SMatrix{3, 3, FT, 9}, ν∇³h_tot::SVector{3, FT})
 
 function atmos_init_aux!(
diff --git a/src/Atmos/Model/linear.jl b/src/Atmos/Model/linear.jl
index 985466cadbe..b518d67d36c 100644
--- a/src/Atmos/Model/linear.jl
+++ b/src/Atmos/Model/linear.jl
@@ -64,22 +64,19 @@ end
 
 abstract type AtmosLinearModel <: BalanceLaw end
 
-function vars_state_conservative(lm::AtmosLinearModel, FT)
+function vars_state(lm::AtmosLinearModel, ::Conservative, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
         ρe::FT
-        turbulence::vars_state_conservative(lm.atmos.turbulence, FT)
-        hyperdiffusion::vars_state_conservative(lm.atmos.hyperdiffusion, FT)
-        moisture::vars_state_conservative(lm.atmos.moisture, FT)
+        turbulence::vars_state(lm.atmos.turbulence, Conservative(), FT)
+        hyperdiffusion::vars_state(lm.atmos.hyperdiffusion, Conservative(), FT)
+        moisture::vars_state(lm.atmos.moisture, Conservative(), FT)
     end
 end
-vars_state_gradient(lm::AtmosLinearModel, FT) = @vars()
-vars_state_gradient_flux(lm::AtmosLinearModel, FT) = @vars()
-vars_state_auxiliary(lm::AtmosLinearModel, FT) =
-    vars_state_auxiliary(lm.atmos, FT)
-vars_integrals(lm::AtmosLinearModel, FT) = @vars()
-vars_reverse_integrals(lm::AtmosLinearModel, FT) = @vars()
+vars_state(lm::AtmosLinearModel, ::AbstractStateType, FT) = @vars()
+vars_state(lm::AtmosLinearModel, ::Auxiliary, FT) =
+    vars_state(lm.atmos, Auxiliary(), FT)
 
 
 function update_auxiliary_state!(
diff --git a/src/Atmos/Model/moisture.jl b/src/Atmos/Model/moisture.jl
index 7ac1281c7a9..58c49b07156 100644
--- a/src/Atmos/Model/moisture.jl
+++ b/src/Atmos/Model/moisture.jl
@@ -3,10 +3,7 @@ export DryModel, EquilMoist
 #### Moisture component in atmosphere model
 abstract type MoistureModel end
 
-vars_state_conservative(::MoistureModel, FT) = @vars()
-vars_state_gradient(::MoistureModel, FT) = @vars()
-vars_state_gradient_flux(::MoistureModel, FT) = @vars()
-vars_state_auxiliary(::MoistureModel, FT) = @vars()
+vars_state(::MoistureModel, ::AbstractStateType, FT) = @vars()
 
 function atmos_nodal_update_auxiliary_state!(
     ::MoistureModel,
@@ -93,7 +90,7 @@ Assumes the moisture components is in the dry limit.
 """
 struct DryModel <: MoistureModel end
 
-vars_state_auxiliary(::DryModel, FT) = @vars(θ_v::FT, air_T::FT)
+vars_state(::DryModel, ::Auxiliary, FT) = @vars(θ_v::FT, air_T::FT)
 @inline function atmos_nodal_update_auxiliary_state!(
     moist::DryModel,
     atmos::AtmosModel,
@@ -139,10 +136,10 @@ EquilMoist{FT}(;
 ) where {FT <: AbstractFloat, IT <: Int} = EquilMoist{FT}(maxiter, tolerance)
 
 
-vars_state_conservative(::EquilMoist, FT) = @vars(ρq_tot::FT)
-vars_state_gradient(::EquilMoist, FT) = @vars(q_tot::FT)
-vars_state_gradient_flux(::EquilMoist, FT) = @vars(∇q_tot::SVector{3, FT})
-vars_state_auxiliary(::EquilMoist, FT) =
+vars_state(::EquilMoist, ::Conservative, FT) = @vars(ρq_tot::FT)
+vars_state(::EquilMoist, ::Gradient, FT) = @vars(q_tot::FT)
+vars_state(::EquilMoist, ::GradientFlux, FT) = @vars(∇q_tot::SVector{3, FT})
+vars_state(::EquilMoist, ::Auxiliary, FT) =
     @vars(temperature::FT, θ_v::FT, q_liq::FT)
 
 @inline function atmos_nodal_update_auxiliary_state!(
diff --git a/src/Atmos/Model/precipitation.jl b/src/Atmos/Model/precipitation.jl
index 6378180fdf7..eecc2e5c41f 100644
--- a/src/Atmos/Model/precipitation.jl
+++ b/src/Atmos/Model/precipitation.jl
@@ -5,10 +5,7 @@ export NoPrecipitation, Rain
 
 using ..Microphysics
 
-vars_state_conservative(::PrecipitationModel, FT) = @vars()
-vars_state_gradient(::PrecipitationModel, FT) = @vars()
-vars_state_gradient_flux(::PrecipitationModel, FT) = @vars()
-vars_state_auxiliary(::PrecipitationModel, FT) = @vars()
+vars_state(::PrecipitationModel, ::AbstractStateType, FT) = @vars()
 
 function atmos_nodal_update_auxiliary_state!(
     ::PrecipitationModel,
@@ -66,10 +63,10 @@ Precipitation model with rain only.
 """
 struct Rain <: PrecipitationModel end
 
-vars_state_conservative(::Rain, FT) = @vars(ρq_rain::FT)
-vars_state_gradient(::Rain, FT) = @vars(q_rain::FT)
-vars_state_gradient_flux(::Rain, FT) = @vars(ρd_q_rain::SVector{3, FT})
-vars_state_auxiliary(::Rain, FT) =
+vars_state(::Rain, ::Conservative, FT) = @vars(ρq_rain::FT)
+vars_state(::Rain, ::Gradient, FT) = @vars(q_rain::FT)
+vars_state(::Rain, ::GradientFlux, FT) = @vars(ρd_q_rain::SVector{3, FT})
+vars_state(::Rain, ::Auxiliary, FT) =
     @vars(terminal_velocity::FT, src_q_rai_tot::FT)
 
 function atmos_nodal_update_auxiliary_state!(
diff --git a/src/Atmos/Model/radiation.jl b/src/Atmos/Model/radiation.jl
index 1d3c05a36fc..5528294cd0b 100644
--- a/src/Atmos/Model/radiation.jl
+++ b/src/Atmos/Model/radiation.jl
@@ -2,10 +2,7 @@ export RadiationModel, NoRadiation
 
 abstract type RadiationModel end
 
-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,
diff --git a/src/Atmos/Model/ref_state.jl b/src/Atmos/Model/ref_state.jl
index 413edee8b91..0a47315cf3b 100644
--- a/src/Atmos/Model/ref_state.jl
+++ b/src/Atmos/Model/ref_state.jl
@@ -13,10 +13,8 @@ condition or for linearization.
 """
 abstract type ReferenceState end
 
-vars_state_conservative(m::ReferenceState, FT) = @vars()
-vars_state_gradient(m::ReferenceState, FT) = @vars()
-vars_state_gradient_flux(m::ReferenceState, FT) = @vars()
-vars_state_auxiliary(m::ReferenceState, FT) = @vars()
+vars_state(m::ReferenceState, ::AbstractStateType, FT) = @vars()
+
 atmos_init_aux!(
     ::ReferenceState,
     ::AtmosModel,
@@ -49,7 +47,7 @@ function HydrostaticState(
     )
 end
 
-vars_state_auxiliary(m::HydrostaticState, FT) =
+vars_state(m::HydrostaticState, ::Auxiliary, FT) =
     @vars(ρ::FT, p::FT, T::FT, ρe::FT, ρq_tot::FT)
 
 
diff --git a/src/Atmos/Model/tracers.jl b/src/Atmos/Model/tracers.jl
index a2165a92a9a..41299323ad9 100644
--- a/src/Atmos/Model/tracers.jl
+++ b/src/Atmos/Model/tracers.jl
@@ -37,10 +37,7 @@ abstract type TracerModel end
 
 export NoTracers, NTracers
 
-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,
@@ -148,11 +145,11 @@ struct NTracers{N, FT} <: TracerModel
     δ_χ::SVector{N, FT}
 end
 
-vars_state_conservative(tr::NTracers, FT) = @vars(ρχ::typeof(tr.δ_χ))
-vars_state_gradient(tr::NTracers, FT) = @vars(χ::typeof(tr.δ_χ))
-vars_state_gradient_flux(tr::NTracers, FT) =
+vars_state(tr::NTracers, ::Conservative, FT) = @vars(ρχ::typeof(tr.δ_χ))
+vars_state(tr::NTracers, ::Gradient, FT) = @vars(χ::typeof(tr.δ_χ))
+vars_state(tr::NTracers, ::GradientFlux, FT) =
     @vars(∇χ::SMatrix{3, length(tr.δ_χ), FT, 3 * length(tr.δ_χ)})
-vars_state_auxiliary(tr::NTracers, FT) = @vars(δ_χ::typeof(tr.δ_χ))
+vars_state(tr::NTracers, ::Auxiliary, FT) = @vars(δ_χ::typeof(tr.δ_χ))
 
 function atmos_init_aux!(
     tr::NTracers,
diff --git a/src/BalanceLaws/BalanceLaws.jl b/src/BalanceLaws/BalanceLaws.jl
index 6d57fbc9b6b..1d58242a36f 100644
--- a/src/BalanceLaws/BalanceLaws.jl
+++ b/src/BalanceLaws/BalanceLaws.jl
@@ -12,10 +12,8 @@ using ..VariableTemplates
 using ..Courant
 
 export BalanceLaw,
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
+    number_states,
     init_state_conservative!,
     init_state_auxiliary!,
     flux_first_order!,
@@ -35,6 +33,7 @@ export BalanceLaw,
     reverse_integral_load_auxiliary_state!,
     reverse_integral_set_auxiliary_state!
 
+include("state_types.jl")
 include("interface.jl")
 
 end
diff --git a/src/BalanceLaws/interface.jl b/src/BalanceLaws/interface.jl
index 8e4048fd5ba..412ee3a3744 100644
--- a/src/BalanceLaws/interface.jl
+++ b/src/BalanceLaws/interface.jl
@@ -16,69 +16,15 @@ Subtypes `L` should define the methods below
 abstract type BalanceLaw end # PDE part
 
 """
-    vars_state_conservative(::L, FT)
+    vars_state(::L, ::Conservative, FT)
 
 a tuple of symbols containing the state variables
 given a float type `FT`.
 """
-function vars_state_conservative end
+function vars_state end
 
-"""
-    vars_state_auxiliary(::L, FT)
-
-a tuple of symbols containing the auxiliary variables
-given a float type `FT`.
-"""
-function vars_state_auxiliary end
-
-"""
-    vars_state_gradient(::L, FT)
-
-a tuple of symbols containing the transformed variables
-of which gradients are computed given a float type `FT`.
-"""
-function vars_state_gradient end
-
-"""
-    vars_state_gradient_flux(::L, FT)
-
-a tuple of symbols containing the diffusive variables
-given a float type `FT`.
-"""
-function vars_state_gradient_flux end
-
-"""
-    vars_gradient_laplacian(::L, FT)
-
-a tuple of symbols containing the transformed variables
-of which gradients of laplacian are computed, they must
-be a subset of `vars_state_gradient`, given a float type `FT`.
-"""
-vars_gradient_laplacian(::BalanceLaw, FT) = @vars()
-
-"""
-    vars_hyperdiffusive(::L, FT)
-
-a tuple of symbols containing the hyperdiffusive variables
-given a float type `FT`.
-"""
-vars_hyperdiffusive(::BalanceLaw, FT) = @vars()
-
-"""
-    vars_integrals(::L, FT)
-
-a tuple of symbols containing variables to be integrated
-along a vertical stack, given a float type `FT`.
-"""
-vars_integrals(::BalanceLaw, FT) = @vars()
-
-"""
-    vars_reverse_integrals(::L, FT)
-
-a tuple of symbols containing variables to be integrated
-along a vertical stack, in reverse, given a float type `FT`.
-"""
-vars_reverse_integrals(::BalanceLaw, FT) = @vars()
+# Fallback: no variables
+vars_state(::BalanceLaw, ::AbstractStateType, FT) = @vars()
 
 """
     init_state_conservative!(
@@ -316,18 +262,11 @@ Compute reverse indefinite integral along stack.
 function reverse_indefinite_stack_integral! end
 
 # Internal methods
-number_state_conservative(m::BalanceLaw, FT) =
-    varsize(vars_state_conservative(m, FT))
-number_state_auxiliary(m::BalanceLaw, FT) = varsize(vars_state_auxiliary(m, FT))
-number_state_gradient(m::BalanceLaw, FT) = varsize(vars_state_gradient(m, FT))
-number_state_gradient_flux(m::BalanceLaw, FT) =
-    varsize(vars_state_gradient_flux(m, FT))
-num_gradient_laplacian(m::BalanceLaw, FT) =
-    varsize(vars_gradient_laplacian(m, FT))
-num_hyperdiffusive(m::BalanceLaw, FT) = varsize(vars_hyperdiffusive(m, FT))
-num_integrals(m::BalanceLaw, FT) = varsize(vars_integrals(m, FT))
-num_reverse_integrals(m::BalanceLaw, FT) =
-    varsize(vars_reverse_integrals(m, FT))
+number_states(m::BalanceLaw, st::AbstractStateType, FT) =
+    varsize(vars_state(m, st, FT))
+
+rank_multiplier(st::AbstractStateType) = 1
+rank_multiplier(st::GradientLaplacian) = 3
 
 ### split explicit functions
 function initialize_states! end
diff --git a/src/BalanceLaws/state_types.jl b/src/BalanceLaws/state_types.jl
new file mode 100644
index 00000000000..dd832f4a139
--- /dev/null
+++ b/src/BalanceLaws/state_types.jl
@@ -0,0 +1,21 @@
+#### State types
+
+export AbstractStateType,
+    Conservative,
+    Auxiliary,
+    Gradient,
+    GradientFlux,
+    GradientLaplacian,
+    Hyperdiffusive,
+    UpwardIntegrals,
+    DownwardIntegrals
+
+abstract type AbstractStateType end
+struct Conservative <: 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
diff --git a/src/Common/Orientations/Orientations.jl b/src/Common/Orientations/Orientations.jl
index 926e8506c64..242c27ec51c 100644
--- a/src/Common/Orientations/Orientations.jl
+++ b/src/Common/Orientations/Orientations.jl
@@ -26,7 +26,8 @@ using StaticArrays
 using LinearAlgebra
 
 using ..VariableTemplates
-import ..BalanceLaws: BalanceLaw, vars_state_auxiliary
+using ..BalanceLaws
+import ..BalanceLaws: vars_state
 
 export Orientation, NoOrientation, FlatOrientation, SphericalOrientation
 
@@ -47,7 +48,7 @@ abstract type Orientation <: BalanceLaw end
 ##### Fallbacks
 #####
 
-function vars_state_auxiliary(m::Orientation, FT)
+function vars_state(m::Orientation, ::Auxiliary, FT)
     @vars begin
         Φ::FT # gravitational potential
         ∇Φ::SVector{3, FT}
@@ -103,7 +104,8 @@ No gravitational force or potential.
 struct NoOrientation <: Orientation end
 
 init_aux!(::NoOrientation, param_set::APS, aux::Vars) = nothing
-vars_state_auxiliary(m::NoOrientation, FT) = @vars()
+
+vars_state(m::NoOrientation, ::Auxiliary, FT) = @vars()
 
 gravitational_potential(::NoOrientation, aux::Vars) = -zero(eltype(aux))
 ∇gravitational_potential(::NoOrientation, aux::Vars) =
diff --git a/src/Common/TurbulenceClosures/TurbulenceClosures.jl b/src/Common/TurbulenceClosures/TurbulenceClosures.jl
index e90d18b6538..092f55ffe4c 100644
--- a/src/Common/TurbulenceClosures/TurbulenceClosures.jl
+++ b/src/Common/TurbulenceClosures/TurbulenceClosures.jl
@@ -2,7 +2,7 @@
     TurbulenceClosures
 
 Functions for turbulence, sub-grid scale modelling. These include
-viscosity terms, diffusivity and stress tensors. 
+viscosity terms, diffusivity and stress tensors.
 
 - [`ConstantViscosityWithDivergence`](@ref)
 - [`SmagorinskyLilly`](@ref)
@@ -47,12 +47,7 @@ using CLIMAParameters.Atmos.SubgridScale: inv_Pr_turb
 
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_gradient_laplacian,
-    vars_state_gradient_flux,
-    vars_hyperdiffusive,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -62,10 +57,6 @@ import ClimateMachine.BalanceLaws:
     init_state_conservative!,
     update_auxiliary_state!,
     nodal_update_auxiliary_state!,
-    number_state_conservative,
-    num_integrals,
-    vars_integrals,
-    vars_reverse_integrals,
     indefinite_stack_integral!,
     reverse_indefinite_stack_integral!,
     integral_load_auxiliary_state!,
@@ -90,18 +81,12 @@ export TurbulenceClosureModel,
 
 
 """
-    Abstract type with default do-nothing behaviour for 
-arbitrary turbulence closure models. 
+    Abstract type with default do-nothing behaviour for
+arbitrary turbulence closure models.
 """
 abstract type TurbulenceClosureModel end
 
-vars_state_conservative(::TurbulenceClosureModel, FT) = @vars()
-vars_state_auxiliary(::TurbulenceClosureModel, FT) = @vars()
-vars_state_gradient(::TurbulenceClosureModel, FT) = @vars()
-vars_state_gradient_flux(::TurbulenceClosureModel, FT) = @vars()
-vars_gradient_laplacian(::TurbulenceClosureModel, FT) = @vars()
-vars_integrals(::TurbulenceClosureModel, FT) = @vars()
-vars_reverse_integrals(::TurbulenceClosureModel, FT) = @vars()
+vars_state(::TurbulenceClosureModel, ::AbstractStateType, FT) = @vars()
 
 """
     init_aux_turbulence!
@@ -174,10 +159,10 @@ for the returned quantities.
 
 # Arguments
 
-- `::TurbulenceClosureModel` = Struct identifier for turbulence closure model
-- `orientation` = `BalanceLaw.orientation`
-- `param_set` = `BalanceLaw.param_set`
-- `state` = Array of prognostic (state) variables. See `vars_state_conservative` in `BalanceLaw`
+- `::TurbulenceClosure` = Struct identifier for turbulence closure model
+- `orientation` = `AtmosModel.orientation`
+- `param_set` = `AtmosModel.param_set`
+- `state` = Array of prognostic (state) variables. See `vars_state` in `AtmosModel`
 - `diffusive` = Array of diffusive variables
 - `aux` = Array of auxiliary variables
 - `t` = time
@@ -299,8 +284,8 @@ struct ConstantViscosityWithDivergence{FT} <: TurbulenceClosureModel
     ρν::FT
 end
 
-vars_state_gradient(::ConstantViscosityWithDivergence, FT) = @vars()
-vars_state_gradient_flux(::ConstantViscosityWithDivergence, FT) =
+vars_state(::ConstantViscosityWithDivergence, ::Gradient, FT) = @vars()
+vars_state(::ConstantViscosityWithDivergence, ::GradientFlux, FT) =
     @vars(S::SHermitianCompact{3, FT, 6})
 
 function compute_gradient_flux!(
@@ -415,9 +400,9 @@ struct SmagorinskyLilly{FT} <: TurbulenceClosureModel
     C_smag::FT
 end
 
-vars_state_auxiliary(::SmagorinskyLilly, FT) = @vars(Δ::FT)
-vars_state_gradient(::SmagorinskyLilly, FT) = @vars(θ_v::FT)
-vars_state_gradient_flux(::SmagorinskyLilly, FT) =
+vars_state(::SmagorinskyLilly, ::Auxiliary, FT) = @vars(Δ::FT)
+vars_state(::SmagorinskyLilly, ::Gradient, FT) = @vars(θ_v::FT)
+vars_state(::SmagorinskyLilly, ::GradientFlux, FT) =
     @vars(S::SHermitianCompact{3, FT, 6}, N²::FT)
 
 
@@ -540,9 +525,10 @@ struct Vreman{FT} <: TurbulenceClosureModel
     "Smagorinsky Coefficient [dimensionless]"
     C_smag::FT
 end
-vars_state_auxiliary(::Vreman, FT) = @vars(Δ::FT)
-vars_state_gradient(::Vreman, FT) = @vars(θ_v::FT)
-vars_state_gradient_flux(::Vreman, FT) = @vars(∇u::SMatrix{3, 3, FT, 9}, N²::FT)
+vars_state(::Vreman, ::Auxiliary, FT) = @vars(Δ::FT)
+vars_state(::Vreman, ::Gradient, FT) = @vars(θ_v::FT)
+vars_state(::Vreman, ::GradientFlux, FT) =
+    @vars(∇u::SMatrix{3, 3, FT, 9}, N²::FT)
 
 function init_aux_turbulence!(
     ::Vreman,
@@ -646,9 +632,9 @@ $(DocStringExtensions.FIELDS)
 struct AnisoMinDiss{FT} <: TurbulenceClosureModel
     C_poincare::FT
 end
-vars_state_auxiliary(::AnisoMinDiss, FT) = @vars(Δ::FT)
-vars_state_gradient(::AnisoMinDiss, FT) = @vars(θ_v::FT)
-vars_state_gradient_flux(::AnisoMinDiss, FT) =
+vars_state(::AnisoMinDiss, ::Auxiliary, FT) = @vars(Δ::FT)
+vars_state(::AnisoMinDiss, ::Gradient, FT) = @vars(θ_v::FT)
+vars_state(::AnisoMinDiss, ::GradientFlux, FT) =
     @vars(∇u::SMatrix{3, 3, FT, 9}, N²::FT)
 function init_aux_turbulence!(
     ::AnisoMinDiss,
diff --git a/src/Diagnostics/Diagnostics.jl b/src/Diagnostics/Diagnostics.jl
index 9e25a1cee8b..e5b7b75fb3a 100644
--- a/src/Diagnostics/Diagnostics.jl
+++ b/src/Diagnostics/Diagnostics.jl
@@ -24,13 +24,7 @@ import KernelAbstractions: CPU
 
 using ..ConfigTypes
 using ..DGMethods
-using ..BalanceLaws:
-    number_state_conservative,
-    vars_state_conservative,
-    number_state_auxiliary,
-    vars_state_auxiliary,
-    vars_state_gradient_flux,
-    number_state_gradient_flux
+using ..BalanceLaws
 using ..Mesh.Interpolation
 using ..MPIStateArrays
 using ..VariableTemplates
diff --git a/src/Diagnostics/atmos_gcm_default.jl b/src/Diagnostics/atmos_gcm_default.jl
index 02c3d697813..371cc353852 100644
--- a/src/Diagnostics/atmos_gcm_default.jl
+++ b/src/Diagnostics/atmos_gcm_default.jl
@@ -235,8 +235,8 @@ function atmos_gcm_default_collect(dgngrp::DiagnosticsGroup, currtime)
     # Compute thermo variables
     thermo_array = Array{FT}(undef, npoints, num_thermo(atmos, FT), nrealelem)
     @visitQ nhorzelem nvertelem Nqk Nq begin
-        state = extract_state_conservative(dg, state_data, ijk, e)
-        aux = extract_state_auxiliary(dg, aux_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
+        aux = extract_state(dg, aux_data, ijk, e, Auxiliary())
 
         thermo = thermo_vars(atmos, view(thermo_array, ijk, :, e))
         compute_thermo!(atmos, state, aux, thermo)
@@ -244,8 +244,11 @@ function atmos_gcm_default_collect(dgngrp::DiagnosticsGroup, currtime)
 
     # Interpolate the state, thermo and dyn vars to sphere (u and vorticity
     # need projection to zonal, merid). All this may happen on the GPU.
-    istate =
-        ArrayType{FT}(undef, interpol.Npl, number_state_conservative(atmos, FT))
+    istate = ArrayType{FT}(
+        undef,
+        interpol.Npl,
+        number_states(atmos, Conservative(), FT),
+    )
     interpolate_local!(interpol, Q.realdata, istate)
 
     ithermo = ArrayType{FT}(undef, interpol.Npl, num_thermo(atmos, FT))
@@ -283,7 +286,7 @@ function atmos_gcm_default_collect(dgngrp::DiagnosticsGroup, currtime)
         )
 
         @visitI nlong nlat nlevel begin
-            statei = Vars{vars_state_conservative(atmos, FT)}(view(
+            statei = Vars{vars_state(atmos, Conservative(), FT)}(view(
                 all_state_data,
                 lo,
                 la,
diff --git a/src/Diagnostics/atmos_les_core.jl b/src/Diagnostics/atmos_les_core.jl
index 6b76d81ba2e..9defa21175d 100644
--- a/src/Diagnostics/atmos_les_core.jl
+++ b/src/Diagnostics/atmos_les_core.jl
@@ -210,8 +210,8 @@ function atmos_les_core_collect(dgngrp::DiagnosticsGroup, currtime)
     @visitQ nhorzelem nvertelem Nqk Nq begin
         evk = Nqk * (ev - 1) + k
 
-        state = extract_state_conservative(dg, state_data, ijk, e)
-        aux = extract_state_auxiliary(dg, aux_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
+        aux = extract_state(dg, aux_data, ijk, e, Auxiliary())
         MH = vgeo[ijk, grid.MHid, e]
 
         thermo = thermo_vars(bl, thermo_array[ijk, e])
@@ -266,7 +266,7 @@ function atmos_les_core_collect(dgngrp::DiagnosticsGroup, currtime)
     @visitQ nhorzelem nvertelem Nqk Nq begin
         evk = Nqk * (ev - 1) + k
 
-        state = extract_state_conservative(dg, state_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
         thermo = thermo_vars(bl, thermo_array[ijk, e])
         MH = vgeo[ijk, grid.MHid, e]
 
diff --git a/src/Diagnostics/atmos_les_default.jl b/src/Diagnostics/atmos_les_default.jl
index f2761d23b17..63fcddcaffd 100644
--- a/src/Diagnostics/atmos_les_default.jl
+++ b/src/Diagnostics/atmos_les_default.jl
@@ -344,9 +344,9 @@ function atmos_les_default_collect(dgngrp::DiagnosticsGroup, currtime)
     @visitQ nhorzelem nvertelem Nqk Nq begin
         evk = Nqk * (ev - 1) + k
 
-        state = extract_state_conservative(dg, state_data, ijk, e)
-        gradflux = extract_state_gradient_flux(dg, gradflux_data, ijk, e)
-        aux = extract_state_auxiliary(dg, aux_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
+        gradflux = extract_state(dg, gradflux_data, ijk, e, GradientFlux())
+        aux = extract_state(dg, aux_data, ijk, e, Auxiliary())
         MH = vgeo[ijk, grid.MHid, e]
 
         thermo = thermo_vars(bl, thermo_array[ijk, e])
@@ -438,7 +438,7 @@ function atmos_les_default_collect(dgngrp::DiagnosticsGroup, currtime)
     @visitQ nhorzelem nvertelem Nqk Nq begin
         evk = Nqk * (ev - 1) + k
 
-        state = extract_state_conservative(dg, state_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
         thermo = thermo_vars(bl, thermo_array[ijk, e])
         MH = vgeo[ijk, grid.MHid, e]
 
diff --git a/src/Diagnostics/atmos_refstate_perturbations.jl b/src/Diagnostics/atmos_refstate_perturbations.jl
index b1db9b49068..0bbdd3e210a 100644
--- a/src/Diagnostics/atmos_refstate_perturbations.jl
+++ b/src/Diagnostics/atmos_refstate_perturbations.jl
@@ -172,8 +172,8 @@ function atmos_refstate_perturbations_collect(
     # Compute thermo variables
     thermo_array = Array{FT}(undef, npoints, num_thermo(atmos, FT), nrealelem)
     @visitQ nhorzelem nvertelem Nqk Nq begin
-        state = extract_state_conservative(dg, state_data, ijk, e)
-        aux = extract_state_auxiliary(dg, aux_data, ijk, e)
+        state = extract_state(dg, state_data, ijk, e, Conservative())
+        aux = extract_state(dg, aux_data, ijk, e, Auxiliary())
 
         thermo = thermo_vars(atmos, view(thermo_array, ijk, :, e))
         compute_thermo!(atmos, state, aux, thermo)
@@ -181,8 +181,11 @@ function atmos_refstate_perturbations_collect(
 
     # Interpolate the state and thermo variables.
     interpol = dgngrp.interpol
-    istate =
-        ArrayType{FT}(undef, interpol.Npl, number_state_conservative(atmos, FT))
+    istate = ArrayType{FT}(
+        undef,
+        interpol.Npl,
+        number_states(atmos, Conservative(), FT),
+    )
     interpolate_local!(interpol, Q.realdata, istate)
 
     if interpol isa InterpolationCubedSphere
@@ -191,7 +194,11 @@ function atmos_refstate_perturbations_collect(
         project_cubed_sphere!(interpol, istate, (_ρu, _ρv, _ρw))
     end
 
-    iaux = ArrayType{FT}(undef, interpol.Npl, number_state_auxiliary(atmos, FT))
+    iaux = ArrayType{FT}(
+        undef,
+        interpol.Npl,
+        number_states(atmos, Auxiliary(), FT),
+    )
     interpolate_local!(interpol, dg.state_auxiliary.realdata, iaux)
 
     ithermo = ArrayType{FT}(undef, interpol.Npl, num_thermo(atmos, FT))
@@ -220,14 +227,14 @@ function atmos_refstate_perturbations_collect(
         )
 
         @visitI nlong nlat nlevel begin
-            statei = Vars{vars_state_conservative(atmos, FT)}(view(
+            statei = Vars{vars_state(atmos, Conservative(), FT)}(view(
                 all_state_data,
                 lo,
                 la,
                 le,
                 :,
             ))
-            auxi = Vars{vars_state_auxiliary(atmos, FT)}(view(
+            auxi = Vars{vars_state(atmos, Auxiliary(), FT)}(view(
                 all_aux_data,
                 lo,
                 la,
diff --git a/src/Diagnostics/diagnostic_fields.jl b/src/Diagnostics/diagnostic_fields.jl
index 672b1456644..691cde45942 100644
--- a/src/Diagnostics/diagnostic_fields.jl
+++ b/src/Diagnostics/diagnostic_fields.jl
@@ -101,8 +101,8 @@ function VectorGradients(dg::DGModel, Q::MPIStateArray)
     data = similar(Q.realdata, npoints, 9, nrealelem)
 
     ind = [
-        varsindex(vars_state_conservative(bl, FT), :ρ)
-        varsindex(vars_state_conservative(bl, FT), :ρu)
+        varsindex(vars_state(bl, Conservative(), FT), :ρ)
+        varsindex(vars_state(bl, Conservative(), FT), :ρu)
     ]
     _ρ, _ρu, _ρv, _ρw = ind[1], ind[2], ind[3], ind[4]
 
diff --git a/src/Diagnostics/dump_aux.jl b/src/Diagnostics/dump_aux.jl
index d41228fd238..acd6343e2f1 100644
--- a/src/Diagnostics/dump_aux.jl
+++ b/src/Diagnostics/dump_aux.jl
@@ -1,32 +1,4 @@
-function dump_aux_init(dgngrp, currtime)
-    FT = eltype(Settings.Q)
-    bl = Settings.dg.balance_law
-    mpicomm = Settings.mpicomm
-    mpirank = MPI.Comm_rank(mpicomm)
-
-    if mpirank == 0
-        # get dimensions for the interpolated grid
-        dims = dimensions(dgngrp.interpol)
-
-        # set up the variables we're going to be writing
-        vars = OrderedDict()
-        auxnames = flattenednames(vars_state_auxiliary(bl, FT))
-        for varname in auxnames
-            vars[varname] = (tuple(collect(keys(dims))...), FT, Dict())
-        end
-
-        dprefix = @sprintf(
-            "%s_%s-%s",
-            dgngrp.out_prefix,
-            dgngrp.name,
-            Settings.starttime,
-        )
-        dfilename = joinpath(Settings.output_dir, dprefix)
-        init_data(dgngrp.writer, dfilename, dims, vars)
-    end
-
-    return nothing
-end
+dump_aux_init(dgngrp, currtime) = dump_init(dgngrp, currtime, Auxiliary())
 
 function dump_aux_collect(dgngrp, currtime)
     interpol = dgngrp.interpol
@@ -40,14 +12,17 @@ function dump_aux_collect(dgngrp, currtime)
     iaux = similar(
         dg.state_auxiliary.data,
         interpol.Npl,
-        number_state_auxiliary(bl, FT),
+        number_states(bl, Auxiliary(), FT),
     )
+
+
     interpolate_local!(interpol, dg.state_auxiliary.data, iaux)
 
+
     all_aux_data = accumulate_interpolated_data(mpicomm, interpol, iaux)
 
     if mpirank == 0
-        auxnames = flattenednames(vars_state_auxiliary(bl, FT))
+        auxnames = flattenednames(vars_state(bl, Auxiliary(), FT))
         varvals = OrderedDict()
         for (vari, varname) in enumerate(auxnames)
             varvals[varname] = all_aux_data[:, :, :, vari]
diff --git a/src/Diagnostics/dump_init.jl b/src/Diagnostics/dump_init.jl
new file mode 100644
index 00000000000..d8d1b581060
--- /dev/null
+++ b/src/Diagnostics/dump_init.jl
@@ -0,0 +1,29 @@
+function dump_init(dgngrp, currtime, st::AbstractStateType)
+    FT = eltype(Settings.Q)
+    bl = Settings.dg.balance_law
+    mpicomm = Settings.mpicomm
+    mpirank = MPI.Comm_rank(mpicomm)
+
+    if mpirank == 0
+        # get dimensions for the interpolated grid
+        dims = dimensions(dgngrp.interpol)
+
+        # set up the variables we're going to be writing
+        vars = OrderedDict()
+        statenames = flattenednames(vars_state(bl, st, FT))
+        for varname in statenames
+            vars[varname] = (tuple(collect(keys(dims))...), FT, Dict())
+        end
+
+        dprefix = @sprintf(
+            "%s_%s-%s",
+            dgngrp.out_prefix,
+            dgngrp.name,
+            Settings.starttime,
+        )
+        dfilename = joinpath(Settings.output_dir, dprefix)
+        init_data(dgngrp.writer, dfilename, dims, vars)
+    end
+
+    return nothing
+end
diff --git a/src/Diagnostics/dump_state.jl b/src/Diagnostics/dump_state.jl
index fb5deadf56d..aeec6a51447 100644
--- a/src/Diagnostics/dump_state.jl
+++ b/src/Diagnostics/dump_state.jl
@@ -1,32 +1,4 @@
-function dump_state_init(dgngrp, currtime)
-    FT = eltype(Settings.Q)
-    bl = Settings.dg.balance_law
-    mpicomm = Settings.mpicomm
-    mpirank = MPI.Comm_rank(mpicomm)
-
-    if mpirank == 0
-        # get dimensions for the interpolated grid
-        dims = dimensions(dgngrp.interpol)
-
-        # set up the variables we're going to be writing
-        vars = OrderedDict()
-        statenames = flattenednames(vars_state_conservative(bl, FT))
-        for varname in statenames
-            vars[varname] = (tuple(collect(keys(dims))...), FT, Dict())
-        end
-
-        dprefix = @sprintf(
-            "%s_%s-%s",
-            dgngrp.out_prefix,
-            dgngrp.name,
-            Settings.starttime,
-        )
-        dfilename = joinpath(Settings.output_dir, dprefix)
-        init_data(dgngrp.writer, dfilename, dims, vars)
-    end
-
-    return nothing
-end
+dump_state_init(dgngrp, currtime) = dump_init(dgngrp, currtime, Conservative())
 
 function dump_state_collect(dgngrp, currtime)
     interpol = dgngrp.interpol
@@ -37,7 +9,8 @@ function dump_state_collect(dgngrp, currtime)
     bl = dg.balance_law
     mpirank = MPI.Comm_rank(mpicomm)
 
-    istate = similar(Q.data, interpol.Npl, number_state_conservative(bl, FT))
+    istate =
+        similar(Q.data, interpol.Npl, number_states(bl, Conservative(), FT))
     interpolate_local!(interpol, Q.data, istate)
 
     if interpol isa InterpolationCubedSphere
@@ -49,7 +22,7 @@ function dump_state_collect(dgngrp, currtime)
     all_state_data = accumulate_interpolated_data(mpicomm, interpol, istate)
 
     if mpirank == 0
-        statenames = flattenednames(vars_state_conservative(bl, FT))
+        statenames = flattenednames(vars_state(bl, Conservative(), FT))
         varvals = OrderedDict()
         for (vari, varname) in enumerate(statenames)
             varvals[varname] = all_state_data[:, :, :, vari]
diff --git a/src/Diagnostics/groups.jl b/src/Diagnostics/groups.jl
index 63aaee0777c..44363afffef 100644
--- a/src/Diagnostics/groups.jl
+++ b/src/Diagnostics/groups.jl
@@ -211,6 +211,7 @@ function setup_atmos_refstate_perturbations(
     )
 end
 
+include("dump_init.jl")
 include("dump_state.jl")
 """
     setup_dump_state_diagnostics(
diff --git a/src/Diagnostics/helpers.jl b/src/Diagnostics/helpers.jl
index 9ed31adad2d..3221672c138 100644
--- a/src/Diagnostics/helpers.jl
+++ b/src/Diagnostics/helpers.jl
@@ -38,34 +38,13 @@ macro visitI(nlong, nlat, nlevel, expr)
 end
 
 # Helpers to extract data from the various state arrays
-function extract_state_conservative(dg, state_conservative, ijk, e)
+function extract_state(dg, state, ijk, e, st::AbstractStateType)
     bl = dg.balance_law
-    FT = eltype(state_conservative)
-    num_state_conservative = number_state_conservative(bl, FT)
-    local_state_conservative = MArray{Tuple{num_state_conservative}, FT}(undef)
-    for s in 1:num_state_conservative
-        local_state_conservative[s] = state_conservative[ijk, s, e]
+    FT = eltype(state)
+    num_state = number_states(bl, st, FT)
+    local_state = MArray{Tuple{num_state}, FT}(undef)
+    for s in 1:num_state
+        local_state[s] = state[ijk, s, e]
     end
-    return Vars{vars_state_conservative(bl, FT)}(local_state_conservative)
-end
-function extract_state_auxiliary(dg, state_auxiliary, ijk, e)
-    bl = dg.balance_law
-    FT = eltype(state_auxiliary)
-    num_state_auxiliary = number_state_auxiliary(bl, FT)
-    local_state_auxiliary = MArray{Tuple{num_state_auxiliary}, FT}(undef)
-    for s in 1:num_state_auxiliary
-        local_state_auxiliary[s] = state_auxiliary[ijk, s, e]
-    end
-    return Vars{vars_state_auxiliary(bl, FT)}(local_state_auxiliary)
-end
-function extract_state_gradient_flux(dg, state_gradient_flux, ijk, e)
-    bl = dg.balance_law
-    FT = eltype(state_gradient_flux)
-    num_state_gradient_flux = number_state_gradient_flux(bl, FT)
-    local_state_gradient_flux =
-        MArray{Tuple{num_state_gradient_flux}, FT}(undef)
-    for s in 1:num_state_gradient_flux
-        local_state_gradient_flux[s] = state_gradient_flux[ijk, s, e]
-    end
-    return Vars{vars_state_gradient_flux(bl, FT)}(local_state_gradient_flux)
+    return Vars{vars_state(bl, st, FT)}(local_state)
 end
diff --git a/src/Driver/Driver.jl b/src/Driver/Driver.jl
index c4713e7dab3..bf70920f938 100644
--- a/src/Driver/Driver.jl
+++ b/src/Driver/Driver.jl
@@ -17,11 +17,7 @@ using ..SystemSolvers
 using ..ConfigTypes
 using ..Diagnostics
 using ..DGMethods
-using ..BalanceLaws:
-    BalanceLaw,
-    vars_state_conservative,
-    vars_state_auxiliary,
-    update_auxiliary_state!
+using ..BalanceLaws: BalanceLaw, vars_state, update_auxiliary_state!
 using ..DGMethods: remainder_DGModel
 using ..DGMethods.NumericalFluxes
 using ..HydrostaticBoussinesq
diff --git a/src/Numerics/DGMethods/DGMethods.jl b/src/Numerics/DGMethods/DGMethods.jl
index 5362aee1971..136b7ee56b1 100644
--- a/src/Numerics/DGMethods/DGMethods.jl
+++ b/src/Numerics/DGMethods/DGMethods.jl
@@ -12,23 +12,18 @@ using ..VariableTemplates
 using ..Courant
 using ..BalanceLaws:
     BalanceLaw,
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_state_gradient,
-    vars_state_gradient_flux,
-    vars_gradient_laplacian,
-    vars_hyperdiffusive,
-    vars_integrals,
-    vars_reverse_integrals,
-    number_state_conservative,
-    number_state_gradient_flux,
-    number_state_conservative,
-    number_state_auxiliary,
-    number_state_gradient,
-    num_gradient_laplacian,
-    num_hyperdiffusive,
-    num_integrals,
-    num_reverse_integrals
+    AbstractStateType,
+    Conservative,
+    Auxiliary,
+    Gradient,
+    GradientFlux,
+    GradientLaplacian,
+    Hyperdiffusive,
+    UpwardIntegrals,
+    DownwardIntegrals,
+    vars_state,
+    rank_multiplier,
+    number_states
 
 import ..BalanceLaws:
     BalanceLaw,
diff --git a/src/Numerics/DGMethods/DGModel.jl b/src/Numerics/DGMethods/DGModel.jl
index ea276cfb7b0..c8446789b6b 100644
--- a/src/Numerics/DGMethods/DGModel.jl
+++ b/src/Numerics/DGMethods/DGModel.jl
@@ -21,13 +21,18 @@ function DGModel(
     numerical_flux_first_order,
     numerical_flux_second_order,
     numerical_flux_gradient;
-    state_auxiliary = create_auxiliary_state(balance_law, grid),
-    state_gradient_flux = create_gradient_state(balance_law, grid),
-    states_higher_order = create_higher_order_states(balance_law, grid),
+    state_auxiliary = create_state(balance_law, grid, Auxiliary()),
+    state_gradient_flux = create_state(balance_law, grid, GradientFlux()),
+    states_higher_order = (
+        create_state(balance_law, grid, GradientLaplacian()),
+        create_state(balance_law, grid, Hyperdiffusive()),
+    ),
     direction = EveryDirection(),
     diffusion_direction = direction,
     modeldata = nothing,
 )
+    state_auxiliary =
+        init_state(state_auxiliary, balance_law, grid, Auxiliary())
     DGModel(
         balance_law,
         grid,
@@ -72,9 +77,9 @@ function (dg::DGModel)(
     state_auxiliary = dg.state_auxiliary
 
     FT = eltype(state_conservative)
-    num_state_conservative = number_state_conservative(balance_law, FT)
-    num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-    nhyperviscstate = num_hyperdiffusive(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
+    num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+    nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
     num_state_tendency = size(tendency, 2)
 
     @assert num_state_conservative ≤ num_state_tendency
@@ -528,7 +533,7 @@ function init_ode_state(dg::DGModel, args...; init_on_cpu = false)
     balance_law = dg.balance_law
     grid = dg.grid
 
-    state_conservative = create_conservative_state(balance_law, grid)
+    state_conservative = create_state(balance_law, grid, Conservative())
 
     topology = grid.topology
     Np = dofs_per_element(grid)
@@ -590,7 +595,7 @@ function restart_ode_state(dg::DGModel, state_data; init_on_cpu = false)
     bl = dg.balance_law
     grid = dg.grid
 
-    state = create_conservative_state(bl, grid)
+    state = create_state(bl, grid, Conservative())
     state .= state_data
 
     device = arraytype(dg.grid) <: Array ? CPU() : CUDADevice()
@@ -603,7 +608,8 @@ function restart_ode_state(dg::DGModel, state_data; init_on_cpu = false)
 end
 
 function restart_auxiliary_state(bl, grid, aux_data)
-    state_auxiliary = create_auxiliary_state(bl, grid)
+    state_auxiliary = create_state(bl, grid, Auxiliary())
+    state_auxiliary = init_state(state_auxiliary, bl, grid, Auxiliary())
     state_auxiliary .= aux_data
     return state_auxiliary
 end
@@ -883,7 +889,7 @@ function MPIStateArrays.MPIStateArray(dg::DGModel)
     balance_law = dg.balance_law
     grid = dg.grid
 
-    state_conservative = create_conservative_state(balance_law, grid)
+    state_conservative = create_state(balance_law, grid, Conservative())
 
     return state_conservative
 end
@@ -896,8 +902,8 @@ function create_hypervisc_indexmap(balance_law::BalanceLaw)
         collect(Iterators.Flatten(_getvars.(fields, fieldtypes(T))))
     end
 
-    gradvars = vars_state_gradient(balance_law, Int)
-    gradlapvars = vars_gradient_laplacian(balance_law, Int)
+    gradvars = vars_state(balance_law, Gradient(), Int)
+    gradlapvars = vars_state(balance_law, GradientLaplacian(), Int)
     indices = Vars{gradvars}(1:varsize(gradvars))
     SVector{varsize(gradlapvars)}(_getvars(indices, gradlapvars))
 end
diff --git a/src/Numerics/DGMethods/DGModel_kernels.jl b/src/Numerics/DGMethods/DGModel_kernels.jl
index 83b2e44c183..70e154fbd62 100644
--- a/src/Numerics/DGMethods/DGModel_kernels.jl
+++ b/src/Numerics/DGMethods/DGModel_kernels.jl
@@ -52,12 +52,12 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
     @uniform begin
         N = polyorder
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
 
         Nq = N + 1
 
@@ -132,11 +132,11 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
             fill!(local_flux, -zero(eltype(local_flux)))
             flux_first_order!(
                 balance_law,
-                Grad{vars_state_conservative(balance_law, FT)}(local_flux),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Grad{vars_state(balance_law, Conservative(), FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -152,17 +152,17 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
             fill!(local_flux, -zero(eltype(local_flux)))
             flux_second_order!(
                 balance_law,
-                Grad{vars_state_conservative(balance_law, FT)}(local_flux),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Grad{vars_state(balance_law, Conservative(), FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -198,13 +198,13 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
                     fill!(local_flux, -zero(eltype(local_flux)))
                     flux_first_order!(
                         balance_law,
-                        Grad{vars_state_conservative(balance_law, FT)}(
+                        Grad{vars_state(balance_law, Conservative(), FT)}(
                             local_flux,
                         ),
-                        Vars{vars_state_conservative(balance_law, FT)}(
+                        Vars{vars_state(balance_law, Conservative(), FT)}(
                             local_state_conservative,
                         ),
-                        Vars{vars_state_auxiliary(balance_law, FT)}(
+                        Vars{vars_state(balance_law, Auxiliary(), FT)}(
                             local_state_auxiliary,
                         ),
                         t,
@@ -225,13 +225,13 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
                     fill!(local_flux, -zero(eltype(local_flux)))
                     flux_first_order!(
                         balance_law,
-                        Grad{vars_state_conservative(balance_law, FT)}(
+                        Grad{vars_state(balance_law, Conservative(), FT)}(
                             local_flux,
                         ),
-                        Vars{vars_state_conservative(balance_law, FT)}(
+                        Vars{vars_state(balance_law, Conservative(), FT)}(
                             local_state_conservative,
                         ),
-                        Vars{vars_state_auxiliary(balance_law, FT)}(
+                        Vars{vars_state(balance_law, Auxiliary(), FT)}(
                             local_state_auxiliary,
                         ),
                         t,
@@ -263,14 +263,14 @@ Computational kernel: Evaluate the volume integrals on right-hand side of a
             fill!(local_source, -zero(eltype(local_source)))
             source!(
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_source),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_source),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -328,12 +328,12 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
 
         Nq = N + 1
 
@@ -408,11 +408,11 @@ end
             fill!(local_flux, -zero(eltype(local_flux)))
             flux_first_order!(
                 balance_law,
-                Grad{vars_state_conservative(balance_law, FT)}(local_flux),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Grad{vars_state(balance_law, Conservative(), FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -428,17 +428,17 @@ end
             fill!(local_flux, -zero(eltype(local_flux)))
             flux_second_order!(
                 balance_law,
-                Grad{vars_state_conservative(balance_law, FT)}(local_flux),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Grad{vars_state(balance_law, Conservative(), FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -476,14 +476,14 @@ end
             fill!(local_source, -zero(eltype(local_source)))
             source!(
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_source),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_source),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -549,11 +549,11 @@ Computational kernel: Evaluate the surface integrals on right-hand side of a
     @uniform begin
         N = polyorder
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
 
         if dim == 1
             Np = (N + 1)
@@ -681,18 +681,18 @@ Computational kernel: Evaluate the surface integrals on right-hand side of a
             numerical_flux_first_order!(
                 numerical_flux_first_order,
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_flux),
                 SVector(normal_vector),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺nondiff,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺nondiff,
                 ),
                 t,
@@ -701,30 +701,30 @@ Computational kernel: Evaluate the surface integrals on right-hand side of a
             numerical_flux_second_order!(
                 numerical_flux_second_order,
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_flux),
                 normal_vector,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux⁻,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺diff,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux⁺,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺diff,
                 ),
                 t,
@@ -748,68 +748,68 @@ Computational kernel: Evaluate the surface integrals on right-hand side of a
             numerical_boundary_flux_first_order!(
                 numerical_flux_first_order,
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_flux),
                 SVector(normal_vector),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺nondiff,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺nondiff,
                 ),
                 bctype,
                 t,
                 face_direction,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative_bottom1,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary_bottom1,
                 ),
             )
             numerical_boundary_flux_second_order!(
                 numerical_flux_second_order,
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(local_flux),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_flux),
                 normal_vector,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux⁻,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺diff,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux⁺,
                 ),
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺diff,
                 ),
                 bctype,
                 t,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative_bottom1,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux_bottom1,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary_bottom1,
                 ),
             )
@@ -844,11 +844,11 @@ end
         N = polyorder
 
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradstate = number_state_gradient(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradstate = number_states(balance_law, Gradient(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
         Nq = N + 1
 
@@ -895,12 +895,12 @@ end
             fill!(local_transform, -zero(eltype(local_transform)))
             compute_gradient_argument!(
                 balance_law,
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform),
-                Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[
                     :,
                     k,
                 ]),
-                Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[
                     :,
                     k,
                 ]),
@@ -975,19 +975,19 @@ end
                 )
                 compute_gradient_flux!(
                     balance_law,
-                    Vars{vars_state_gradient_flux(balance_law, FT)}(
+                    Vars{vars_state(balance_law, GradientFlux(), FT)}(
                         local_state_gradient_flux,
                     ),
-                    Grad{vars_state_gradient(balance_law, FT)}(local_transform_gradient[
+                    Grad{vars_state(balance_law, Gradient(), FT)}(local_transform_gradient[
                         :,
                         :,
                         k,
                     ]),
-                    Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[
+                    Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[
                         :,
                         k,
                     ]),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[
                         :,
                         k,
                     ]),
@@ -1022,11 +1022,11 @@ end
         N = polyorder
 
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradstate = number_state_gradient(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradstate = number_states(balance_law, Gradient(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
         Nq = N + 1
 
@@ -1087,12 +1087,12 @@ end
             fill!(local_transform, -zero(eltype(local_transform)))
             compute_gradient_argument!(
                 balance_law,
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform),
-                Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform),
+                Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[
                     :,
                     k,
                 ]),
-                Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[
                     :,
                     k,
                 ]),
@@ -1152,19 +1152,19 @@ end
                 )
                 compute_gradient_flux!(
                     balance_law,
-                    Vars{vars_state_gradient_flux(balance_law, FT)}(
+                    Vars{vars_state(balance_law, GradientFlux(), FT)}(
                         local_state_gradient_flux,
                     ),
-                    Grad{vars_state_gradient(balance_law, FT)}(local_transform_gradient[
+                    Grad{vars_state(balance_law, Gradient(), FT)}(local_transform_gradient[
                         :,
                         :,
                         k,
                     ]),
-                    Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[
+                    Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[
                         :,
                         k,
                     ]),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[
                         :,
                         k,
                     ]),
@@ -1202,11 +1202,11 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradstate = number_state_gradient(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradstate = number_states(balance_law, Gradient(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
         if dim == 1
             Np = (N + 1)
@@ -1288,11 +1288,13 @@ end
         fill!(local_transform⁻, -zero(eltype(local_transform⁻)))
         compute_gradient_argument!(
             balance_law,
-            Vars{vars_state_gradient(balance_law, FT)}(local_transform⁻),
-            Vars{vars_state_conservative(balance_law, FT)}(
+            Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁻),
+            Vars{vars_state(balance_law, Conservative(), FT)}(
                 local_state_conservative⁻,
             ),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary⁻),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary⁻,
+            ),
             t,
         )
 
@@ -1308,11 +1310,13 @@ end
         fill!(local_transform⁺, -zero(eltype(local_transform⁺)))
         compute_gradient_argument!(
             balance_law,
-            Vars{vars_state_gradient(balance_law, FT)}(local_transform⁺),
-            Vars{vars_state_conservative(balance_law, FT)}(
+            Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁺),
+            Vars{vars_state(balance_law, Conservative(), FT)}(
                 local_state_conservative⁺,
             ),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary⁺),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary⁺,
+            ),
             t,
         )
 
@@ -1327,18 +1331,18 @@ end
                 balance_law,
                 local_transform_gradient,
                 SVector(normal_vector),
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform⁻),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁻),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform⁺),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁺),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺,
                 ),
                 t,
@@ -1346,16 +1350,16 @@ end
             if num_state_gradient_flux > 0
                 compute_gradient_flux!(
                     balance_law,
-                    Vars{vars_state_gradient_flux(balance_law, FT)}(
+                    Vars{vars_state(balance_law, GradientFlux(), FT)}(
                         local_state_gradient_flux,
                     ),
-                    Grad{vars_state_gradient(balance_law, FT)}(
+                    Grad{vars_state(balance_law, Gradient(), FT)}(
                         local_transform_gradient,
                     ),
-                    Vars{vars_state_conservative(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Conservative(), FT)}(
                         local_state_conservative⁻,
                     ),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(
                         local_state_auxiliary⁻,
                     ),
                     t,
@@ -1378,42 +1382,42 @@ end
                 balance_law,
                 local_transform_gradient,
                 SVector(normal_vector),
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform⁻),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁻),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_state_gradient(balance_law, FT)}(local_transform⁺),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Gradient(), FT)}(local_transform⁺),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺,
                 ),
                 bctype,
                 t,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative_bottom1,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary_bottom1,
                 ),
             )
             if num_state_gradient_flux > 0
                 compute_gradient_flux!(
                     balance_law,
-                    Vars{vars_state_gradient_flux(balance_law, FT)}(
+                    Vars{vars_state(balance_law, GradientFlux(), FT)}(
                         local_state_gradient_flux,
                     ),
-                    Grad{vars_state_gradient(balance_law, FT)}(
+                    Grad{vars_state(balance_law, Gradient(), FT)}(
                         local_transform_gradient,
                     ),
-                    Vars{vars_state_conservative(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Conservative(), FT)}(
                         local_state_conservative⁻,
                     ),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(
                         local_state_auxiliary⁻,
                     ),
                     t,
@@ -1439,14 +1443,16 @@ end
 
         compute_gradient_flux!(
             balance_law,
-            Vars{vars_state_gradient_flux(balance_law, FT)}(
+            Vars{vars_state(balance_law, GradientFlux(), FT)}(
                 local_state_conservative⁻visc,
             ),
-            Grad{vars_state_gradient(balance_law, FT)}(l_nG⁻),
-            Vars{vars_state_conservative(balance_law, FT)}(
+            Grad{vars_state(balance_law, Gradient(), FT)}(l_nG⁻),
+            Vars{vars_state(balance_law, Conservative(), FT)}(
                 local_state_conservative⁻,
             ),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary⁻),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary⁻,
+            ),
             t,
         )
 
@@ -1477,8 +1483,8 @@ end
 ) where {dim, polyorder}
     N = polyorder
     FT = eltype(state_auxiliary)
-    num_state_auxiliary = number_state_auxiliary(balance_law, FT)
-    num_state_conservative = number_state_conservative(balance_law, FT)
+    num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
 
     Nq = N + 1
     Nqk = dim == 2 ? 1 : Nq
@@ -1501,8 +1507,10 @@ end
         end
         init_state_conservative!(
             balance_law,
-            Vars{vars_state_conservative(balance_law, FT)}(l_state),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary),
+            Vars{vars_state(balance_law, Conservative(), FT)}(l_state),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary,
+            ),
             coords,
             args...,
         )
@@ -1530,7 +1538,7 @@ See [`BalanceLaw`](@ref) for usage.
 ) where {dim, polyorder}
     N = polyorder
     FT = eltype(state_auxiliary)
-    num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+    num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
     Nq = N + 1
     Nqk = dim == 2 ? 1 : Nq
@@ -1549,7 +1557,9 @@ See [`BalanceLaw`](@ref) for usage.
 
         init_state_auxiliary!(
             balance_law,
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary,
+            ),
             LocalGeometry(Val(polyorder), vgeo, n, e),
         )
 
@@ -1577,8 +1587,8 @@ Update the auxiliary state array
     activedofs,
 ) where {dim, N}
     FT = eltype(state_conservative)
-    num_state_conservative = number_state_conservative(balance_law, FT)
-    num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
+    num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
     Nq = N + 1
 
@@ -1609,10 +1619,10 @@ Update the auxiliary state array
 
             f!(
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
                 t,
@@ -1638,9 +1648,9 @@ end
     activedofs,
 ) where {dim, N}
     FT = eltype(state_conservative)
-    num_state_conservative = number_state_conservative(balance_law, FT)
-    num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-    num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
+    num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+    num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
     Nq = N + 1
 
@@ -1677,13 +1687,13 @@ end
 
             f!(
                 balance_law,
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary,
                 ),
-                Vars{vars_state_gradient_flux(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientFlux(), FT)}(
                     local_state_gradient_flux,
                 ),
                 t,
@@ -1716,9 +1726,9 @@ See [`BalanceLaw`](@ref) for usage.
 ) where {dim, N, nvertelem}
     @uniform begin
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
-        nout = num_integrals(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
+        nout = number_states(balance_law, UpwardIntegrals(), FT)
 
         Nq = N + 1
         Nqj = dim == 2 ? 1 : Nq
@@ -1769,15 +1779,15 @@ See [`BalanceLaw`](@ref) for usage.
 
                 integral_load_auxiliary_state!(
                     balance_law,
-                    Vars{vars_integrals(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, UpwardIntegrals(), FT)}(view(
                         local_kernel,
                         :,
                         k,
                     )),
-                    Vars{vars_state_conservative(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Conservative(), FT)}(
                         local_state_conservative,
                     ),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(
                         local_state_auxiliary,
                     ),
                 )
@@ -1805,13 +1815,13 @@ See [`BalanceLaw`](@ref) for usage.
                 ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                 integral_set_auxiliary_state!(
                     balance_law,
-                    Vars{vars_state_auxiliary(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(view(
                         state_auxiliary,
                         ijk,
                         :,
                         e,
                     )),
-                    Vars{vars_integrals(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, UpwardIntegrals(), FT)}(view(
                         local_kernel,
                         :,
                         k,
@@ -1839,7 +1849,7 @@ end
 
         Nq = N + 1
         Nqj = dim == 2 ? 1 : Nq
-        nout = num_reverse_integrals(balance_law, FT)
+        nout = number_states(balance_law, DownwardIntegrals(), FT)
 
         # note that k is the second not 4th index (since this is scratch memory and k
         # needs to be persistent across threads)
@@ -1858,14 +1868,14 @@ end
         et = nvertelem + (eh - 1) * nvertelem
         reverse_integral_load_auxiliary_state!(
             balance_law,
-            Vars{vars_reverse_integrals(balance_law, FT)}(l_T),
-            Vars{vars_state_conservative(balance_law, FT)}(view(
+            Vars{vars_state(balance_law, DownwardIntegrals(), FT)}(l_T),
+            Vars{vars_state(balance_law, Conservative(), FT)}(view(
                 state,
                 ijk,
                 :,
                 et,
             )),
-            Vars{vars_state_auxiliary(balance_law, FT)}(view(
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(view(
                 state_auxiliary,
                 ijk,
                 :,
@@ -1880,14 +1890,14 @@ end
                 ijk = i + Nq * ((j - 1) + Nqj * (k - 1))
                 reverse_integral_load_auxiliary_state!(
                     balance_law,
-                    Vars{vars_reverse_integrals(balance_law, FT)}(l_V),
-                    Vars{vars_state_conservative(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, DownwardIntegrals(), FT)}(l_V),
+                    Vars{vars_state(balance_law, Conservative(), FT)}(view(
                         state,
                         ijk,
                         :,
                         e,
                     )),
-                    Vars{vars_state_auxiliary(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(view(
                         state_auxiliary,
                         ijk,
                         :,
@@ -1897,13 +1907,13 @@ end
                 l_V .= l_T .- l_V
                 reverse_integral_set_auxiliary_state!(
                     balance_law,
-                    Vars{vars_state_auxiliary(balance_law, FT)}(view(
+                    Vars{vars_state(balance_law, Auxiliary(), FT)}(view(
                         state_auxiliary,
                         ijk,
                         :,
                         e,
                     )),
-                    Vars{vars_reverse_integrals(balance_law, FT)}(l_V),
+                    Vars{vars_state(balance_law, DownwardIntegrals(), FT)}(l_V),
                 )
             end
         end
@@ -1962,7 +1972,7 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(Qhypervisc_grad)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
 
         Nq = N + 1
 
@@ -2053,7 +2063,7 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(Qhypervisc_grad)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
 
         Nq = N + 1
 
@@ -2134,7 +2144,7 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(Qhypervisc_grad)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
 
         if dim == 1
             Np = (N + 1)
@@ -2202,19 +2212,19 @@ end
             numerical_flux_divergence!(
                 divgradnumpenalty,
                 balance_law,
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_div),
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_div),
                 normal_vector,
-                Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad⁻),
-                Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad⁺),
+                Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad⁻),
+                Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad⁺),
             )
         else
             numerical_boundary_flux_divergence!(
                 divgradnumpenalty,
                 balance_law,
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_div),
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_div),
                 normal_vector,
-                Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad⁻),
-                Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad⁺),
+                Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad⁻),
+                Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad⁺),
                 bctype,
             )
         end
@@ -2246,10 +2256,10 @@ end
         N = polyorder
 
         FT = eltype(Qhypervisc_grad)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
         ngradtransformstate = num_state_conservative
 
         Nq = N + 1
@@ -2340,12 +2350,12 @@ end
         )
         transform_post_gradient_laplacian!(
             balance_law,
-            Vars{vars_hyperdiffusive(balance_law, FT)}(
+            Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                 local_state_hyperdiffusion,
             ),
-            Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad_lap),
-            Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[:]),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[:]),
+            Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad_lap),
+            Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[:]),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[:]),
             t,
         )
         @unroll for s in 1:nhyperviscstate
@@ -2374,10 +2384,10 @@ end
         N = polyorder
 
         FT = eltype(Qhypervisc_grad)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
         ngradtransformstate = num_state_conservative
 
         Nq = N + 1
@@ -2448,12 +2458,12 @@ end
         )
         transform_post_gradient_laplacian!(
             balance_law,
-            Vars{vars_hyperdiffusive(balance_law, FT)}(
+            Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                 local_state_hyperdiffusion,
             ),
-            Grad{vars_gradient_laplacian(balance_law, FT)}(l_grad_lap),
-            Vars{vars_state_conservative(balance_law, FT)}(local_state_conservative[:]),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary[:]),
+            Grad{vars_state(balance_law, GradientLaplacian(), FT)}(l_grad_lap),
+            Vars{vars_state(balance_law, Conservative(), FT)}(local_state_conservative[:]),
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(local_state_auxiliary[:]),
             t,
         )
         @unroll for s in 1:nhyperviscstate
@@ -2484,10 +2494,10 @@ end
     @uniform begin
         N = polyorder
         FT = eltype(Qhypervisc_grad)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        ngradlapstate = num_gradient_laplacian(balance_law, FT)
-        nhyperviscstate = num_hyperdiffusive(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        ngradlapstate = number_states(balance_law, GradientLaplacian(), FT)
+        nhyperviscstate = number_states(balance_law, Hyperdiffusive(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
         ngradtransformstate = num_state_conservative
 
         if dim == 1
@@ -2576,22 +2586,22 @@ end
             numerical_flux_higher_order!(
                 hyperviscnumflux,
                 balance_law,
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion,
                 ),
                 normal_vector,
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_lap⁻),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_lap⁻),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_lap⁺),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_lap⁺),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺,
                 ),
                 t,
@@ -2600,22 +2610,22 @@ end
             numerical_boundary_flux_higher_order!(
                 hyperviscnumflux,
                 balance_law,
-                Vars{vars_hyperdiffusive(balance_law, FT)}(
+                Vars{vars_state(balance_law, Hyperdiffusive(), FT)}(
                     local_state_hyperdiffusion,
                 ),
                 normal_vector,
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_lap⁻),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_lap⁻),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁻,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁻,
                 ),
-                Vars{vars_gradient_laplacian(balance_law, FT)}(l_lap⁺),
-                Vars{vars_state_conservative(balance_law, FT)}(
+                Vars{vars_state(balance_law, GradientLaplacian(), FT)}(l_lap⁺),
+                Vars{vars_state(balance_law, Conservative(), FT)}(
                     local_state_conservative⁺,
                 ),
-                Vars{vars_state_auxiliary(balance_law, FT)}(
+                Vars{vars_state(balance_law, Auxiliary(), FT)}(
                     local_state_auxiliary⁺,
                 ),
                 bctype,
@@ -2648,9 +2658,9 @@ end
 ) where {dim, N}
     @uniform begin
         FT = eltype(state_conservative)
-        num_state_conservative = number_state_conservative(balance_law, FT)
-        num_state_gradient_flux = number_state_gradient_flux(balance_law, FT)
-        num_state_auxiliary = number_state_auxiliary(balance_law, FT)
+        num_state_conservative = number_states(balance_law, Conservative(), FT)
+        num_state_gradient_flux = number_states(balance_law, GradientFlux(), FT)
+        num_state_auxiliary = number_states(balance_law, Auxiliary(), FT)
 
         Nq = N + 1
 
@@ -2685,11 +2695,13 @@ end
         Δx = pointwise_courant[n, e]
         c = local_courant(
             balance_law,
-            Vars{vars_state_conservative(balance_law, FT)}(
+            Vars{vars_state(balance_law, Conservative(), FT)}(
                 local_state_conservative,
             ),
-            Vars{vars_state_auxiliary(balance_law, FT)}(local_state_auxiliary),
-            Vars{vars_state_gradient_flux(balance_law, FT)}(
+            Vars{vars_state(balance_law, Auxiliary(), FT)}(
+                local_state_auxiliary,
+            ),
+            Vars{vars_state(balance_law, GradientFlux(), FT)}(
                 local_state_gradient_flux,
             ),
             Δx,
diff --git a/src/Numerics/DGMethods/NumericalFluxes.jl b/src/Numerics/DGMethods/NumericalFluxes.jl
index 8f2774af9ac..fdab66e60af 100644
--- a/src/Numerics/DGMethods/NumericalFluxes.jl
+++ b/src/Numerics/DGMethods/NumericalFluxes.jl
@@ -14,23 +14,15 @@ export NumericalFluxGradient,
 using StaticArrays, LinearAlgebra
 using ClimateMachine.VariableTemplates
 using KernelAbstractions.Extras: @unroll
+using ...BalanceLaws
 import ...BalanceLaws:
-    BalanceLaw,
-    vars_state_conservative,
-    vars_state_gradient_flux,
-    vars_state_auxiliary,
-    vars_state_gradient,
+    vars_state,
     boundary_state!,
     wavespeed,
     flux_first_order!,
     flux_second_order!,
     compute_gradient_flux!,
-    number_state_conservative,
-    number_state_gradient,
     compute_gradient_argument!,
-    num_gradient_laplacian,
-    vars_gradient_laplacian,
-    vars_hyperdiffusive,
     transform_post_gradient_laplacian!
 
 """
@@ -287,7 +279,7 @@ function numerical_flux_first_order!(
 ) where {S, A}
 
     FT = eltype(fluxᵀn)
-    num_state_conservative = number_state_conservative(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
     fluxᵀn = parent(fluxᵀn)
 
     flux⁻ = similar(fluxᵀn, Size(3, num_state_conservative))
@@ -373,7 +365,7 @@ function numerical_flux_second_order!(
 ) where {S, D, HD, A}
 
     FT = eltype(fluxᵀn)
-    num_state_conservative = number_state_conservative(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
     fluxᵀn = parent(fluxᵀn)
 
     flux⁻ = similar(fluxᵀn, Size(3, num_state_conservative))
@@ -572,7 +564,7 @@ function normal_boundary_flux_second_order!(
     aux1⁻,
 ) where {S}
     FT = eltype(fluxᵀn)
-    num_state_conservative = number_state_conservative(balance_law, FT)
+    num_state_conservative = number_states(balance_law, Conservative(), FT)
     fluxᵀn = parent(fluxᵀn)
 
     flux = similar(fluxᵀn, Size(3, num_state_conservative))
diff --git a/src/Numerics/DGMethods/create_states.jl b/src/Numerics/DGMethods/create_states.jl
index d2646ba235a..1d7107a9fd3 100644
--- a/src/Numerics/DGMethods/create_states.jl
+++ b/src/Numerics/DGMethods/create_states.jl
@@ -1,22 +1,24 @@
 #### Create states
 
-function create_conservative_state(balance_law::BalanceLaw, grid)
+function create_state(balance_law, grid, st::AbstractStateType)
     topology = grid.topology
-    # FIXME: Remove after updating CUDA
+    Np = dofs_per_element(grid)
+
     h_vgeo = Array(grid.vgeo)
     FT = eltype(h_vgeo)
-    Np = dofs_per_element(grid)
     DA = arraytype(grid)
 
     weights = view(h_vgeo, :, grid.Mid, :)
     weights = reshape(weights, size(weights, 1), 1, size(weights, 2))
 
-    V = vars_state_conservative(balance_law, FT)
-    state_conservative = MPIStateArray{FT, V}(
+    # TODO: Clean up this MPIStateArray interface...
+    ns = number_states(balance_law, st, FT)
+    V = vars_state(balance_law, st, FT)
+    state = MPIStateArray{FT, V}(
         topology.mpicomm,
         DA,
         Np,
-        number_state_conservative(balance_law, FT),
+        rank_multiplier(st) * ns,
         length(topology.elems),
         realelems = topology.realelems,
         ghostelems = topology.ghostelems,
@@ -27,142 +29,30 @@ function create_conservative_state(balance_law::BalanceLaw, grid)
         nabrtovmapsend = grid.nabrtovmapsend,
         weights = weights,
     )
-    return state_conservative
+    return state
 end
 
-function create_auxiliary_state(balance_law, grid)
+function init_state(state, balance_law, grid, ::Auxiliary)
     topology = grid.topology
     Np = dofs_per_element(grid)
-
-    h_vgeo = Array(grid.vgeo)
-    FT = eltype(h_vgeo)
-    DA = arraytype(grid)
-
-    weights = view(h_vgeo, :, grid.Mid, :)
-    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))
-
-    V = vars_state_auxiliary(balance_law, FT)
-    state_auxiliary = MPIStateArray{FT, V}(
-        topology.mpicomm,
-        DA,
-        Np,
-        number_state_auxiliary(balance_law, FT),
-        length(topology.elems),
-        realelems = topology.realelems,
-        ghostelems = topology.ghostelems,
-        vmaprecv = grid.vmaprecv,
-        vmapsend = grid.vmapsend,
-        nabrtorank = topology.nabrtorank,
-        nabrtovmaprecv = grid.nabrtovmaprecv,
-        nabrtovmapsend = grid.nabrtovmapsend,
-        weights = weights,
-    )
-
     dim = dimensionality(grid)
     polyorder = polynomialorder(grid)
     vgeo = grid.vgeo
-    device = array_device(state_auxiliary)
+    device = array_device(state)
     nrealelem = length(topology.realelems)
     event = Event(device)
     event = kernel_init_state_auxiliary!(device, min(Np, 1024), Np * nrealelem)(
         balance_law,
         Val(dim),
         Val(polyorder),
-        state_auxiliary.data,
+        state.data,
         vgeo,
         topology.realelems,
         dependencies = (event,),
     )
-    event = MPIStateArrays.begin_ghost_exchange!(
-        state_auxiliary;
-        dependencies = event,
-    )
-    event = MPIStateArrays.end_ghost_exchange!(
-        state_auxiliary;
-        dependencies = event,
-    )
+    event = MPIStateArrays.begin_ghost_exchange!(state; dependencies = event)
+    event = MPIStateArrays.end_ghost_exchange!(state; dependencies = event)
     wait(device, event)
 
-    return state_auxiliary
-end
-
-function create_gradient_state(balance_law, grid)
-    topology = grid.topology
-    Np = dofs_per_element(grid)
-
-    h_vgeo = Array(grid.vgeo)
-    FT = eltype(h_vgeo)
-    DA = arraytype(grid)
-
-    weights = view(h_vgeo, :, grid.Mid, :)
-    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))
-
-    # TODO: Clean up this MPIStateArray interface...
-    V = vars_state_gradient_flux(balance_law, FT)
-    state_gradient_flux = MPIStateArray{FT, V}(
-        topology.mpicomm,
-        DA,
-        Np,
-        number_state_gradient_flux(balance_law, FT),
-        length(topology.elems),
-        realelems = topology.realelems,
-        ghostelems = topology.ghostelems,
-        vmaprecv = grid.vmaprecv,
-        vmapsend = grid.vmapsend,
-        nabrtorank = topology.nabrtorank,
-        nabrtovmaprecv = grid.nabrtovmaprecv,
-        nabrtovmapsend = grid.nabrtovmapsend,
-        weights = weights,
-    )
-
-    return state_gradient_flux
-end
-
-function create_higher_order_states(balance_law, grid)
-    topology = grid.topology
-    Np = dofs_per_element(grid)
-
-    h_vgeo = Array(grid.vgeo)
-    FT = eltype(h_vgeo)
-    DA = arraytype(grid)
-
-    weights = view(h_vgeo, :, grid.Mid, :)
-    weights = reshape(weights, size(weights, 1), 1, size(weights, 2))
-
-    ngradlapstate = num_gradient_laplacian(balance_law, FT)
-    # TODO: Clean up this MPIStateArray interface...
-    V = vars_gradient_laplacian(balance_law, FT)
-    Qhypervisc_grad = MPIStateArray{FT, V}(
-        topology.mpicomm,
-        DA,
-        Np,
-        3ngradlapstate,
-        length(topology.elems),
-        realelems = topology.realelems,
-        ghostelems = topology.ghostelems,
-        vmaprecv = grid.vmaprecv,
-        vmapsend = grid.vmapsend,
-        nabrtorank = topology.nabrtorank,
-        nabrtovmaprecv = grid.nabrtovmaprecv,
-        nabrtovmapsend = grid.nabrtovmapsend,
-        weights = weights,
-    )
-
-    V = vars_hyperdiffusive(balance_law, FT)
-    Qhypervisc_div = MPIStateArray{FT, V}(
-        topology.mpicomm,
-        DA,
-        Np,
-        ngradlapstate,
-        length(topology.elems),
-        realelems = topology.realelems,
-        ghostelems = topology.ghostelems,
-        vmaprecv = grid.vmaprecv,
-        vmapsend = grid.vmapsend,
-        nabrtorank = topology.nabrtorank,
-        nabrtovmaprecv = grid.nabrtovmaprecv,
-        nabrtovmapsend = grid.nabrtovmapsend,
-        weights = weights,
-    )
-    return Qhypervisc_grad, Qhypervisc_div
+    return state
 end
diff --git a/src/Numerics/DGMethods/remainder.jl b/src/Numerics/DGMethods/remainder.jl
index f8d80ab019c..a61b004a9ce 100644
--- a/src/Numerics/DGMethods/remainder.jl
+++ b/src/Numerics/DGMethods/remainder.jl
@@ -2,14 +2,7 @@ using StaticNumbers
 export remainder_DGModel
 
 import ..BalanceLaws:
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_state_gradient,
-    vars_state_gradient_flux,
-    vars_integrals,
-    vars_reverse_integrals,
-    vars_gradient_laplacian,
-    vars_hyperdiffusive,
+    vars_state,
     init_state_conservative!,
     init_state_auxiliary!,
     flux_first_order!,
@@ -113,25 +106,25 @@ function remainder_DGModel(
     # If any of these asserts fail, the remainder model will need to be extended
     # to allow for it; see `flux_first_order!` and `source!` below.
     for subdg in subsdg
-        @assert number_state_conservative(subdg.balance_law, FT) <=
-                number_state_conservative(maindg.balance_law, FT)
+        @assert number_states(subdg.balance_law, Conservative(), FT) <=
+                number_states(maindg.balance_law, Conservative(), FT)
 
-        @assert number_state_auxiliary(subdg.balance_law, FT) ==
-                number_state_auxiliary(maindg.balance_law, FT)
+        @assert number_states(subdg.balance_law, Auxiliary(), FT) ==
+                number_states(maindg.balance_law, Auxiliary(), FT)
 
-        @assert number_state_gradient(subdg.balance_law, FT) == 0
-        @assert number_state_gradient_flux(subdg.balance_law, FT) == 0
+        @assert number_states(subdg.balance_law, Gradient(), FT) == 0
+        @assert number_states(subdg.balance_law, GradientFlux(), FT) == 0
 
-        @assert num_gradient_laplacian(subdg.balance_law, FT) == 0
-        @assert num_hyperdiffusive(subdg.balance_law, FT) == 0
+        @assert number_states(subdg.balance_law, GradientLaplacian(), FT) == 0
+        @assert number_states(subdg.balance_law, Hyperdiffusive(), FT) == 0
 
         # Do not currenlty support nested remainder models
         # For this to work the way directions and numerical fluxes are handled
         # would need to be updated.
         @assert !(subdg.balance_law isa RemBL)
 
-        @assert num_integrals(subdg.balance_law, FT) == 0
-        @assert num_reverse_integrals(subdg.balance_law, FT) == 0
+        @assert number_states(subdg.balance_law, UpwardIntegrals(), FT) == 0
+        @assert number_states(subdg.balance_law, DownwardIntegrals(), FT) == 0
 
         # The remainder model requires that the subcomponent direction be
         # included in the main model directions
@@ -164,29 +157,8 @@ function remainder_DGModel(
 end
 
 # Inherit most of the functionality from the main model
-vars_state_conservative(rem_balance_law::RemBL, FT) =
-    vars_state_conservative(rem_balance_law.main, FT)
-
-vars_state_gradient(rem_balance_law::RemBL, FT) =
-    vars_state_gradient(rem_balance_law.main, FT)
-
-vars_state_gradient_flux(rem_balance_law::RemBL, FT) =
-    vars_state_gradient_flux(rem_balance_law.main, FT)
-
-vars_state_auxiliary(rem_balance_law::RemBL, FT) =
-    vars_state_auxiliary(rem_balance_law.main, FT)
-
-vars_integrals(rem_balance_law::RemBL, FT) =
-    vars_integrals(rem_balance_law.main, FT)
-
-vars_reverse_integrals(rem_balance_law::RemBL, FT) =
-    vars_integrals(rem_balance_law.main, FT)
-
-vars_gradient_laplacian(rem_balance_law::RemBL, FT) =
-    vars_gradient_laplacian(rem_balance_law.main, FT)
-
-vars_hyperdiffusive(rem_balance_law::RemBL, FT) =
-    vars_hyperdiffusive(rem_balance_law.main, FT)
+vars_state(rem_balance_law::RemBL, st::AbstractStateType, FT) =
+    vars_state(rem_balance_law.main, st, FT)
 
 update_auxiliary_state!(dg::DGModel, rem_balance_law::RemBL, args...) =
     update_auxiliary_state!(dg, rem_balance_law.main, args...)
@@ -381,11 +353,11 @@ function wavespeed(
 
     ws = fill(
         -zero(FT),
-        MVector{number_state_conservative(rem_balance_law.main, FT), FT},
+        MVector{number_states(rem_balance_law.main, Conservative(), FT), FT},
     )
     rs = fill(
         -zero(FT),
-        MVector{number_state_conservative(rem_balance_law.main, FT), FT},
+        MVector{number_states(rem_balance_law.main, Conservative(), FT), FT},
     )
 
     # Compute the main components wavespeed
@@ -403,7 +375,7 @@ function wavespeed(
     # Compute the sub components wavespeed
     for (sub, subdir) in zip(rem_balance_law.subs, rem_balance_law.subsdir)
         @inbounds if subdir isa Union{Dirs.types...}
-            num_state = static(number_state_conservative(sub, Float32))
+            num_state = static(number_states(sub, Conservative(), Float32))
             rs[static(1):num_state] .+=
                 wavespeed(sub, nM, state, aux, t, (subdir,))
         end
@@ -488,7 +460,7 @@ function numerical_flux_first_order!(
             nf = numerical_fluxes[2][k]
 
             FT = eltype(a_fluxᵀn)
-            num_state_conservative = number_state_conservative(sub, FT)
+            num_state_conservative = number_states(sub, Conservative(), FT)
 
             a_sub_fluxᵀn = MVector{num_state_conservative, FT}(undef)
             a_sub_state_conservative⁻ =
@@ -496,7 +468,7 @@ function numerical_flux_first_order!(
             a_sub_state_conservative⁺ =
                 MVector{num_state_conservative, FT}(undef)
 
-            state_rng = static(1):static(number_state_conservative(sub, FT))
+            state_rng = static(1):static(number_states(sub, Conservative(), FT))
             a_sub_fluxᵀn .= a_fluxᵀn[state_rng]
             a_sub_state_conservative⁻ .= parent(state_conservative⁻)[state_rng]
             a_sub_state_conservative⁺ .= parent(state_conservative⁺)[state_rng]
@@ -506,13 +478,13 @@ function numerical_flux_first_order!(
             numerical_flux_first_order!(
                 nf,
                 sub,
-                Vars{vars_state_conservative(sub, FT)}(a_sub_fluxᵀn),
+                Vars{vars_state(sub, Conservative(), FT)}(a_sub_fluxᵀn),
                 normal_vector,
-                Vars{vars_state_conservative(sub, FT)}(
+                Vars{vars_state(sub, Conservative(), FT)}(
                     a_sub_state_conservative⁻,
                 ),
                 state_auxiliary⁻,
-                Vars{vars_state_conservative(sub, FT)}(
+                Vars{vars_state(sub, Conservative(), FT)}(
                     a_sub_state_conservative⁺,
                 ),
                 state_auxiliary⁺,
@@ -614,7 +586,7 @@ function numerical_boundary_flux_first_order!(
             a_state_auxiliary⁺ .= a_back_state_auxiliary⁺
 
             FT = eltype(a_fluxᵀn)
-            num_state_conservative = number_state_conservative(sub, FT)
+            num_state_conservative = number_states(sub, Conservative(), FT)
 
             a_sub_fluxᵀn = MVector{num_state_conservative, FT}(undef)
             a_sub_state_conservative⁻ =
@@ -624,7 +596,7 @@ function numerical_boundary_flux_first_order!(
             a_sub_state_conservative1⁻ =
                 MVector{num_state_conservative, FT}(undef)
 
-            state_rng = static(1):static(number_state_conservative(sub, FT))
+            state_rng = static(1):static(number_states(sub, Conservative(), FT))
             a_sub_fluxᵀn .= a_fluxᵀn[state_rng]
             a_sub_state_conservative⁻ .= parent(state_conservative⁻)[state_rng]
             a_sub_state_conservative⁺ .= parent(state_conservative⁺)[state_rng]
@@ -637,20 +609,20 @@ function numerical_boundary_flux_first_order!(
             numerical_boundary_flux_first_order!(
                 nf,
                 sub,
-                Vars{vars_state_conservative(sub, FT)}(a_sub_fluxᵀn),
+                Vars{vars_state(sub, Conservative(), FT)}(a_sub_fluxᵀn),
                 normal_vector,
-                Vars{vars_state_conservative(sub, FT)}(
+                Vars{vars_state(sub, Conservative(), FT)}(
                     a_sub_state_conservative⁻,
                 ),
                 state_auxiliary⁻,
-                Vars{vars_state_conservative(sub, FT)}(
+                Vars{vars_state(sub, Conservative(), FT)}(
                     a_sub_state_conservative⁺,
                 ),
                 state_auxiliary⁺,
                 bctype,
                 t,
                 (rem_balance_law.subsdir[k],),
-                Vars{vars_state_conservative(sub, FT)}(
+                Vars{vars_state(sub, Conservative(), FT)}(
                     a_sub_state_conservative1⁻,
                 ),
                 state_auxiliary1⁻,
diff --git a/src/Numerics/SystemSolvers/SystemSolvers.jl b/src/Numerics/SystemSolvers/SystemSolvers.jl
index 2e2a85d360b..011cd567d9c 100644
--- a/src/Numerics/SystemSolvers/SystemSolvers.jl
+++ b/src/Numerics/SystemSolvers/SystemSolvers.jl
@@ -8,8 +8,7 @@ import ..Mesh.Grids: polynomialorder, dimensionality
 using ..Mesh.Topologies
 using ..DGMethods
 using ..DGMethods: DGModel
-using ..BalanceLaws:
-    BalanceLaw, number_state_conservative, number_state_gradient_flux
+using ..BalanceLaws
 
 using Adapt
 using LinearAlgebra
diff --git a/src/Numerics/SystemSolvers/columnwise_lu_solver.jl b/src/Numerics/SystemSolvers/columnwise_lu_solver.jl
index c71d6e7a338..8b2294b8d96 100644
--- a/src/Numerics/SystemSolvers/columnwise_lu_solver.jl
+++ b/src/Numerics/SystemSolvers/columnwise_lu_solver.jl
@@ -269,13 +269,13 @@ function banded_matrix(
     FT = eltype(Q.data)
     device = array_device(Q)
 
-    nstate = number_state_conservative(bl, FT)
+    nstate = number_states(bl, Conservative(), FT)
     N = polynomialorder(grid)
     Nq = N + 1
 
     # p is lower bandwidth
     # q is upper bandwidth
-    eband = number_state_gradient_flux(bl, FT) == 0 ? 1 : 2
+    eband = number_states(bl, GradientFlux(), FT) == 0 ? 1 : 2
     p = q = nstate * Nq * eband - 1
 
     nrealelem = length(topology.realelems)
diff --git a/src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl b/src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl
index b0e896643e3..32563994c92 100644
--- a/src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl
+++ b/src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesqModel.jl
@@ -10,18 +10,13 @@ using ..VariableTemplates
 using ..MPIStateArrays
 using ..Mesh.Filters: apply!
 using ..Mesh.Grids: VerticalDirection
-using ..BalanceLaws: BalanceLaw
+using ..BalanceLaws
 import ..BalanceLaws: nodal_update_auxiliary_state!
 using ..DGMethods.NumericalFluxes: RusanovNumericalFlux
 
 import ..DGMethods.NumericalFluxes: update_penalty!
-import ..DGMethods:
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_state_gradient,
-    vars_integrals,
-    vars_state_gradient_flux,
-    vars_reverse_integrals,
+import ..BalanceLaws:
+    vars_state,
     init_state_conservative!,
     init_state_auxiliary!,
     compute_gradient_argument!,
@@ -122,7 +117,7 @@ end
 HBModel = HydrostaticBoussinesqModel
 
 """
-    vars_state_conservative(::HBModel)
+    vars_state(::HBModel, ::Conservative)
 
 prognostic variables evolved forward in time
 
@@ -130,7 +125,7 @@ u = (u,v) = (zonal velocity, meridional velocity)
 η = sea surface height
 θ = temperature
 """
-function vars_state_conservative(m::HBModel, T)
+function vars_state(m::HBModel, ::Conservative, T)
     @vars begin
         u::SVector{2, T}
         η::T # real a 2-D variable TODO: should be 2D
@@ -150,7 +145,7 @@ function init_state_conservative!(m::HBModel, Q::Vars, A::Vars, coords, t)
 end
 
 """
-    vars_state_auxiliary(::HBModel)
+    vars_state(::HBModel, ::Auxiliary)
 helper variables for computation
 
 second half is because there is no dedicated integral kernels
@@ -164,7 +159,7 @@ first half of these are fields that are used for computation
 y = north-south coordinate
 
 """
-function vars_state_auxiliary(m::HBModel, T)
+function vars_state(m::HBModel, ::Auxiliary, T)
     @vars begin
         y::T     # y-coordinate of the box
         w::T     # ∫(-∇⋅u)
@@ -185,12 +180,12 @@ function init_state_auxiliary!(m::HBModel, A::Vars, geom::LocalGeometry)
 end
 
 """
-    vars_state_gradient(::HBModel)
+    vars_state(::HBModel, ::Gradient)
 
 variables that you want to take a gradient of
 these are just copies in our model
 """
-function vars_state_gradient(m::HBModel, T)
+function vars_state(m::HBModel, ::Gradient, T)
     @vars begin
         ∇u::SVector{2, T}
         ∇θ::T
@@ -218,12 +213,12 @@ this computation is done pointwise at each nodal point
 end
 
 """
-    vars_state_gradient_flux(::HBModel)
+    vars_state(::HBModel, ::GradientFlux, FT)
 
 the output of the gradient computations
 multiplies ∇u by viscosity tensor and ∇θ by the diffusivity tensor
 """
-function vars_state_gradient_flux(m::HBModel, T)
+function vars_state(m::HBModel, ::GradientFlux, T)
     @vars begin
         ν∇u::SMatrix{3, 2, T, 6}
         κ∇θ::SVector{3, T}
@@ -294,7 +289,7 @@ end
 location to store integrands for bottom up integrals
 ∇hu = the horizontal divegence of u, e.g. dw/dz
 """
-function vars_integrals(m::HBModel, T)
+function vars_state(m::HBModel, ::UpwardIntegrals, T)
     @vars begin
         ∇hu::T
         αᵀθ::T
@@ -349,7 +344,7 @@ end
 location to store integrands for top down integrals
 αᵀθ = density perturbation
 """
-function vars_reverse_integrals(m::HBModel, T)
+function vars_state(m::HBModel, ::DownwardIntegrals, T)
     @vars begin
         αᵀθ::T
     end
@@ -639,8 +634,8 @@ function update_auxiliary_state_gradient!(
     # [1] to convert from range to integer
     # copy_stack_field_down! doesn't like ranges
     # eventually replace this with a reshape and broadcast
-    index_w = varsindex(vars_state_auxiliary(m, FT), :w)[1]
-    index_wz0 = varsindex(vars_state_auxiliary(m, FT), :wz0)[1]
+    index_w = varsindex(vars_state(m, Auxiliary(), FT), :w)[1]
+    index_wz0 = varsindex(vars_state(m, Auxiliary(), FT), :wz0)[1]
     copy_stack_field_down!(dg, m, A, index_w, index_wz0, elems)
 
     return true
diff --git a/src/Ocean/HydrostaticBoussinesq/LinearHBModel.jl b/src/Ocean/HydrostaticBoussinesq/LinearHBModel.jl
index adad63c3c20..3b2c6d771f3 100644
--- a/src/Ocean/HydrostaticBoussinesq/LinearHBModel.jl
+++ b/src/Ocean/HydrostaticBoussinesq/LinearHBModel.jl
@@ -24,13 +24,13 @@ end
 """
     Copy over state, aux, and diff variables from HBModel
 """
-vars_state_conservative(lm::LinearHBModel, FT) =
-    vars_state_conservative(lm.ocean, FT)
-vars_state_gradient(lm::LinearHBModel, FT) = vars_state_gradient(lm.ocean, FT)
-vars_state_gradient_flux(lm::LinearHBModel, FT) =
-    vars_state_gradient_flux(lm.ocean, FT)
-vars_state_auxiliary(lm::LinearHBModel, FT) = vars_state_auxiliary(lm.ocean, FT)
-vars_integrals(lm::LinearHBModel, FT) = @vars()
+vars_state(lm::LinearHBModel, ::Conservative, FT) =
+    vars_state(lm.ocean, Conservative(), FT)
+vars_state(lm::LinearHBModel, st::Gradient, FT) = vars_state(lm.ocean, st, FT)
+vars_state(lm::LinearHBModel, ::GradientFlux, FT) =
+    vars_state(lm.ocean, GradientFlux(), FT)
+vars_state(lm::LinearHBModel, st::Auxiliary, FT) = vars_state(lm.ocean, st, FT)
+vars_state(lm::LinearHBModel, ::UpwardIntegrals, FT) = @vars()
 
 """
     No integration, hyperbolic flux, or source terms
diff --git a/src/Ocean/ShallowWater/ShallowWaterModel.jl b/src/Ocean/ShallowWater/ShallowWaterModel.jl
index a45efb93d09..0aefac02e2c 100644
--- a/src/Ocean/ShallowWater/ShallowWaterModel.jl
+++ b/src/Ocean/ShallowWater/ShallowWaterModel.jl
@@ -7,13 +7,9 @@ using ..VariableTemplates
 using LinearAlgebra: Diagonal, dot
 using CLIMAParameters.Planet: grav
 
+using ..BalanceLaws
 import ..BalanceLaws:
-    BalanceLaw,
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
-    vars_integrals,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -55,27 +51,27 @@ end
 
 SWModel = ShallowWaterModel
 
-function vars_state_conservative(m::SWModel, T)
+function vars_state(m::SWModel, ::Conservative, T)
     @vars begin
         η::T
         U::SVector{2, T}
     end
 end
 
-function vars_state_auxiliary(m::SWModel, T)
+function vars_state(m::SWModel, ::Auxiliary, T)
     @vars begin
         f::T
         τ::SVector{2, T}  # value includes τₒ, g, and ρ
     end
 end
 
-function vars_state_gradient(m::SWModel, T)
+function vars_state(m::SWModel, ::Gradient, T)
     @vars begin
         ∇U::SVector{2, T}
     end
 end
 
-function vars_state_gradient_flux(m::SWModel, T)
+function vars_state(m::SWModel, ::GradientFlux, T)
     @vars begin
         ν∇U::SMatrix{3, 2, T, 6}
     end
diff --git a/src/Utilities/Callbacks/Callbacks.jl b/src/Utilities/Callbacks/Callbacks.jl
index d3fcf04cb89..4c3c8235b4a 100644
--- a/src/Utilities/Callbacks/Callbacks.jl
+++ b/src/Utilities/Callbacks/Callbacks.jl
@@ -14,7 +14,7 @@ using CLIMAParameters.Planet: day
 using ..Courant
 using ..Checkpoint
 using ..DGMethods: courant
-using ..BalanceLaws: vars_state_conservative, vars_state_auxiliary
+using ..BalanceLaws: vars_state
 using ..Diagnostics
 using ..GenericCallbacks
 using ..MPIStateArrays
@@ -151,8 +151,8 @@ function vtk(vtk_opt, solver_config, output_dir)
             )
             outprefix = joinpath(output_dir, vprefix)
 
-            statenames = flattenednames(vars_state_conservative(bl, FT))
-            auxnames = flattenednames(vars_state_auxiliary(bl, FT))
+            statenames = flattenednames(vars_state(bl, Conservative(), FT))
+            auxnames = flattenednames(vars_state(bl, Auxiliary(), FT))
 
             writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)
 
diff --git a/src/Utilities/SingleStackUtils/SingleStackUtils.jl b/src/Utilities/SingleStackUtils/SingleStackUtils.jl
index 6b244c4ec9e..a45d72abf13 100644
--- a/src/Utilities/SingleStackUtils/SingleStackUtils.jl
+++ b/src/Utilities/SingleStackUtils/SingleStackUtils.jl
@@ -25,7 +25,7 @@ using ..VariableTemplates
     ) where {T, dim, N}
 
 Return a dictionary whose keys are the `flattenednames()` of the variables
-specified in `vars` (as returned by e.g. `vars_state_conservative()`), and
+specified in `vars` (as returned by e.g. `vars_state`), and
 whose values are arrays of the values for that variable along the vertical
 dimension in `Q`. Only a single element is expected in the horizontal as
 this is intended for the single stack configuration and `i` and `j` identify
@@ -120,7 +120,7 @@ end
     ) where {T, dim, N}
 
 Return a dictionary whose keys are the `flattenednames()` of the variables
-specified in `vars` (as returned by e.g. `vars_state_conservative()`), and
+specified in `vars` (as returned by e.g. `vars_state`), and
 whose values are arrays of the horizontal averages for that variable along
 the vertical dimension in `Q`. Only a single element is expected in the
 horizontal as this is intended for the single stack configuration.
@@ -165,7 +165,7 @@ end
     ) where {T, dim, N}
 
 Return a dictionary whose keys are the `flattenednames()` of the variables
-specified in `vars` (as returned by e.g. `vars_state_conservative()`), and
+specified in `vars` (as returned by e.g. `vars_state`), and
 whose values are arrays of the horizontal variance for that variable along
 the vertical dimension in `Q`. Only a single element is expected in the
 horizontal as this is intended for the single stack configuration.
diff --git a/test/Atmos/Parameterizations/Microphysics/KM_ice.jl b/test/Atmos/Parameterizations/Microphysics/KM_ice.jl
index f3163fe4f0c..c680099f82e 100644
--- a/test/Atmos/Parameterizations/Microphysics/KM_ice.jl
+++ b/test/Atmos/Parameterizations/Microphysics/KM_ice.jl
@@ -2,7 +2,7 @@ using Dierckx
 
 include("KinematicModel.jl")
 
-function vars_state_conservative(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Conservative, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
@@ -15,7 +15,7 @@ function vars_state_conservative(m::KinematicModel, FT)
     end
 end
 
-function vars_state_auxiliary(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Auxiliary, FT)
     @vars begin
         # defined in init_state_auxiliary
         p::FT
@@ -876,20 +876,20 @@ function main()
     mpicomm = MPI.COMM_WORLD
 
     # get state variables indices for filtering
-    ρq_liq_ind = varsindex(vars_state_conservative(model, FT), :ρq_liq)
-    ρq_ice_ind = varsindex(vars_state_conservative(model, FT), :ρq_ice)
-    ρq_rai_ind = varsindex(vars_state_conservative(model, FT), :ρq_rai)
-    ρq_sno_ind = varsindex(vars_state_conservative(model, FT), :ρq_sno)
+    ρq_liq_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_liq)
+    ρq_ice_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_ice)
+    ρq_rai_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_rai)
+    ρq_sno_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_sno)
     # get aux variables indices for testing
-    q_tot_ind = varsindex(vars_state_auxiliary(model, FT), :q_tot)
-    q_vap_ind = varsindex(vars_state_auxiliary(model, FT), :q_vap)
-    q_liq_ind = varsindex(vars_state_auxiliary(model, FT), :q_liq)
-    q_ice_ind = varsindex(vars_state_auxiliary(model, FT), :q_ice)
-    q_rai_ind = varsindex(vars_state_auxiliary(model, FT), :q_rai)
-    q_sno_ind = varsindex(vars_state_auxiliary(model, FT), :q_sno)
-    S_ind = varsindex(vars_state_auxiliary(model, FT), :S)
-    rain_w_ind = varsindex(vars_state_auxiliary(model, FT), :rain_w)
-    snow_w_ind = varsindex(vars_state_auxiliary(model, FT), :snow_w)
+    q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)
+    q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)
+    q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)
+    q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)
+    q_rai_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_rai)
+    q_sno_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_sno)
+    S_ind = varsindex(vars_state(model, Auxiliary(), FT), :S)
+    rain_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :rain_w)
+    snow_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :snow_w)
 
     # filter out negative values
     cb_tmar_filter =
@@ -926,9 +926,9 @@ function main()
                 out_path_prefix,
                 solver_config.Q,
                 solver_config.dg,
-                flattenednames(vars_state_conservative(model, FT)),
+                flattenednames(vars_state(model, Conservative(), FT)),
                 solver_config.dg.state_auxiliary,
-                flattenednames(vars_state_auxiliary(model, FT)),
+                flattenednames(vars_state(model, Auxiliary(), FT)),
             )
             step[1] += 1
             nothing
diff --git a/test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl b/test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl
index b77475d6bf5..f45df8e3d24 100644
--- a/test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl
+++ b/test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl
@@ -1,6 +1,6 @@
 include("KinematicModel.jl")
 
-function vars_state_conservative(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Conservative, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
@@ -9,7 +9,7 @@ function vars_state_conservative(m::KinematicModel, FT)
     end
 end
 
-function vars_state_auxiliary(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Auxiliary, FT)
     @vars begin
         # defined in init_state_auxiliary
         p::FT
@@ -253,20 +253,20 @@ function main()
             out_path_prefix,
             solver_config.Q,
             solver_config.dg,
-            flattenednames(vars_state_conservative(model, FT)),
+            flattenednames(vars_state(model, Conservative(), FT)),
             solver_config.dg.state_auxiliary,
-            flattenednames(vars_state_auxiliary(model, FT)),
+            flattenednames(vars_state(model, Auxiliary(), FT)),
         )
         step[1] += 1
         nothing
     end
 
     # get aux variables indices for testing
-    q_tot_ind = varsindex(vars_state_auxiliary(model, FT), :q_tot)
-    q_vap_ind = varsindex(vars_state_auxiliary(model, FT), :q_vap)
-    q_liq_ind = varsindex(vars_state_auxiliary(model, FT), :q_liq)
-    q_ice_ind = varsindex(vars_state_auxiliary(model, FT), :q_ice)
-    S_ind = varsindex(vars_state_auxiliary(model, FT), :S)
+    q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)
+    q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)
+    q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)
+    q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)
+    S_ind = varsindex(vars_state(model, Auxiliary(), FT), :S)
 
     # call solve! function for time-integrator
     result = ClimateMachine.invoke!(
diff --git a/test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl b/test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl
index 4e3bf5e1c5e..7dd5d44ab87 100644
--- a/test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl
+++ b/test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl
@@ -1,6 +1,6 @@
 include("KinematicModel.jl")
 
-function vars_state_conservative(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Conservative, FT)
     @vars begin
         ρ::FT
         ρu::SVector{3, FT}
@@ -12,7 +12,7 @@ function vars_state_conservative(m::KinematicModel, FT)
     end
 end
 
-function vars_state_auxiliary(m::KinematicModel, FT)
+function vars_state(m::KinematicModel, ::Auxiliary, FT)
     @vars begin
         # defined in init_state_auxiliary
         p::FT
@@ -399,17 +399,17 @@ function main()
     mpicomm = MPI.COMM_WORLD
 
     # get state variables indices for filtering
-    ρq_liq_ind = varsindex(vars_state_conservative(model, FT), :ρq_liq)
-    ρq_ice_ind = varsindex(vars_state_conservative(model, FT), :ρq_ice)
-    ρq_rai_ind = varsindex(vars_state_conservative(model, FT), :ρq_rai)
+    ρq_liq_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_liq)
+    ρq_ice_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_ice)
+    ρq_rai_ind = varsindex(vars_state(model, Conservative(), FT), :ρq_rai)
     # get aux variables indices for testing
-    q_tot_ind = varsindex(vars_state_auxiliary(model, FT), :q_tot)
-    q_vap_ind = varsindex(vars_state_auxiliary(model, FT), :q_vap)
-    q_liq_ind = varsindex(vars_state_auxiliary(model, FT), :q_liq)
-    q_ice_ind = varsindex(vars_state_auxiliary(model, FT), :q_ice)
-    q_rai_ind = varsindex(vars_state_auxiliary(model, FT), :q_rai)
-    S_ind = varsindex(vars_state_auxiliary(model, FT), :S)
-    rain_w_ind = varsindex(vars_state_auxiliary(model, FT), :rain_w)
+    q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)
+    q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)
+    q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)
+    q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)
+    q_rai_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_rai)
+    S_ind = varsindex(vars_state(model, Auxiliary(), FT), :S)
+    rain_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :rain_w)
 
     # filter out negative values
     cb_tmar_filter =
@@ -446,9 +446,9 @@ function main()
                 out_path_prefix,
                 solver_config.Q,
                 solver_config.dg,
-                flattenednames(vars_state_conservative(model, FT)),
+                flattenednames(vars_state(model, Conservative(), FT)),
                 solver_config.dg.state_auxiliary,
-                flattenednames(vars_state_auxiliary(model, FT)),
+                flattenednames(vars_state(model, Auxiliary(), FT)),
             )
             step[1] += 1
             nothing
diff --git a/test/Atmos/Parameterizations/Microphysics/KinematicModel.jl b/test/Atmos/Parameterizations/Microphysics/KinematicModel.jl
index 42240fb3b19..029878ca4d2 100644
--- a/test/Atmos/Parameterizations/Microphysics/KinematicModel.jl
+++ b/test/Atmos/Parameterizations/Microphysics/KinematicModel.jl
@@ -86,13 +86,11 @@ const ice_param_set = param_set.microphys_param_set.ice
 const rain_param_set = param_set.microphys_param_set.rain
 const snow_param_set = param_set.microphys_param_set.snow
 
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, Gradient, GradientFlux, Hyperdiffusive
 
 import ClimateMachine.BalanceLaws:
-    vars_state_conservative,
-    vars_state_auxiliary,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     init_state_conservative!,
     init_state_auxiliary!,
     update_auxiliary_state!,
@@ -167,9 +165,9 @@ function KinematicModel{FT}(
     return KinematicModel{FT, typeof.(atmos)...}(atmos...)
 end
 
-vars_state_gradient(m::KinematicModel, FT) = @vars()
+vars_state(m::KinematicModel, ::Gradient, FT) = @vars()
 
-vars_state_gradient_flux(m::KinematicModel, FT) = @vars()
+vars_state(m::KinematicModel, ::GradientFlux, FT) = @vars()
 
 function init_state_auxiliary!(
     m::KinematicModel,
diff --git a/test/Diagnostics/diagnostic_fields_test.jl b/test/Diagnostics/diagnostic_fields_test.jl
index b952310b45a..2c23ec924f6 100644
--- a/test/Diagnostics/diagnostic_fields_test.jl
+++ b/test/Diagnostics/diagnostic_fields_test.jl
@@ -9,6 +9,7 @@ using ClimateMachine.Atmos
 using ClimateMachine.Orientations
 using ClimateMachine.ConfigTypes
 using ClimateMachine.Diagnostics
+using ClimateMachine.BalanceLaws
 using ClimateMachine.GenericCallbacks
 using ClimateMachine.ODESolvers
 using ClimateMachine.Mesh.Filters
@@ -23,7 +24,7 @@ struct EarthParameterSet <: AbstractEarthParameterSet end
 const param_set = EarthParameterSet()
 
 import ClimateMachine.Mesh.Grids: _x1, _x2, _x3
-import ClimateMachine.BalanceLaws: vars_state_conservative
+import ClimateMachine.BalanceLaws: vars_state
 import ClimateMachine.VariableTemplates.varsindex
 
 # ------------------------ Description ------------------------- #
@@ -175,8 +176,8 @@ function run_brick_diagostics_fields_test()
         Npl = size(Q.realdata, 1)
 
         ind = [
-            varsindex(vars_state_conservative(model, FT), :ρ)
-            varsindex(vars_state_conservative(model, FT), :ρu)
+            varsindex(vars_state(model, Conservative(), FT), :ρ)
+            varsindex(vars_state(model, Conservative(), FT), :ρu)
         ]
         _ρ, _ρu, _ρv, _ρw = ind[1], ind[2], ind[3], ind[4]
 
diff --git a/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl b/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl
index d1c7da4ffe9..fcbc2471056 100644
--- a/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl
+++ b/test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl
@@ -28,8 +28,7 @@ using ClimateMachine.Atmos:
     NoPrecipitation,
     NoRadiation,
     NTracers,
-    vars_state_conservative,
-    vars_state_auxiliary,
+    vars_state,
     Gravity,
     HydrostaticState,
     AtmosAcousticGravityLinearModel
@@ -312,8 +311,8 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
-    auxnames = flattenednames(vars_state_auxiliary(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
+    auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
     writevtk(filename, Q, dg, statenames, dg.state_auxiliary, auxnames)
 
     ## Generate the pvtu file for these vtk files
diff --git a/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl b/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl
index 26af6349e90..2024ec1316e 100644
--- a/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl
+++ b/test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl
@@ -33,8 +33,7 @@ using ClimateMachine.Atmos:
     NoRadiation,
     NTracers,
     ConstantViscosityWithDivergence,
-    vars_state_conservative,
-    vars_state_auxiliary,
+    vars_state,
     Gravity,
     HydrostaticState,
     AtmosAcousticGravityLinearModel,
@@ -353,8 +352,8 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
-    auxnames = flattenednames(vars_state_auxiliary(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
+    auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
     writevtk(filename, Q, dg, statenames, dg.state_auxiliary, auxnames)
 
     ## Generate the pvtu file for these vtk files
diff --git a/test/Numerics/DGMethods/Euler/isentropicvortex.jl b/test/Numerics/DGMethods/Euler/isentropicvortex.jl
index f6d75fcc47e..d5d4bda159d 100644
--- a/test/Numerics/DGMethods/Euler/isentropicvortex.jl
+++ b/test/Numerics/DGMethods/Euler/isentropicvortex.jl
@@ -21,10 +21,11 @@ using ClimateMachine.Atmos:
     DryModel,
     NoPrecipitation,
     NoRadiation,
-    vars_state_conservative
+    vars_state
 using ClimateMachine.Orientations: NoOrientation
 using ClimateMachine.VariableTemplates: flattenednames
 using ClimateMachine.TurbulenceClosures
+using ClimateMachine.BalanceLaws: Conservative
 
 using CLIMAParameters
 using CLIMAParameters.Planet: kappa_d
@@ -355,7 +356,7 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl b/test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl
index 3cd2d45e72a..c2b2ff191fe 100644
--- a/test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl
+++ b/test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl
@@ -24,11 +24,12 @@ using ClimateMachine.Atmos:
     DryModel,
     NoPrecipitation,
     NoRadiation,
-    vars_state_conservative
+    vars_state
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.Orientations: NoOrientation
 using ClimateMachine.VariableTemplates: @vars, Vars, flattenednames
-import ClimateMachine.Atmos: atmos_init_aux!, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: Conservative, Auxiliary
+import ClimateMachine.Atmos: atmos_init_aux!, vars_state
 
 using CLIMAParameters
 using CLIMAParameters.Planet: kappa_d
@@ -291,7 +292,7 @@ end
 struct IsentropicVortexReferenceState{FT} <: ReferenceState
     setup::IsentropicVortexSetup{FT}
 end
-vars_state_auxiliary(::IsentropicVortexReferenceState, FT) =
+vars_state(::IsentropicVortexReferenceState, ::Auxiliary, FT) =
     @vars(ρ::FT, ρe::FT, p::FT, T::FT)
 function atmos_init_aux!(
     m::IsentropicVortexReferenceState,
@@ -368,7 +369,7 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl b/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl
index cebe4ad7597..77ec45aea0b 100644
--- a/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl
+++ b/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl
@@ -24,11 +24,12 @@ using ClimateMachine.Atmos:
     DryModel,
     NoPrecipitation,
     NoRadiation,
-    vars_state_conservative
+    vars_state
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.Orientations: NoOrientation
 using ClimateMachine.VariableTemplates: @vars, Vars, flattenednames
-import ClimateMachine.Atmos: atmos_init_aux!, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: Conservative, Auxiliary
+import ClimateMachine.Atmos: atmos_init_aux!, vars_state
 
 using CLIMAParameters
 using CLIMAParameters.Planet: kappa_d
@@ -276,7 +277,7 @@ end
 struct IsentropicVortexReferenceState{FT} <: ReferenceState
     setup::IsentropicVortexSetup{FT}
 end
-vars_state_auxiliary(::IsentropicVortexReferenceState, FT) =
+vars_state(::IsentropicVortexReferenceState, ::Auxiliary, FT) =
     @vars(ρ::FT, ρe::FT, p::FT, T::FT)
 function atmos_init_aux!(
     m::IsentropicVortexReferenceState,
@@ -353,7 +354,7 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl b/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl
index 2cc081deec3..8b14f359d7a 100644
--- a/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl
+++ b/test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl
@@ -24,11 +24,12 @@ using ClimateMachine.Atmos:
     DryModel,
     NoPrecipitation,
     NoRadiation,
-    vars_state_conservative
+    vars_state
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.Orientations: NoOrientation
 using ClimateMachine.VariableTemplates: @vars, Vars, flattenednames
-import ClimateMachine.Atmos: atmos_init_aux!, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: Conservative, Auxiliary
+import ClimateMachine.Atmos: atmos_init_aux!, vars_state
 
 using CLIMAParameters
 using CLIMAParameters.Planet: kappa_d
@@ -286,7 +287,7 @@ end
 struct IsentropicVortexReferenceState{FT} <: ReferenceState
     setup::IsentropicVortexSetup{FT}
 end
-vars_state_auxiliary(::IsentropicVortexReferenceState, FT) =
+vars_state(::IsentropicVortexReferenceState, ::Auxiliary, FT) =
     @vars(ρ::FT, ρe::FT, p::FT, T::FT)
 function atmos_init_aux!(
     m::IsentropicVortexReferenceState,
@@ -363,7 +364,7 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl b/test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl
index 1b8691fbacd..a1a2b9c096e 100644
--- a/test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl
+++ b/test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl
@@ -24,11 +24,12 @@ using ClimateMachine.Atmos:
     DryModel,
     NoPrecipitation,
     NoRadiation,
-    vars_state_conservative
+    vars_state
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.Orientations: NoOrientation
 using ClimateMachine.VariableTemplates: @vars, Vars, flattenednames
-import ClimateMachine.Atmos: atmos_init_aux!, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: Conservative, Auxiliary
+import ClimateMachine.Atmos: atmos_init_aux!, vars_state
 
 using CLIMAParameters
 using CLIMAParameters.Planet: kappa_d
@@ -299,7 +300,7 @@ end
 struct IsentropicVortexReferenceState{FT} <: ReferenceState
     setup::IsentropicVortexSetup{FT}
 end
-vars_state_auxiliary(::IsentropicVortexReferenceState, FT) =
+vars_state(::IsentropicVortexReferenceState, ::Auxiliary, FT) =
     @vars(ρ::FT, ρe::FT, p::FT, T::FT)
 function atmos_init_aux!(
     m::IsentropicVortexReferenceState,
@@ -375,7 +376,7 @@ function do_output(
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model.jl b/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model.jl
index 8c934b22b19..215eb8baf45 100644
--- a/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model.jl
@@ -1,13 +1,12 @@
 using StaticArrays
 using ClimateMachine.VariableTemplates
 using ClimateMachine.DGMethods: nodal_update_auxiliary_state!
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, Gradient, GradientFlux
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
+    number_states,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -17,9 +16,7 @@ import ClimateMachine.BalanceLaws:
     update_auxiliary_state!,
     init_state_conservative!,
     boundary_state!,
-    wavespeed,
-    number_state_conservative,
-    number_state_gradient
+    wavespeed
 
 using ClimateMachine.Mesh.Geometry: LocalGeometry
 using ClimateMachine.DGMethods.NumericalFluxes:
@@ -55,10 +52,11 @@ end
 #   `coord` coordinate points (needed for BCs)
 #   `u` advection velocity
 #   `D` Diffusion tensor
-vars_state_auxiliary(::AdvectionDiffusion, FT) =
+vars_state(::AdvectionDiffusion, ::Auxiliary, FT) =
     @vars(coord::SVector{3, FT}, u::SVector{3, FT}, D::SMatrix{3, 3, FT, 9})
-function vars_state_auxiliary(
+function vars_state(
     ::AdvectionDiffusion{dim, P, fluxBC, true},
+    ::Auxiliary,
     FT,
 ) where {dim, P, fluxBC}
     @vars begin
@@ -68,19 +66,21 @@ function vars_state_auxiliary(
 end
 
 # Density is only state
-vars_state_conservative(::AdvectionDiffusion, FT) = @vars(ρ::FT)
+vars_state(::AdvectionDiffusion, ::Conservative, FT) = @vars(ρ::FT)
 
 # Take the gradient of density
-vars_state_gradient(::AdvectionDiffusion, FT) = @vars(ρ::FT)
-vars_state_gradient(
+vars_state(::AdvectionDiffusion, ::Gradient, FT) = @vars(ρ::FT)
+vars_state(
     ::AdvectionDiffusion{dim, P, fluxBC, true},
+    ::Gradient,
     FT,
 ) where {dim, P, fluxBC} = @vars()
 
 # The DG auxiliary variable: D ∇ρ
-vars_state_gradient_flux(::AdvectionDiffusion, FT) = @vars(σ::SVector{3, FT})
-vars_state_gradient_flux(
+vars_state(::AdvectionDiffusion, ::GradientFlux, FT) = @vars(σ::SVector{3, FT})
+vars_state(
     ::AdvectionDiffusion{dim, P, fluxBC, true},
+    ::GradientFlux,
     FT,
 ) where {dim, P, fluxBC} = @vars()
 
@@ -299,8 +299,8 @@ function boundary_state!(
         diff⁺.σ = diff⁻.σ
     elseif bctype == 2 # Neumann with data
         FT = eltype(diff⁺)
-        ngrad = number_state_gradient(m, FT)
-        ∇state = Grad{vars_state_gradient(m, FT)}(similar(
+        ngrad = number_states(m, Gradient(), FT)
+        ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(
             parent(diff⁺),
             Size(3, ngrad),
         ))
@@ -310,8 +310,8 @@ function boundary_state!(
         # compute the diffusive flux using the boundary state
     elseif bctype == 4 # zero Neumann
         FT = eltype(diff⁺)
-        ngrad = number_state_gradient(m, FT)
-        ∇state = Grad{vars_state_gradient(m, FT)}(similar(
+        ngrad = number_states(m, Gradient(), FT)
+        ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(
             parent(diff⁺),
             Size(3, ngrad),
         ))
@@ -361,8 +361,8 @@ function boundary_flux_second_order!(
         flux_second_order!(m, F, state⁻, diff⁻, hyperdiff⁻, aux⁻, t)
     elseif bctype == 2 # Neumann data
         FT = eltype(diff⁺)
-        ngrad = number_state_gradient(m, FT)
-        ∇state = Grad{vars_state_gradient(m, FT)}(similar(
+        ngrad = number_states(m, Gradient(), FT)
+        ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(
             parent(diff⁺),
             Size(3, ngrad),
         ))
diff --git a/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl b/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl
index 6e10adeeb2a..0aca6e64bb0 100644
--- a/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl
@@ -82,7 +82,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl b/test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl
index dfeac4675ae..e1f3a834702 100644
--- a/test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl
@@ -165,7 +165,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_model.jl b/test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_model.jl
index f44ab85ba54..e114719f48d 100644
--- a/test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_model.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_model.jl
@@ -1,11 +1,15 @@
 using StaticArrays
 using ClimateMachine.VariableTemplates
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw,
+    Conservative,
+    Auxiliary,
+    Gradient,
+    GradientFlux,
+    GradientLaplacian,
+    Hyperdiffusive
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -15,8 +19,6 @@ import ClimateMachine.BalanceLaws:
     init_state_conservative!,
     boundary_state!,
     wavespeed,
-    vars_gradient_laplacian,
-    vars_hyperdiffusive,
     transform_post_gradient_laplacian!
 
 using ClimateMachine.Mesh.Geometry: LocalGeometry
@@ -33,19 +35,19 @@ struct HyperDiffusion{dim, P} <: BalanceLaw
     end
 end
 
-vars_state_auxiliary(::HyperDiffusion, FT) = @vars(D::SMatrix{3, 3, FT, 9})
+vars_state(::HyperDiffusion, ::Auxiliary, FT) = @vars(D::SMatrix{3, 3, FT, 9})
 #
 # Density is only state
-vars_state_conservative(::HyperDiffusion, FT) = @vars(ρ::FT)
+vars_state(::HyperDiffusion, ::Conservative, FT) = @vars(ρ::FT)
 
 # Take the gradient of density
-vars_state_gradient(::HyperDiffusion, FT) = @vars(ρ::FT)
+vars_state(::HyperDiffusion, ::Gradient, FT) = @vars(ρ::FT)
 # Take the gradient of laplacian of density
-vars_gradient_laplacian(::HyperDiffusion, FT) = @vars(ρ::FT)
+vars_state(::HyperDiffusion, ::GradientLaplacian, FT) = @vars(ρ::FT)
 
-vars_state_gradient_flux(::HyperDiffusion, FT) = @vars()
+vars_state(::HyperDiffusion, ::GradientFlux, FT) = @vars()
 # The hyperdiffusion DG auxiliary variable: D ∇ Δρ
-vars_hyperdiffusive(::HyperDiffusion, FT) = @vars(σ::SVector{3, FT})
+vars_state(::HyperDiffusion, ::Hyperdiffusive, FT) = @vars(σ::SVector{3, FT})
 
 function flux_first_order!(m::HyperDiffusion, _...) end
 
diff --git a/test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl b/test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl
index a9a3ddfcf3c..38ec060ada0 100644
--- a/test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl
@@ -71,7 +71,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl
index 546ae43c3ce..d094799088c 100644
--- a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl
@@ -76,7 +76,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl
index 0e8c87c4f2b..52b18707b0f 100644
--- a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl
@@ -82,7 +82,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl
index 473f90aa077..a05914aa75b 100644
--- a/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl
+++ b/test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl
@@ -82,7 +82,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
     exactnames = statenames .* "_exact"
 
     writevtk(filename, Q, dg, statenames, Qe, exactnames)
diff --git a/test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl b/test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl
index 0f24981f081..9295b78b7ae 100644
--- a/test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl
+++ b/test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl
@@ -22,7 +22,7 @@ using StaticArrays
 using Logging, Printf, Dates
 using ClimateMachine.VTK
 using Random
-using ClimateMachine.Atmos: vars_state_conservative, vars_state_auxiliary
+using ClimateMachine.Atmos: vars_state
 
 using CLIMAParameters
 using CLIMAParameters.Planet: R_d, cp_d, cv_d, grav, MSLP
@@ -194,9 +194,9 @@ function run(
             outprefix,
             Q,
             dg,
-            flattenednames(vars_state_conservative(model, FT)),
+            flattenednames(vars_state(model, Conservative(), FT)),
             dg.state_auxiliary,
-            flattenednames(vars_state_auxiliary(model, FT)),
+            flattenednames(vars_state(model, Auxiliary(), FT)),
         )
         step[1] += 1
         nothing
diff --git a/test/Numerics/DGMethods/compressible_Navier_Stokes/mms_model.jl b/test/Numerics/DGMethods/compressible_Navier_Stokes/mms_model.jl
index 41d1ae32b54..f4e7ed81e8b 100644
--- a/test/Numerics/DGMethods/compressible_Navier_Stokes/mms_model.jl
+++ b/test/Numerics/DGMethods/compressible_Navier_Stokes/mms_model.jl
@@ -1,12 +1,11 @@
 using ClimateMachine.VariableTemplates
 
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, Gradient, GradientFlux
+
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -23,10 +22,11 @@ using ClimateMachine.Mesh.Geometry: LocalGeometry
 
 struct MMSModel{dim} <: BalanceLaw end
 
-vars_state_auxiliary(::MMSModel, T) = @vars(x1::T, x2::T, x3::T)
-vars_state_conservative(::MMSModel, T) = @vars(ρ::T, ρu::T, ρv::T, ρw::T, ρe::T)
-vars_state_gradient(::MMSModel, T) = @vars(u::T, v::T, w::T)
-vars_state_gradient_flux(::MMSModel, T) =
+vars_state(::MMSModel, ::Auxiliary, T) = @vars(x1::T, x2::T, x3::T)
+vars_state(::MMSModel, ::Conservative, T) =
+    @vars(ρ::T, ρu::T, ρv::T, ρw::T, ρe::T)
+vars_state(::MMSModel, ::Gradient, T) = @vars(u::T, v::T, w::T)
+vars_state(::MMSModel, ::GradientFlux, T) =
     @vars(τ11::T, τ22::T, τ33::T, τ12::T, τ13::T, τ23::T)
 
 function flux_first_order!(
diff --git a/test/Numerics/DGMethods/compressible_Navier_Stokes/ref_state.jl b/test/Numerics/DGMethods/compressible_Navier_Stokes/ref_state.jl
index 24097576e30..71d4880254b 100644
--- a/test/Numerics/DGMethods/compressible_Navier_Stokes/ref_state.jl
+++ b/test/Numerics/DGMethods/compressible_Navier_Stokes/ref_state.jl
@@ -8,6 +8,7 @@ using ClimateMachine.DGMethods.NumericalFluxes
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.ODESolvers
 using ClimateMachine.GenericCallbacks
+using ClimateMachine.BalanceLaws
 using ClimateMachine.Atmos
 using ClimateMachine.Orientations
 using ClimateMachine.VariableTemplates
@@ -36,7 +37,7 @@ import ClimateMachine.Atmos: MoistureModel, temperature, pressure, soundspeed
 init_state_conservative!(bl, state, aux, coords, t) = nothing
 
 # initial condition
-using ClimateMachine.Atmos: vars_state_auxiliary
+using ClimateMachine.Atmos: vars_state
 
 function run1(mpicomm, ArrayType, dim, topl, N, timeend, FT, dt)
 
@@ -72,7 +73,7 @@ function run1(mpicomm, ArrayType, dim, topl, N, timeend, FT, dt)
         outprefix,
         dg.state_auxiliary,
         dg,
-        flattenednames(vars_state_auxiliary(model, FT)),
+        flattenednames(vars_state(model, Auxiliary(), FT)),
     )
     return FT(0)
 end
@@ -110,7 +111,7 @@ function run2(mpicomm, ArrayType, dim, topl, N, timeend, FT, dt)
         outprefix,
         dg.state_auxiliary,
         dg,
-        flattenednames(vars_state_auxiliary(model, FT)),
+        flattenednames(vars_state(model, Auxiliary(), FT)),
     )
     return FT(0)
 end
diff --git a/test/Numerics/DGMethods/conservation/sphere.jl b/test/Numerics/DGMethods/conservation/sphere.jl
index 6be72477b18..cec4cc97258 100644
--- a/test/Numerics/DGMethods/conservation/sphere.jl
+++ b/test/Numerics/DGMethods/conservation/sphere.jl
@@ -19,19 +19,19 @@ using ClimateMachine.DGMethods.NumericalFluxes
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.ODESolvers
 using ClimateMachine.GenericCallbacks
+using ClimateMachine.BalanceLaws:
+    Conservative, Auxiliary, AbstractStateType, BalanceLaw
+
 using LinearAlgebra
 using StaticArrays
 using Logging, Printf, Dates
 using Random
 
 using ClimateMachine.VariableTemplates
-using ClimateMachine.DGMethods: BalanceLaw
+
+import ClimateMachine.BalanceLaws: vars_state
 
 import ClimateMachine.DGMethods:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -48,11 +48,10 @@ import ClimateMachine.DGMethods.NumericalFluxes:
 
 struct ConservationTestModel <: BalanceLaw end
 
-vars_state_auxiliary(::ConservationTestModel, T) = @vars(vel::SVector{3, T})
-vars_state_conservative(::ConservationTestModel, T) = @vars(q::T, p::T)
+vars_state(::ConservationTestModel, ::Auxiliary, T) = @vars(vel::SVector{3, T})
+vars_state(::ConservationTestModel, ::Conservative, T) = @vars(q::T, p::T)
 
-vars_state_gradient(::ConservationTestModel, T) = @vars()
-vars_state_gradient_flux(::ConservationTestModel, T) = @vars()
+vars_state(::ConservationTestModel, ::AbstractStateType, T) = @vars()
 
 function init_state_auxiliary!(
     ::ConservationTestModel,
diff --git a/test/Numerics/DGMethods/courant.jl b/test/Numerics/DGMethods/courant.jl
index 2009abb2aa3..4f47d5d51ff 100644
--- a/test/Numerics/DGMethods/courant.jl
+++ b/test/Numerics/DGMethods/courant.jl
@@ -26,7 +26,8 @@ using ClimateMachine.Atmos:
     NoPrecipitation,
     Gravity,
     HydrostaticState,
-    vars_state_conservative,
+    ConstantViscosityWithDivergence,
+    vars_state,
     soundspeed
 using ClimateMachine.TurbulenceClosures
 using ClimateMachine.Orientations
diff --git a/test/Numerics/DGMethods/integral_test.jl b/test/Numerics/DGMethods/integral_test.jl
index 5161b716c12..5ef58de241c 100644
--- a/test/Numerics/DGMethods/integral_test.jl
+++ b/test/Numerics/DGMethods/integral_test.jl
@@ -11,13 +11,10 @@ using Printf
 using LinearAlgebra
 using Logging
 
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -26,8 +23,6 @@ import ClimateMachine.BalanceLaws:
     init_state_auxiliary!,
     init_state_conservative!,
     update_auxiliary_state!,
-    vars_integrals,
-    vars_reverse_integrals,
     indefinite_stack_integral!,
     reverse_indefinite_stack_integral!,
     integral_load_auxiliary_state!,
@@ -41,11 +36,11 @@ using ClimateMachine.Mesh.Geometry: LocalGeometry
 
 struct IntegralTestModel{dim} <: BalanceLaw end
 
-vars_reverse_integrals(::IntegralTestModel, T) = @vars(a::T, b::T)
-vars_integrals(::IntegralTestModel, T) = @vars(a::T, b::T)
-vars_state_auxiliary(m::IntegralTestModel, T) = @vars(
-    int::vars_integrals(m, T),
-    rev_int::vars_reverse_integrals(m, T),
+vars_state(::IntegralTestModel, ::DownwardIntegrals, T) = @vars(a::T, b::T)
+vars_state(::IntegralTestModel, ::UpwardIntegrals, T) = @vars(a::T, b::T)
+vars_state(m::IntegralTestModel, ::Auxiliary, T) = @vars(
+    int::vars_state(m, UpwardIntegrals(), T),
+    rev_int::vars_state(m, DownwardIntegrals(), T),
     coord::SVector{3, T},
     a::T,
     b::T,
@@ -53,8 +48,7 @@ vars_state_auxiliary(m::IntegralTestModel, T) = @vars(
     rev_b::T
 )
 
-vars_state_conservative(::IntegralTestModel, T) = @vars()
-vars_state_gradient_flux(::IntegralTestModel, T) = @vars()
+vars_state(::IntegralTestModel, ::AbstractStateType, T) = @vars()
 
 flux_first_order!(::IntegralTestModel, _...) = nothing
 flux_second_order!(::IntegralTestModel, _...) = nothing
diff --git a/test/Numerics/DGMethods/integral_test_sphere.jl b/test/Numerics/DGMethods/integral_test_sphere.jl
index a83d77b82cb..38c8ce4465d 100644
--- a/test/Numerics/DGMethods/integral_test_sphere.jl
+++ b/test/Numerics/DGMethods/integral_test_sphere.jl
@@ -11,13 +11,16 @@ using Printf
 using LinearAlgebra
 using Logging
 
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw,
+    Conservative,
+    Auxiliary,
+    GradientFlux,
+    UpwardIntegrals,
+    DownwardIntegrals
+
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
-    vars_integrals,
+    vars_state,
     integral_load_auxiliary_state!,
     flux_first_order!,
     flux_second_order!,
@@ -31,7 +34,6 @@ import ClimateMachine.BalanceLaws:
     init_state_auxiliary!,
     init_state_conservative!,
     integral_set_auxiliary_state!,
-    vars_reverse_integrals,
     reverse_integral_load_auxiliary_state!,
     reverse_integral_set_auxiliary_state!
 
@@ -56,13 +58,17 @@ function update_auxiliary_state!(
     return true
 end
 
-vars_integrals(::IntegralTestSphereModel, T) = @vars(v::T)
-vars_reverse_integrals(::IntegralTestSphereModel, T) = @vars(v::T)
-vars_state_auxiliary(m::IntegralTestSphereModel, T) =
-    @vars(int::vars_integrals(m, T), rev_int::vars_integrals(m, T), r::T, a::T)
+vars_state(::IntegralTestSphereModel, ::UpwardIntegrals, T) = @vars(v::T)
+vars_state(::IntegralTestSphereModel, ::DownwardIntegrals, T) = @vars(v::T)
+vars_state(m::IntegralTestSphereModel, ::Auxiliary, T) = @vars(
+    int::vars_state(m, UpwardIntegrals(), T),
+    rev_int::vars_state(m, UpwardIntegrals(), T),
+    r::T,
+    a::T
+)
 
-vars_state_conservative(::IntegralTestSphereModel, T) = @vars()
-vars_state_gradient_flux(::IntegralTestSphereModel, T) = @vars()
+vars_state(::IntegralTestSphereModel, ::Conservative, T) = @vars()
+vars_state(::IntegralTestSphereModel, ::GradientFlux, T) = @vars()
 
 flux_first_order!(::IntegralTestSphereModel, _...) = nothing
 flux_second_order!(::IntegralTestSphereModel, _...) = nothing
diff --git a/test/Numerics/DGMethods/remainder_model.jl b/test/Numerics/DGMethods/remainder_model.jl
index bcbcecb9d83..7067b174522 100644
--- a/test/Numerics/DGMethods/remainder_model.jl
+++ b/test/Numerics/DGMethods/remainder_model.jl
@@ -12,10 +12,9 @@ using ClimateMachine.DGMethods:
     init_ode_state,
     remainder_DGModel,
     wavespeed,
-    number_state_conservative,
-    number_state_auxiliary,
-    vars_state_conservative,
-    vars_state_auxiliary
+    number_states,
+    vars_state
+using ClimateMachine.BalanceLaws: Conservative, Auxiliary
 using ClimateMachine.VariableTemplates: Vars
 using ClimateMachine.DGMethods.NumericalFluxes:
     RusanovNumericalFlux,
@@ -157,13 +156,14 @@ function run(
     # Create some random data to check the wavespeed function with
     nM = rand(3)
     nM /= norm(nM)
-    state_conservative = Vars{vars_state_conservative(dg.balance_law, FT)}(rand(
-        FT,
-        number_state_conservative(dg.balance_law, FT),
-    ))
-    state_auxiliary = Vars{vars_state_auxiliary(dg.balance_law, FT)}(rand(
+    state_conservative =
+        Vars{vars_state(dg.balance_law, Conservative(), FT)}(rand(
+            FT,
+            number_states(dg.balance_law, Conservative(), FT),
+        ))
+    state_auxiliary = Vars{vars_state(dg.balance_law, Auxiliary(), FT)}(rand(
         FT,
-        number_state_auxiliary(dg.balance_law, FT),
+        number_states(dg.balance_law, Auxiliary(), FT),
     ))
     full_wavespeed = wavespeed(
         dg.balance_law,
diff --git a/test/Numerics/DGMethods/vars_test.jl b/test/Numerics/DGMethods/vars_test.jl
index df14386552c..c693c1b7132 100644
--- a/test/Numerics/DGMethods/vars_test.jl
+++ b/test/Numerics/DGMethods/vars_test.jl
@@ -11,14 +11,11 @@ using Printf
 using LinearAlgebra
 using Logging
 
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, GradientFlux
 
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
-    vars_integrals,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -34,10 +31,11 @@ using ClimateMachine.Mesh.Geometry: LocalGeometry
 
 struct VarsTestModel{dim} <: BalanceLaw end
 
-vars_state_conservative(::VarsTestModel, T) = @vars(x::T, coord::SVector{3, T})
-vars_state_auxiliary(m::VarsTestModel, T) =
+vars_state(::VarsTestModel, ::Conservative, T) =
+    @vars(x::T, coord::SVector{3, T})
+vars_state(m::VarsTestModel, ::Auxiliary, T) =
     @vars(coord::SVector{3, T}, polynomial::T)
-vars_state_gradient_flux(m::VarsTestModel, T) = @vars()
+vars_state(m::VarsTestModel, ::GradientFlux, T) = @vars()
 
 flux_first_order!(::VarsTestModel, _...) = nothing
 flux_second_order!(::VarsTestModel, _...) = nothing
diff --git a/test/Numerics/Mesh/filter.jl b/test/Numerics/Mesh/filter.jl
index 119925b4e65..6a887a78509 100644
--- a/test/Numerics/Mesh/filter.jl
+++ b/test/Numerics/Mesh/filter.jl
@@ -4,6 +4,7 @@ using ClimateMachine.VariableTemplates: @vars, Vars
 using ClimateMachine.Mesh.Grids:
     EveryDirection, HorizontalDirection, VerticalDirection
 using ClimateMachine.MPIStateArrays: weightedsum
+using ClimateMachine.BalanceLaws
 
 import GaussQuadrature
 using MPI
@@ -101,7 +102,8 @@ ClimateMachine.init()
 end
 
 struct FilterTestModel{N} <: ClimateMachine.BalanceLaws.BalanceLaw end
-ClimateMachine.BalanceLaws.vars_state_auxiliary(::FilterTestModel, FT) = @vars()
+ClimateMachine.BalanceLaws.vars_state(::FilterTestModel, ::Auxiliary, FT) =
+    @vars()
 ClimateMachine.BalanceLaws.init_state_auxiliary!(::FilterTestModel, _...) =
     nothing
 
@@ -121,8 +123,11 @@ filtered(::VerticalDirection, dim, x, y, z) =
 filtered(::HorizontalDirection, dim, x, y, z) =
     (dim == 2) ? l2(x) * l3(y) + l3(x) : l2(x) * l3(y) + l3(x) + l2(y)
 
-ClimateMachine.BalanceLaws.vars_state_conservative(::FilterTestModel{4}, FT) =
-    @vars(q1::FT, q2::FT, q3::FT, q4::FT)
+ClimateMachine.BalanceLaws.vars_state(
+    ::FilterTestModel{4},
+    ::Conservative,
+    FT,
+) = @vars(q1::FT, q2::FT, q3::FT, q4::FT)
 function ClimateMachine.BalanceLaws.init_state_conservative!(
     ::FilterTestModel{4},
     state::Vars,
@@ -209,8 +214,9 @@ end
     end
 end
 
-ClimateMachine.BalanceLaws.vars_state_conservative(
+ClimateMachine.BalanceLaws.vars_state(
     ::FilterTestModel{1},
+    ::Conservative,
     FT,
 ) where {N} = @vars(q::FT)
 function ClimateMachine.BalanceLaws.init_state_conservative!(
diff --git a/test/Numerics/Mesh/interpolation.jl b/test/Numerics/Mesh/interpolation.jl
index 24372144966..ec41959316b 100644
--- a/test/Numerics/Mesh/interpolation.jl
+++ b/test/Numerics/Mesh/interpolation.jl
@@ -14,7 +14,7 @@ using ClimateMachine
 ClimateMachine.init()
 using ClimateMachine.ConfigTypes
 using ClimateMachine.Atmos
-using ClimateMachine.Atmos: vars_state_conservative, vars_state_auxiliary
+using ClimateMachine.Atmos: vars_state
 using ClimateMachine.Orientations
 using ClimateMachine.DGMethods
 using ClimateMachine.DGMethods.NumericalFluxes
diff --git a/test/Numerics/SystemSolvers/bandedsystem.jl b/test/Numerics/SystemSolvers/bandedsystem.jl
index 3a9a782a6d0..17a7fc99d65 100644
--- a/test/Numerics/SystemSolvers/bandedsystem.jl
+++ b/test/Numerics/SystemSolvers/bandedsystem.jl
@@ -7,11 +7,9 @@ using Logging
 using LinearAlgebra
 using Random
 using StaticArrays
-using ClimateMachine.BalanceLaws: BalanceLaw
-import ClimateMachine.BalanceLaws:
-    vars_state_conservative, number_state_conservative
-using ClimateMachine.DGMethods:
-    DGModel, init_ode_state, create_conservative_state
+using ClimateMachine.BalanceLaws
+import ClimateMachine.BalanceLaws: vars_state, number_states
+using ClimateMachine.DGMethods: DGModel, init_ode_state, create_state
 using ClimateMachine.SystemSolvers: banded_matrix, banded_matrix_vector_product!
 using ClimateMachine.DGMethods.NumericalFluxes:
     RusanovNumericalFlux,
@@ -41,7 +39,7 @@ function init_velocity_diffusion!(
 end
 
 struct BigAdvectionDiffusion <: BalanceLaw end
-function vars_state_conservative(::BigAdvectionDiffusion, FT)
+function vars_state(::BigAdvectionDiffusion, ::Conservative, FT)
     @vars begin
         ρ::FT
         X::SVector{3, FT}
@@ -168,10 +166,16 @@ let
                         atol = 100 * eps(FT),
                     ))
 
-                    big_Q =
-                        create_conservative_state(BigAdvectionDiffusion(), grid)
-                    big_dQ =
-                        create_conservative_state(BigAdvectionDiffusion(), grid)
+                    big_Q = create_state(
+                        BigAdvectionDiffusion(),
+                        grid,
+                        Conservative(),
+                    )
+                    big_dQ = create_state(
+                        BigAdvectionDiffusion(),
+                        grid,
+                        Conservative(),
+                    )
 
                     big_Q .= NaN
                     big_dQ .= NaN
diff --git a/test/Numerics/SystemSolvers/poisson.jl b/test/Numerics/SystemSolvers/poisson.jl
index f8c76da02d6..d2ba3288074 100644
--- a/test/Numerics/SystemSolvers/poisson.jl
+++ b/test/Numerics/SystemSolvers/poisson.jl
@@ -11,12 +11,11 @@ using ClimateMachine.DGMethods.NumericalFluxes
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.VariableTemplates
 using ClimateMachine.DGMethods
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, Gradient, GradientFlux
+
 import ClimateMachine.DGMethods:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     flux_first_order!,
     flux_second_order!,
     source!,
@@ -41,10 +40,10 @@ end
 
 struct PoissonModel{dim} <: BalanceLaw end
 
-vars_state_auxiliary(::PoissonModel, T) = @vars(rhs_ϕ::T)
-vars_state_conservative(::PoissonModel, T) = @vars(ϕ::T)
-vars_state_gradient(::PoissonModel, T) = @vars(ϕ::T)
-vars_state_gradient_flux(::PoissonModel, T) = @vars(∇ϕ::SVector{3, T})
+vars_state(::PoissonModel, ::Auxiliary, T) = @vars(rhs_ϕ::T)
+vars_state(::PoissonModel, ::Conservative, T) = @vars(ϕ::T)
+vars_state(::PoissonModel, ::Gradient, T) = @vars(ϕ::T)
+vars_state(::PoissonModel, ::GradientFlux, T) = @vars(∇ϕ::SVector{3, T})
 
 boundary_state!(nf, bl::PoissonModel, _...) = nothing
 
diff --git a/test/Ocean/HydrostaticBoussinesq/test_divergence_free.jl b/test/Ocean/HydrostaticBoussinesq/test_divergence_free.jl
index a44d7b47cac..40ac27238f7 100644
--- a/test/Ocean/HydrostaticBoussinesq/test_divergence_free.jl
+++ b/test/Ocean/HydrostaticBoussinesq/test_divergence_free.jl
@@ -6,7 +6,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, Conservative
 using ClimateMachine.HydrostaticBoussinesq
 
 using CLIMAParameters
@@ -89,11 +89,11 @@ function test_divergence_free(; 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, Conservative(), 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, Conservative(), FT)}(minimum(
         solver_config.Q,
         dims = (1, 3),
     ))
diff --git a/test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl b/test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl
index 5902b9954f7..5058d09bc6f 100644
--- a/test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl
+++ b/test/Ocean/ShallowWater/2D_hydrostatic_spindown.jl
@@ -13,7 +13,7 @@ using ClimateMachine.ShallowWater
 using ClimateMachine.Mesh.Topologies
 using ClimateMachine.Mesh.Grids
 using ClimateMachine.DGMethods
-using ClimateMachine.BalanceLaws: vars_state_conservative, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: vars_state, Conservative, Auxiliary
 using ClimateMachine.DGMethods.NumericalFluxes
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.VTK
@@ -186,8 +186,9 @@ function make_callbacks(
             step
         )
         @info "doing VTK output" outprefix
-        statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
-        auxnames = flattenednames(vars_state_auxiliary(model, eltype(Q)))
+        statenames =
+            flattenednames(vars_state(model, Conservative(), eltype(Q)))
+        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
         writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)
     end
 
diff --git a/test/Ocean/ShallowWater/GyreDriver.jl b/test/Ocean/ShallowWater/GyreDriver.jl
index a4740074635..aa1cda545fe 100644
--- a/test/Ocean/ShallowWater/GyreDriver.jl
+++ b/test/Ocean/ShallowWater/GyreDriver.jl
@@ -119,13 +119,15 @@ function run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
 
     if test > 2
         outprefix = @sprintf("ic_mpirank%04d_ic", MPI.Comm_rank(mpicomm))
-        statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
-        auxnames = flattenednames(vars_state_auxiliary(model, eltype(Q)))
+        statenames =
+            flattenednames(vars_state(model, Conservative(), eltype(Q)))
+        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
         writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)
 
         outprefix = @sprintf("exact_mpirank%04d", MPI.Comm_rank(mpicomm))
-        statenames = flattenednames(vars_state_conservative(model, eltype(Qe)))
-        auxnames = flattenednames(vars_state_auxiliary(model, eltype(Qe)))
+        statenames =
+            flattenednames(vars_state(model, Conservative(), eltype(Qe)))
+        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Qe)))
         writevtk(outprefix, Qe, dg, statenames, dg.state_auxiliary, auxnames)
 
         step = [0]
@@ -140,9 +142,9 @@ function run(mpicomm, topl, ArrayType, N, dt, FT, model, test)
             )
             @debug "doing VTK output" outprefix
             statenames =
-                flattenednames(vars_state_conservative(model, eltype(Q)))
+                flattenednames(vars_state(model, Conservative(), eltype(Q)))
             auxiliarynames =
-                flattenednames(vars_state_auxiliary(model, eltype(Q)))
+                flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
             writevtk(
                 outprefix,
                 Q,
diff --git a/test/Ocean/ShallowWater/GyreInABox.jl b/test/Ocean/ShallowWater/GyreInABox.jl
index 10cdcf4ca8d..1f8e54aca0b 100644
--- a/test/Ocean/ShallowWater/GyreInABox.jl
+++ b/test/Ocean/ShallowWater/GyreInABox.jl
@@ -16,8 +16,7 @@ using ClimateMachine.VTK
 import ClimateMachine.ShallowWater:
     shallow_init_state!,
     shallow_init_aux!,
-    vars_state_conservative,
-    vars_state_auxiliary,
+    vars_state,
     shallow_boundary_state!,
     TurbulenceClosure,
     LinearDrag,
diff --git a/test/Ocean/SplitExplicit/hydrostatic_spindown.jl b/test/Ocean/SplitExplicit/hydrostatic_spindown.jl
index 34a38342aa1..b599d872a8d 100644
--- a/test/Ocean/SplitExplicit/hydrostatic_spindown.jl
+++ b/test/Ocean/SplitExplicit/hydrostatic_spindown.jl
@@ -12,7 +12,7 @@ using ClimateMachine.ShallowWater
 
 using ClimateMachine.Mesh.Topologies
 using ClimateMachine.Mesh.Grids
-using ClimateMachine.BalanceLaws: vars_state_conservative, vars_state_auxiliary
+using ClimateMachine.BalanceLaws: vars_state, Conservative, Auxiliary
 using ClimateMachine.DGMethods
 using ClimateMachine.DGMethods.NumericalFluxes
 using ClimateMachine.MPIStateArrays
@@ -248,8 +248,9 @@ function make_callbacks(
             step
         )
         @info "doing VTK output" outprefix
-        statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
-        auxnames = flattenednames(vars_state_auxiliary(model, eltype(Q)))
+        statenames =
+            flattenednames(vars_state(model, Conservative(), eltype(Q)))
+        auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))
         writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)
     end
 
diff --git a/tutorials/Atmos/burgers_single_stack.jl b/tutorials/Atmos/burgers_single_stack.jl
index bbbb589c47c..6a3e91e9f4c 100644
--- a/tutorials/Atmos/burgers_single_stack.jl
+++ b/tutorials/Atmos/burgers_single_stack.jl
@@ -69,7 +69,6 @@ using ClimateMachine.Mesh.Grids
 using ClimateMachine.Writers
 using ClimateMachine.DGMethods
 using ClimateMachine.DGMethods.NumericalFluxes
-using ClimateMachine.BalanceLaws: BalanceLaw
 using ClimateMachine.Mesh.Geometry: LocalGeometry
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.GenericCallbacks
@@ -79,12 +78,9 @@ using ClimateMachine.SingleStackUtils
 
 #  - import necessary ClimateMachine modules: (`import`ing enables us to
 #  provide implementations of these structs/methods)
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     source!,
     flux_second_order!,
     flux_first_order!,
@@ -155,18 +151,19 @@ m = BurgersEquation{FT}();
 # by the solver.
 
 # Specify auxiliary variables for `BurgersEquation`
-vars_state_auxiliary(::BurgersEquation, FT) = @vars(z::FT, T::FT);
+vars_state(::BurgersEquation, ::Auxiliary, FT) = @vars(z::FT, T::FT);
 
 # Specify state variables, the variables solved for in the PDEs, for
 # `BurgersEquation`
-vars_state_conservative(::BurgersEquation, FT) =
+vars_state(::BurgersEquation, ::Conservative, FT) =
     @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);
 
 # Specify state variables whose gradients are needed for `BurgersEquation`
-vars_state_gradient(::BurgersEquation, FT) = @vars(u::SVector{3, FT}, ρcT::FT);
+vars_state(::BurgersEquation, ::Gradient, FT) =
+    @vars(u::SVector{3, FT}, ρcT::FT);
 
 # Specify gradient variables for `BurgersEquation`
-vars_state_gradient_flux(::BurgersEquation, FT) =
+vars_state(::BurgersEquation, ::GradientFlux, FT) =
     @vars(μ∇u::SMatrix{3, 3, FT, 9}, α∇ρcT::SVector{3, FT});
 
 # ## Define the compute kernels
@@ -175,7 +172,7 @@ vars_state_gradient_flux(::BurgersEquation, FT) =
 # `init_state_conservative!`. Note that
 # - this method is only called at `t=0`
 # - `aux.z` and `aux.T` are available here because we've specified `z` and `T`
-# in `vars_state_auxiliary`
+# in `vars_state`
 function init_state_auxiliary!(
     m::BurgersEquation,
     aux::Vars,
@@ -188,7 +185,7 @@ end;
 # Specify the initial values in `state::Vars`. Note that
 # - this method is only called at `t=0`
 # - `state.ρ`, `state.ρu` and`state.ρcT` are available here because
-# we've specified `ρ`, `ρu` and `ρcT` in `vars_state_conservative`
+# we've specified `ρ`, `ρu` and `ρcT` in `vars_state`
 function init_state_conservative!(
     m::BurgersEquation,
     state::Vars,
@@ -227,7 +224,7 @@ end;
 
 # Compute/update all auxiliary variables at each node. Note that
 # - `aux.T` is available here because we've specified `T` in
-# `vars_state_auxiliary`
+# `vars_state`
 function heat_eq_nodal_update_aux!(
     m::BurgersEquation,
     state::Vars,
@@ -240,7 +237,7 @@ end;
 # Since we have second-order fluxes, we must tell `ClimateMachine` to compute
 # the gradient of `ρcT` and `u`. Here, we specify how `ρcT`, `u` are computed. Note that
 # `transform.ρcT` and `transform.u` are available here because we've specified `ρcT`
-# and `u`in `vars_state_gradient`
+# and `u`in `vars_state`
 function compute_gradient_argument!(
     m::BurgersEquation,
     transform::Vars,
@@ -255,9 +252,9 @@ end;
 # Specify where in `diffusive::Vars` to store the computed gradient from
 # `compute_gradient_argument!`. Note that:
 #  - `diffusive.μ∇u` is available here because we've specified `μ∇u` in
-#  `vars_state_gradient_flux`
+#  `vars_state`
 #  - `∇transform.u` is available here because we've specified `u` in
-#  `vars_state_gradient`
+#  `vars_state`
 #  - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor
 function compute_gradient_flux!(
     m::BurgersEquation,
@@ -297,7 +294,7 @@ end;
 # Compute advective flux.
 # Note that:
 # - `state.ρu` is available here because we've specified `ρu` in
-# `vars_state_conservative`
+# `vars_state`
 function flux_first_order!(
     m::BurgersEquation,
     flux::Grad,
@@ -316,7 +313,7 @@ end;
 # Compute diffusive flux (e.g. ``F(μ, \mathbf{u}, t) = -μ∇\mathbf{u}`` in the original PDE).
 # Note that:
 # - `diffusive.μ∇u` is available here because we've specified `μ∇u` in
-# `vars_state_gradient_flux`
+# `vars_state`
 function flux_second_order!(
     m::BurgersEquation,
     flux::Grad,
@@ -450,14 +447,14 @@ z = get_z(driver_config.grid, z_scale)
 state_vars = get_vars_from_nodal_stack(
     driver_config.grid,
     solver_config.Q,
-    vars_state_conservative(m, FT),
+    vars_state(m, Conservative(), FT),
     i = 1,
     j = 1,
 );
 aux_vars = get_vars_from_nodal_stack(
     driver_config.grid,
     solver_config.dg.state_auxiliary,
-    vars_state_auxiliary(m, FT),
+    vars_state(m, Auxiliary(), FT),
     i = 1,
     j = 1,
     exclude = [z_key],
@@ -493,13 +490,13 @@ export_plot_snapshot(
 state_vars_var = get_horizontal_variance(
     driver_config.grid,
     solver_config.Q,
-    vars_state_conservative(m, FT),
+    vars_state(m, Conservative(), FT),
 );
 
 state_vars_avg = get_horizontal_mean(
     driver_config.grid,
     solver_config.Q,
-    vars_state_conservative(m, FT),
+    vars_state(m, Conservative(), FT),
 );
 
 export_plot_snapshot(
@@ -545,12 +542,12 @@ callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
     state_vars_var = get_horizontal_variance(
         driver_config.grid,
         solver_config.Q,
-        vars_state_conservative(m, FT),
+        vars_state(m, Conservative(), FT),
     )
     state_vars_avg = get_horizontal_mean(
         driver_config.grid,
         solver_config.Q,
-        vars_state_conservative(m, FT),
+        vars_state(m, Conservative(), FT),
     )
     step[1] += 1
     push!(data_var, state_vars_var)
diff --git a/tutorials/Land/Heat/heat_equation.jl b/tutorials/Land/Heat/heat_equation.jl
index b53a43ec13b..b3c1d37d354 100644
--- a/tutorials/Land/Heat/heat_equation.jl
+++ b/tutorials/Land/Heat/heat_equation.jl
@@ -61,7 +61,9 @@ using ClimateMachine.Mesh.Topologies
 using ClimateMachine.Mesh.Grids
 using ClimateMachine.DGMethods
 using ClimateMachine.DGMethods.NumericalFluxes
-using ClimateMachine.BalanceLaws: BalanceLaw
+using ClimateMachine.BalanceLaws:
+    BalanceLaw, Conservative, Auxiliary, Gradient, GradientFlux
+
 using ClimateMachine.Mesh.Geometry: LocalGeometry
 using ClimateMachine.MPIStateArrays
 using ClimateMachine.GenericCallbacks
@@ -72,10 +74,7 @@ using ClimateMachine.SingleStackUtils
 #  - import necessary ClimateMachine modules: (`import`ing enables us to
 #  provide implementations of these structs/methods)
 import ClimateMachine.BalanceLaws:
-    vars_state_auxiliary,
-    vars_state_conservative,
-    vars_state_gradient,
-    vars_state_gradient_flux,
+    vars_state,
     source!,
     flux_second_order!,
     flux_first_order!,
@@ -136,17 +135,17 @@ m = HeatModel{FT}();
 # the solver.
 
 # Specify auxiliary variables for `HeatModel`
-vars_state_auxiliary(::HeatModel, FT) = @vars(z::FT, T::FT);
+vars_state(::HeatModel, ::Auxiliary, FT) = @vars(z::FT, T::FT);
 
 # Specify state variables, the variables solved for in the PDEs, for
 # `HeatModel`
-vars_state_conservative(::HeatModel, FT) = @vars(ρcT::FT);
+vars_state(::HeatModel, ::Conservative, FT) = @vars(ρcT::FT);
 
 # Specify state variables whose gradients are needed for `HeatModel`
-vars_state_gradient(::HeatModel, FT) = @vars(ρcT::FT);
+vars_state(::HeatModel, ::Gradient, FT) = @vars(ρcT::FT);
 
 # Specify gradient variables for `HeatModel`
-vars_state_gradient_flux(::HeatModel, FT) = @vars(α∇ρcT::SVector{3, FT});
+vars_state(::HeatModel, ::GradientFlux, FT) = @vars(α∇ρcT::SVector{3, FT});
 
 # ## Define the compute kernels
 
@@ -154,7 +153,7 @@ vars_state_gradient_flux(::HeatModel, FT) = @vars(α∇ρcT::SVector{3, FT});
 # `init_state_conservative!`. Note that
 # - this method is only called at `t=0`
 # - `aux.z` and `aux.T` are available here because we've specified `z` and `T`
-# in `vars_state_auxiliary`
+# in `vars_state`
 function init_state_auxiliary!(m::HeatModel, aux::Vars, geom::LocalGeometry)
     aux.z = geom.coord[3]
     aux.T = m.initialT
@@ -163,7 +162,7 @@ end;
 # Specify the initial values in `state::Vars`. Note that
 # - this method is only called at `t=0`
 # - `state.ρcT` is available here because we've specified `ρcT` in
-# `vars_state_conservative`
+# `vars_state`
 function init_state_conservative!(
     m::HeatModel,
     state::Vars,
@@ -193,7 +192,7 @@ end;
 
 # Compute/update all auxiliary variables at each node. Note that
 # - `aux.T` is available here because we've specified `T` in
-# `vars_state_auxiliary`
+# `vars_state`
 function heat_eq_nodal_update_aux!(
     m::HeatModel,
     state::Vars,
@@ -206,7 +205,7 @@ end;
 # Since we have second-order fluxes, we must tell `ClimateMachine` to compute
 # the gradient of `ρcT`. Here, we specify how `ρcT` is computed. Note that
 #  - `transform.ρcT` is available here because we've specified `ρcT` in
-#  `vars_state_gradient`
+#  `vars_state`
 function compute_gradient_argument!(
     m::HeatModel,
     transform::Vars,
@@ -220,9 +219,9 @@ end;
 # Specify where in `diffusive::Vars` to store the computed gradient from
 # `compute_gradient_argument!`. Note that:
 #  - `diffusive.α∇ρcT` is available here because we've specified `α∇ρcT` in
-#  `vars_state_gradient_flux`
+#  `vars_state`
 #  - `∇transform.ρcT` is available here because we've specified `ρcT`  in
-#  `vars_state_gradient`
+#  `vars_state`
 function compute_gradient_flux!(
     m::HeatModel,
     diffusive::Vars,
@@ -241,7 +240,7 @@ function flux_first_order!(m::HeatModel, _...) end;
 # Compute diffusive flux (``F(α, ρcT, t) = -α ∇ρcT`` in the original PDE).
 # Note that:
 # - `diffusive.α∇ρcT` is available here because we've specified `α∇ρcT` in
-# `vars_state_gradient_flux`
+# `vars_state`
 function flux_second_order!(
     m::HeatModel,
     flux::Grad,
@@ -369,12 +368,12 @@ z = get_z(grid, z_scale)
 state_vars = SingleStackUtils.get_vars_from_nodal_stack(
     grid,
     Q,
-    vars_state_conservative(m, FT),
+    vars_state(m, Conservative(), FT),
 )
 aux_vars = SingleStackUtils.get_vars_from_nodal_stack(
     grid,
     aux,
-    vars_state_auxiliary(m, FT),
+    vars_state(m, Auxiliary(), FT),
 )
 all_vars = OrderedDict(state_vars..., aux_vars...);
 export_plot_snapshot(
@@ -410,12 +409,12 @@ callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do
     state_vars = SingleStackUtils.get_vars_from_nodal_stack(
         grid,
         Q,
-        vars_state_conservative(m, FT),
+        vars_state(m, Conservative(), FT),
     )
     aux_vars = SingleStackUtils.get_vars_from_nodal_stack(
         grid,
         aux,
-        vars_state_auxiliary(m, FT);
+        vars_state(m, Auxiliary(), FT);
         exclude = [z_key],
     )
     all_vars = OrderedDict(state_vars..., aux_vars...)
diff --git a/tutorials/Numerics/DGMethods/nonnegative.jl b/tutorials/Numerics/DGMethods/nonnegative.jl
index aee089a4974..937310c9ff3 100644
--- a/tutorials/Numerics/DGMethods/nonnegative.jl
+++ b/tutorials/Numerics/DGMethods/nonnegative.jl
@@ -92,7 +92,7 @@ function do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, testname)
         vtkstep
     )
 
-    statenames = flattenednames(vars_state_conservative(model, eltype(Q)))
+    statenames = flattenednames(vars_state(model, Conservative(), eltype(Q)))
 
     writevtk(filename, Q, dg, statenames)