From 74826a896e9cb102a3fb601a0f93662f1bf24c66 Mon Sep 17 00:00:00 2001 From: Charles Kawczynski Date: Tue, 23 Jun 2020 11:39:49 -0700 Subject: [PATCH] Add types to BalanceLaw, and unify several funcs: - Unifies `vars_state` (BalanceLaws) into a single function - Unifies `number_states` (BalanceLaws) into a single function - Unifies `create_state` (DGMethods) into a single function - Unifies `extract_state` (Diagnostics) into a single function --- docs/src/APIs/BalanceLaws/BalanceLaws.md | 7 +- .../DGMethods/how_to_make_a_balance_law.md | 14 +- docs/src/Theory/Atmos/Model/tracers.md | 5 +- docs/src/Theory/Common/Turbulence.md | 6 +- experiments/AtmosLES/bomex.jl | 2 +- experiments/AtmosLES/dycoms.jl | 20 +- experiments/AtmosLES/taylor-green.jl | 5 +- experiments/OceanBoxGCM/homogeneous_box.jl | 6 +- src/Atmos/Model/AtmosModel.jl | 102 ++-- src/Atmos/Model/filters.jl | 2 +- src/Atmos/Model/hyperdiffusion.jl | 16 +- src/Atmos/Model/linear.jl | 17 +- src/Atmos/Model/moisture.jl | 15 +- src/Atmos/Model/precipitation.jl | 13 +- src/Atmos/Model/radiation.jl | 5 +- src/Atmos/Model/ref_state.jl | 8 +- src/Atmos/Model/tracers.jl | 13 +- src/BalanceLaws/BalanceLaws.jl | 7 +- src/BalanceLaws/interface.jl | 79 +-- src/BalanceLaws/state_types.jl | 21 + src/Common/Orientations/Orientations.jl | 8 +- .../TurbulenceClosures/TurbulenceClosures.jl | 56 +-- src/Diagnostics/Diagnostics.jl | 8 +- src/Diagnostics/atmos_gcm_default.jl | 13 +- src/Diagnostics/atmos_les_core.jl | 6 +- src/Diagnostics/atmos_les_default.jl | 8 +- .../atmos_refstate_perturbations.jl | 21 +- src/Diagnostics/diagnostic_fields.jl | 4 +- src/Diagnostics/dump_aux.jl | 37 +- src/Diagnostics/dump_init.jl | 29 ++ src/Diagnostics/dump_state.jl | 35 +- src/Diagnostics/groups.jl | 1 + src/Diagnostics/helpers.jl | 35 +- src/Driver/Driver.jl | 6 +- src/Numerics/DGMethods/DGMethods.jl | 29 +- src/Numerics/DGMethods/DGModel.jl | 30 +- src/Numerics/DGMethods/DGModel_kernels.jl | 454 +++++++++--------- src/Numerics/DGMethods/NumericalFluxes.jl | 18 +- src/Numerics/DGMethods/create_states.jl | 140 +----- src/Numerics/DGMethods/remainder.jl | 82 ++-- src/Numerics/SystemSolvers/SystemSolvers.jl | 3 +- .../SystemSolvers/columnwise_lu_solver.jl | 4 +- .../HydrostaticBoussinesqModel.jl | 35 +- .../HydrostaticBoussinesq/LinearHBModel.jl | 14 +- src/Ocean/ShallowWater/ShallowWaterModel.jl | 16 +- src/Utilities/Callbacks/Callbacks.jl | 6 +- .../SingleStackUtils/SingleStackUtils.jl | 6 +- .../Parameterizations/Microphysics/KM_ice.jl | 34 +- .../Microphysics/KM_saturation_adjustment.jl | 18 +- .../Microphysics/KM_warm_rain.jl | 28 +- .../Microphysics/KinematicModel.jl | 12 +- test/Diagnostics/diagnostic_fields_test.jl | 7 +- .../DGMethods/Euler/acousticwave_1d_imex.jl | 7 +- .../DGMethods/Euler/acousticwave_mrigark.jl | 7 +- .../DGMethods/Euler/isentropicvortex.jl | 5 +- .../DGMethods/Euler/isentropicvortex_imex.jl | 9 +- .../Euler/isentropicvortex_mrigark.jl | 9 +- .../isentropicvortex_mrigark_implicit.jl | 9 +- .../Euler/isentropicvortex_multirate.jl | 9 +- .../advection_diffusion_model.jl | 42 +- ...advection_diffusion_model_1dimex_bgmres.jl | 2 +- .../advection_diffusion/advection_sphere.jl | 2 +- .../hyperdiffusion_model.jl | 28 +- .../periodic_3D_hyperdiffusion.jl | 2 +- .../pseudo1D_advection_diffusion.jl | 2 +- .../pseudo1D_advection_diffusion_1dimex.jl | 2 +- ...1D_advection_diffusion_mrigark_implicit.jl | 2 +- .../density_current_model.jl | 6 +- .../compressible_Navier_Stokes/mms_model.jl | 18 +- .../compressible_Navier_Stokes/ref_state.jl | 7 +- .../Numerics/DGMethods/conservation/sphere.jl | 17 +- test/Numerics/DGMethods/courant.jl | 3 +- test/Numerics/DGMethods/integral_test.jl | 22 +- .../DGMethods/integral_test_sphere.jl | 32 +- test/Numerics/DGMethods/remainder_model.jl | 20 +- test/Numerics/DGMethods/vars_test.jl | 16 +- test/Numerics/Mesh/filter.jl | 14 +- test/Numerics/Mesh/interpolation.jl | 2 +- test/Numerics/SystemSolvers/bandedsystem.jl | 24 +- test/Numerics/SystemSolvers/poisson.jl | 17 +- .../test_divergence_free.jl | 6 +- .../ShallowWater/2D_hydrostatic_spindown.jl | 7 +- test/Ocean/ShallowWater/GyreDriver.jl | 14 +- test/Ocean/ShallowWater/GyreInABox.jl | 3 +- .../SplitExplicit/hydrostatic_spindown.jl | 7 +- .../SingleStackUtils/horizontal_stats_test.jl | 28 +- tutorials/Atmos/burgers_single_stack.jl | 45 +- tutorials/Land/Heat/heat_equation.jl | 39 +- tutorials/Numerics/DGMethods/nonnegative.jl | 2 +- 89 files changed, 891 insertions(+), 1161 deletions(-) create mode 100644 src/BalanceLaws/state_types.jl create mode 100644 src/Diagnostics/dump_init.jl 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 161e23d3622..d3c70f17c30 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 @@ -119,7 +119,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. @@ -164,7 +164,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/test/Utilities/SingleStackUtils/horizontal_stats_test.jl b/test/Utilities/SingleStackUtils/horizontal_stats_test.jl index 8b4b85fa888..810c702ce05 100644 --- a/test/Utilities/SingleStackUtils/horizontal_stats_test.jl +++ b/test/Utilities/SingleStackUtils/horizontal_stats_test.jl @@ -12,13 +12,15 @@ using ClimateMachine.DGMethods.Grids using ClimateMachine.MPIStateArrays using ClimateMachine.VariableTemplates using ClimateMachine.SingleStackUtils -using ClimateMachine.DGMethods: BalanceLaw, LocalGeometry - -import ClimateMachine.DGMethods: - vars_state_auxiliary, - vars_state_conservative, - vars_state_gradient, - vars_state_gradient_flux, +using ClimateMachine.DGMethods: LocalGeometry + +using ClimateMachine.BalanceLaws: BalanceLaw, + Auxiliary, + Conservative, + Gradient, + GradientFlux +import ClimateMachine.BalanceLaws: + vars_state, init_state_auxiliary!, init_state_conservative! @@ -32,10 +34,10 @@ end EmptyBalLaw(param_set, zmax) = EmptyBalLaw{typeof(zmax), typeof(param_set)}(param_set, zmax) -vars_state_auxiliary(::EmptyBalLaw, FT) = @vars(x::FT, y::FT, z::FT) -vars_state_conservative(::EmptyBalLaw, FT) = @vars(ρ::FT) -vars_state_gradient(::EmptyBalLaw, FT) = @vars() -vars_state_gradient_flux(::EmptyBalLaw, FT) = @vars() +vars_state(::EmptyBalLaw, ::Auxiliary, FT) = @vars(x::FT, y::FT, z::FT) +vars_state(::EmptyBalLaw, ::Conservative, FT) = @vars(ρ::FT) +vars_state(::EmptyBalLaw, ::Gradient, FT) = @vars() +vars_state(::EmptyBalLaw, ::GradientFlux, FT) = @vars() function init_state_auxiliary!(m::EmptyBalLaw, aux::Vars, geom::LocalGeometry) aux.x = geom.coord[1] @@ -145,12 +147,12 @@ function main() test_hmean( driver_config.grid, solver_config.Q, - vars_state_conservative(m, FT), + vars_state(m, Conservative(), FT), ) test_hvar( driver_config.grid, solver_config.Q, - vars_state_conservative(m, FT), + vars_state(m, Conservative(), FT), ) end 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)