[ { "path": ".buildkite/docs-pipeline.yml", "content": "env:\n JULIA_VERSION: \"1.5.4\"\n GKSwstype: nul\n OPENBLAS_NUM_THREADS: 1\n CLIMATEMACHINE_SETTINGS_DISABLE_GPU: \"true\"\n CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: \"true\"\n CLIMATEMACHINE_SETTINGS_DISABLE_CUSTOM_LOGGER: \"true\"\n\nsteps:\n - label: \"Build project\"\n command:\n - \"julia --project --color=yes -e 'using Pkg; Pkg.instantiate()'\"\n - \"julia --project=docs/ --color=yes -e 'using Pkg; Pkg.instantiate()'\"\n - \"julia --project=docs/ --color=yes -e 'using Pkg; Pkg.precompile()'\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n slurm_cpus_per_task: 1\n slurm_mem_per_cpu: 6000\n\n - wait\n\n - label: \"Build docs\"\n command:\n # this extracts out the PR number from the bors message on the trying branch\n # to force documenter to deploy the PR branch number gh-preview\n - \"if [ $$BUILDKITE_BRANCH == \\\"trying\\\" ]; then \\\n export BUILDKITE_PULL_REQUEST=\\\"$${BUILDKITE_MESSAGE//[!0-9]/}\\\"; \\\n fi\"\n - \"if [[ ! -z \\\"$${PULL_REQUEST}\\\" ]]; then \\\n export BUILDKITE_PULL_REQUEST=\\\"$${PULL_REQUEST}\\\"; \\\n fi\"\n - \"julia --project=docs/ --color=yes --procs=10 docs/make.jl\"\n env:\n JULIA_PROJECT: \"docs/\"\n agents:\n config: cpu\n queue: central\n slurm_time: 120\n slurm_nodes: 1\n slurm_ntasks: 10\n slurm_cpus_per_task: 1\n slurm_mem_per_cpu: 6000\n\n" }, { "path": ".buildkite/pipeline.yml", "content": "env:\n JULIA_VERSION: \"1.5.4\"\n OPENMPI_VERSION: \"4.0.4\"\n CUDA_VERSION: \"10.2\"\n OPENBLAS_NUM_THREADS: 1\n CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: \"true\"\n\nsteps:\n - label: \"init cpu env\"\n key: \"init_cpu_env\"\n command:\n - \"echo $JULIA_DEPOT_PATH\"\n - \"julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'\"\n - \"julia --project -e 'using Pkg; Pkg.precompile()'\"\n - \"julia --project -e 'using Pkg; Pkg.status()'\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"init gpu env\"\n key: \"init_gpu_env\"\n command:\n - \"echo $JULIA_DEPOT_PATH\"\n - \"julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'\"\n - \"julia --project -e 'using Pkg; Pkg.precompile()'\"\n # force the initialization of the CUDA runtime\n # as it is lazily loaded by default\n - \"julia --project -e 'using CUDA; CUDA.versioninfo()'\"\n # force the initialization of the CUDA Compiler runtime\n # as it is lazily generated by default\n - \"julia --project -e 'using CUDA; CUDA.precompile_runtime()'\"\n - \"julia --project -e 'using Pkg; Pkg.status()'\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - wait\n\n - label: \"cpu_advection_diffusion_model_1dimex_bgmres\"\n key: \"cpu_advection_diffusion_model_1dimex_bgmres\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_fvm_advection_diffusion_model_1dimex_bjfnks\"\n key: \"cpu_fvm_advection_diffusion_model_1dimex_bjfnks\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion_model_1dimex_bjfnks.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_pseudo1D_advection_diffusion\"\n key: \"cpu_pseudo1D_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_pseudo1D_advection_diffusion_1dimex\"\n key: \"cpu_pseudo1D_advection_diffusion_1dimex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_pseudo1D_advection_diffusion_mrigark_implicit\"\n key: \"cpu_pseudo1D_advection_diffusion_mrigark_implicit\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_pseudo1D_heat_eqn\"\n key: \"cpu_pseudo1D_heat_eqn\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_heat_eqn.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_bickley_jet_2D\"\n key: \"cpu_bickley_jet_2D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/test_bickley_jet.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_periodic_3D_hyperdiffusion\"\n key: \"cpu_periodic_3D_hyperdiffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_hyperdiffusion_bc\"\n key: \"cpu_hyperdiffusion_bc\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_bc.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_diffusion_hyperdiffusion_sphere\"\n key: \"cpu_diffusion_hyperdiffusion_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/diffusion_hyperdiffusion_sphere.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_connect_1d\"\n key: \"cpu_mpi_connect_1d\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect_1d.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 5\n\n - label: \"cpu_mpi_connect_sphere\"\n key: \"cpu_mpi_connect_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect_sphere.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 5\n\n - label: \"cpu_mpi_getpartition\"\n key: \"cpu_mpi_getpartition\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_getpartition.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 5\n\n - label: \"cpu_mpi_sortcolumns4\"\n key: \"cpu_mpi_sortcolumns4\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_sortcolumns.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 4\n\n - label: \"cpu_ode_tests_convergence\"\n key: \"cpu_ode_tests_convergence\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/ODESolvers/ode_tests_convergence.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_ode_tests_basic\"\n key: \"cpu_ode_tests_basic\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/ODESolvers/ode_tests_basic.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_varsindex\"\n key: \"cpu_varsindex\"\n command:\n - \"mpiexec julia --color=yes --project test/Arrays/varsindex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_diagnostic_fields_test\"\n key: \"cpu_diagnostic_fields_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Diagnostics/diagnostic_fields_test.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_vars_test\"\n key: \"cpu_vars_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/vars_test.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_grad_test\"\n key: \"cpu_grad_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/grad_test.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_grad_test_sphere\"\n key: \"cpu_grad_test_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/grad_test_sphere.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_horizontal_integral_test\"\n key: \"cpu_horizontal_integral_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/horizontal_integral_test.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_integral_test\"\n key: \"cpu_integral_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/integral_test.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_integral_test_sphere\"\n key: \"cpu_integral_test_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/integral_test_sphere.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_custom_filter\"\n key: \"cpu_custom_filter\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/custom_filter.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_fv_reconstruction_test\"\n key: \"cpu_fv_reconstruction_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/fv_reconstruction_test.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_remainder_model\"\n key: \"cpu_remainder_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/remainder_model.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_isentropicvortex\"\n key: \"cpu_isentropicvortex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_isentropicvortex_imex\"\n key: \"cpu_isentropicvortex_imex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n \n - label: \"cpu_isentropicvortex_lmars\"\n key: \"cpu_isentropicvortex_lmars\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_lmars.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_isentropicvortex_multirate\"\n key: \"cpu_isentropicvortex_multirate\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_isentropicvortex_mrigark\"\n key: \"cpu_isentropicvortex_mrigark\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_isentropicvortex_mrigark_implicit\"\n key: \"cpu_isentropicvortex_mrigark_implicit\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_acousticwave_1d_imex\"\n key: \"cpu_acousticwave_1d_imex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_acousticwave_mrigark\"\n key: \"cpu_acousticwave_mrigark\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_acousticwave_variable_degree\"\n key: \"cpu_acousticwave_variable_degree\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_variable_degree.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_fvm_balance\"\n key: \"cpu_fvm_balance\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/fvm_balance.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mms_bc_atmos\"\n key: \"cpu_mms_bc_atmos\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_atmos.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mms_bc_dgmodel\"\n key: \"cpu_mms_bc_dgmodel\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_dgmodel.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_density_current_model\"\n key: \"cpu_density_current_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_direction_splitting_advection_diffusion\"\n key: \"cpu_direction_splitting_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/direction_splitting_advection_diffusion.jl --fix-rng-seed\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_variable_degree_advection_diffusion\"\n key: \"cpu_variable_degree_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/variable_degree_advection_diffusion.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_sphere\"\n key: \"cpu_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/conservation/sphere.jl --fix-rng-seed\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_advection_sphere\"\n key: \"cpu_advection_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_advection_sphere\"\n key: \"cpu_fvm_advection_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_sphere.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_swirl\"\n key: \"cpu_fvm_swirl\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_swirl.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_advection\"\n key: \"cpu_fvm_advection\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_advection_diffusion\"\n key: \"cpu_fvm_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_advection_diffusion_periodic\"\n key: \"cpu_fvm_advection_diffusion_periodic\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion_periodic.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_fvm_isentropicvortex\"\n key: \"cpu_fvm_isentropicvortex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/fvm_isentropicvortex.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_gcm_driver_test\"\n key: \"cpu_gcm_driver_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Driver/gcm_driver_test.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_poisson\"\n key: \"cpu_poisson\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/poisson.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_columnwiselu\"\n key: \"cpu_columnwiselu\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/columnwiselu.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_cg\"\n key: \"cpu_cg\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/cg.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_bgmres\"\n key: \"cpu_bgmres\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/bgmres.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_bandedsystem\"\n key: \"cpu_bandedsystem\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/bandedsystem.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_interpolation\"\n key: \"cpu_interpolation\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/interpolation.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_dss_mpi\"\n key: \"cpu_dss_mpi\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/DSS_mpi.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_dss\"\n key: \"cpu_dss\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/DSS.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_GyreDriver\"\n key: \"cpu_GyreDriver\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/ShallowWater/GyreDriver.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_windstress_short\"\n key: \"cpu_test_windstress_short\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_windstress_short.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_ocean_gyre_short\"\n key: \"cpu_test_ocean_gyre_short\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_ocean_gyre_short.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_2D_spindown\"\n key: \"cpu_test_2D_spindown\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/ShallowWater/test_2D_spindown.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_3D_spindown\"\n key: \"cpu_test_3D_spindown\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_3D_spindown.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_vertical_integral_model\"\n key: \"cpu_test_vertical_integral_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_vertical_integral_model.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_spindown_long\"\n key: \"cpu_test_spindown_long\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_spindown_long.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_restart\"\n key: \"cpu_test_restart\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_restart.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_test_coriolis\"\n key: \"cpu_test_coriolis\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_coriolis.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_KM_saturation_adjustment\"\n key: \"cpu_KM_saturation_adjustment\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_KM_warm_rain\"\n key: \"cpu_KM_warm_rain\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_haverkamp_test\"\n key: \"cpu_haverkamp_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/haverkamp_test.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_soil_params\"\n key: \"cpu_soil_params\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/soil_heterogeneity.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_heat_analytic_unit_test\"\n key: \"cpu_heat_analytic_unit_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/heat_analytic_unit_test.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_discrete_hydrostatic_balance\"\n key: \"cpu_discrete_hydrostatic_balance\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Model/discrete_hydrostatic_balance.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_soil_test_bc\"\n key: \"cpu_soil_test_bc\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_bc.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_soil_test_bc_3d\"\n key: \"cpu_soil_test_bc_3d\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_bc_3d.jl \"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"gpu_varsindex\"\n key: \"gpu_varsindex\"\n command:\n - \"mpiexec julia --color=yes --project test/Arrays/varsindex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_variable_templates\"\n key: \"gpu_variable_templates\"\n command:\n - \"mpiexec julia --color=yes --project test/Utilities/VariableTemplates/runtests_gpu.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"cpu_land_overland_flow_vcatchment\"\n key: \"cpu_land_overland_flow_vcatchment\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_overland_flow_vcatchment.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"gpu_diagnostic_fields_test\"\n key: \"gpu_diagnostic_fields_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Diagnostics/diagnostic_fields_test.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_vars_test\"\n key: \"gpu_vars_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/vars_test.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_custom_filter\"\n key: \"gpu_custom_filter\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/custom_filter.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_remainder_model\"\n key: \"gpu_remainder_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/remainder_model.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex\"\n key: \"gpu_isentropicvortex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex_imex\"\n key: \"gpu_isentropicvortex_imex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex_lmars\"\n key: \"gpu_isentropicvortex_lmars\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_lmars.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex_multirate\"\n key: \"gpu_isentropicvortex_multirate\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex_mrigark\"\n key: \"gpu_isentropicvortex_mrigark\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isentropicvortex_mrigark_implicit\"\n key: \"gpu_isentropicvortex_mrigark_implicit\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_acousticwave_1d_imex\"\n key: \"gpu_acousticwave_1d_imex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_acousticwave_mrigark\"\n key: \"gpu_acousticwave_mrigark\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_acousticwave_variable_degree\"\n key: \"gpu_acousticwave_variable_degree\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/acousticwave_variable_degree.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_balance\"\n key: \"gpu_fvm_balance\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/fvm_balance.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_mms_bc_atmos\"\n key: \"gpu_mms_bc_atmos\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_atmos.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_mms_bc_dgmodel\"\n key: \"gpu_mms_bc_dgmodel\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_dgmodel.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_density_current_model\"\n key: \"gpu_density_current_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_direction_splitting_advection_diffusion\"\n key: \"gpu_direction_splitting_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/direction_splitting_advection_diffusion.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_variable_degree_advection_diffusion\"\n key: \"gpu_variable_degree_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/variable_degree_advection_diffusion.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_sphere\"\n key: \"gpu_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/conservation/sphere.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_advection_sphere\"\n key: \"gpu_advection_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_advection_sphere\"\n key: \"gpu_fvm_advection_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_sphere.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_swirl\"\n key: \"gpu_fvm_swirl\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_swirl.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_advection\"\n key: \"gpu_fvm_advection\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_advection_diffusion\"\n key: \"gpu_fvm_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_advection_diffusion_periodic\"\n key: \"gpu_fvm_advection_diffusion_periodic\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion_periodic.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_fvm_isentropicvortex\"\n key: \"gpu_fvm_isentropicvortex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/Euler/fvm_isentropicvortex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_gcm_driver_test\"\n key: \"gpu_gcm_driver_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Driver/gcm_driver_test.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_poisson\"\n key: \"gpu_poisson\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/poisson.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 2\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_columnwiselu\"\n key: \"gpu_columnwiselu\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/columnwiselu.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bandedsystem\"\n key: \"gpu_bandedsystem\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/SystemSolvers/bandedsystem.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"cpu_courant\"\n key: \"cpu_courant\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/courant.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_brickmesh\"\n key: \"cpu_brickmesh\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/BrickMesh.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_elements\"\n key: \"cpu_elements\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/Elements.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_metrics\"\n key: \"cpu_metrics\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/Metrics.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_grids\"\n key: \"cpu_grids\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/Grids.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_topology\"\n key: \"cpu_topology\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/topology.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_grid_integral\"\n key: \"cpu_grid_integral\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/grid_integral.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_filter\"\n key: \"cpu_filter\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/filter.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_filter_tmar\"\n key: \"cpu_filter_tmar\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/filter_TMAR.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"cpu_mpi_centroid\"\n key: \"cpu_mpi_centroid\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_centroid.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_connect_ell\"\n key: \"cpu_mpi_connect_ell\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect_ell.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_mpi_connect\"\n key: \"cpu_mpi_connect\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_connectfull\"\n key: \"cpu_mpi_connectfull\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connectfull.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_connect_stacked\"\n key: \"cpu_mpi_connect_stacked\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect_stacked.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_connect_stacked_3d\"\n key: \"cpu_mpi_connect_stacked_3d\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_connect_stacked_3d.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 2\n\n - label: \"cpu_mpi_getpartition3\"\n key: \"cpu_mpi_getpartition3\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_getpartition.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_partition\"\n key: \"cpu_mpi_partition\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_partition.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 3\n\n - label: \"cpu_mpi_sortcolumns\"\n key: \"cpu_mpi_sortcolumns\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/mpi_sortcolumns.jl\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"gpu_interpolation\"\n key: \"gpu_interpolation\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/interpolation.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dss_mpi\"\n key: \"gpu_dss_mpi\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/DSS_mpi.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dss\"\n key: \"gpu_dss\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/Mesh/DSS.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_GyreDriver\"\n key: \"gpu_GyreDriver\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/ShallowWater/GyreDriver.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_windstress_short\"\n key: \"gpu_test_windstress_short\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_windstress_short.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_ocean_gyre_short\"\n key: \"gpu_test_ocean_gyre_short\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_ocean_gyre_short.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_2D_spindown\"\n key: \"gpu_test_2D_spindown\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/ShallowWater/test_2D_spindown.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_3D_spindown\"\n key: \"gpu_test_3D_spindown\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_3D_spindown.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_vertical_integral_model\"\n key: \"gpu_test_vertical_integral_model\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_vertical_integral_model.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_spindown_long\"\n key: \"gpu_test_spindown_long\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_spindown_long.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_restart\"\n key: \"gpu_test_restart\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_restart.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_coriolis\"\n key: \"gpu_test_coriolis\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/test_coriolis.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_KM_saturation_adjustment\"\n key: \"gpu_KM_saturation_adjustment\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_KM_warm_rain\"\n key: \"gpu_KM_warm_rain\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_KM_ice\"\n key: \"gpu_KM_ice\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Parameterizations/Microphysics/KM_ice.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_pseudo1D_advection_diffusion\"\n key: \"gpu_pseudo1D_advection_diffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl --integration-testing\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_pseudo1D_advection_diffusion_1dimex\"\n key: \"gpu_pseudo1D_advection_diffusion_1dimex\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_pseudo1D_advection_diffusion_mrigark_implicit\"\n key: \"gpu_pseudo1D_advection_diffusion_mrigark_implicit\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_pseudo1D_heat_eqn\"\n key: \"gpu_pseudo1D_heat_eqn\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/pseudo1D_heat_eqn.jl --integration-testing\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bickley_jet_2D\"\n key: \"gpu_bickley_jet_2D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/test_bickley_jet.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bickley_jet_3D\"\n key: \"gpu_bickley_jet_3D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/test_bickley_jet.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cnse_buoyancy_3D\"\n key: \"gpu_cnse_buoyancy_3D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/test_buoyancy.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cnse_sphere_3D\"\n key: \"gpu_cnse_sphere_3D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_sphere.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cnse_sphere_heat_3D\"\n key: \"gpu_cnse_sphere_heat_3D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_heat_equation.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cnse_sphere_balance_3D\"\n key: \"gpu_cnse_sphere_balance_3D\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_hydrostatic_balance.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_periodic_3D_hyperdiffusion\"\n key: \"gpu_periodic_3D_hyperdiffusion\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl --integration-testing\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_hyperdiffusion_bc\"\n key: \"gpu_hyperdiffusion_bc\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_bc.jl --integration-testing\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_diffusion_hyperdiffusion_sphere\"\n key: \"gpu_diffusion_hyperdiffusion_sphere\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/DGMethods/advection_diffusion/diffusion_hyperdiffusion_sphere.jl --integration-testing\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_esdg_baroclinic_wave\"\n key: \"gpu_esdg_baroclinic_wave\"\n command:\n - \"mpiexec julia --color=yes --project test/Numerics/ESDGMethods/DryAtmos/baroclinic_wave.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dry_rayleigh_benard\"\n key: \"gpu_dry_rayleigh_benard\"\n command:\n - \"mpiexec julia --color=yes --project tutorials/Atmos/dry_rayleigh_benard.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dry_risingbubble\"\n key: \"gpu_dry_risingbubble\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/risingbubble.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n \n - label: \"gpu_dry_risingbubble_fvm\"\n key: \"gpu_dry_risingbubble_fvm\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/risingbubble_fvm.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_moist_risingbubble\"\n key: \"gpu_moist_risingbubble\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/risingbubble.jl --with-moisture \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_solid_body_rotation\"\n key: \"gpu_solid_body_rotation\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/solid_body_rotation.jl --diagnostics=default \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n \n - label: \"gpu_solid_body_rotation_fvm\"\n key: \"gpu_solid_body_rotation_fvm\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/solid_body_rotation_fvm.jl --diagnostics=default \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_solid_body_rotation_mountain\"\n key: \"gpu_solid_body_rotation_mountain\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/solid_body_rotation_mountain.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dry_baroclinic_wave\"\n key: \"gpu_dry_baroclinic_wave\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/baroclinic_wave.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n \n - label: \"gpu_dry_baroclinic_wave_fvm\"\n key: \"gpu_dry_baroclinic_wave_fvm\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/baroclinic_wave_fvm.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_moist_baroclinic_wave\"\n key: \"gpu_moist_baroclinic_wave\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/baroclinic_wave.jl --with-moisture \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_heldsuarez\"\n key: \"gpu_heldsuarez\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosGCM/heldsuarez.jl --diagnostics=default --fix-rng-seed \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_nonhydrostatic_gravity_wave\"\n key: \"gpu_nonhydrostatic_gravity_wave\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosGCM/nonhydrostatic_gravity_wave.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_isothermal_zonal_flow\"\n key: \"gpu_isothermal_zonal_flow\"\n command:\n - \"mpiexec julia --color=yes --project experiments/TestCase/isothermal_zonal_flow.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_surfacebubble\"\n key: \"gpu_surfacebubble\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/surfacebubble.jl --diagnostics=default\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dycoms\"\n key: \"gpu_dycoms\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/dycoms.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_dycoms_with_precip\"\n key: \"gpu_dycoms_with_precip\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/dycoms.jl --fix-rng-seed --moisture-model nonequilibrium --precipitation-model rain --sim-time 120 --check-asserts yes\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_squall_eq_no_precip\"\n key: \"gpu_squall_eq_no_precip\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/squall_line.jl --fix-rng-seed --sim-time 600 --check-asserts yes --moisture-model equilibrium --precipitation-model noprecipitation\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_squall_eq_rain\"\n key: \"gpu_squall_eq_rain\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/squall_line.jl --fix-rng-seed --sim-time 600 --check-asserts yes --moisture-model equilibrium --precipitation-model rain\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_squall_neq_rain_snow\"\n key: \"gpu_squall_neq_rain_snow\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/squall_line.jl --fix-rng-seed --sim-time 600 --check-asserts yes --moisture-model nonequilibrium --precipitation-model rainsnow\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bomex_les\"\n key: \"gpu_bomex_les\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/bomex_les.jl --diagnostics=default --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_sbl_les\"\n key: \"gpu_sbl_les\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/stable_bl_les.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cbl_les\"\n key: \"gpu_cbl_les\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/convective_bl_les.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_cfsite\"\n key: \"gpu_cfsite\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/cfsite_hadgem2-a_07_amip.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bomex_single_stack\"\n key: \"gpu_bomex_single_stack\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/bomex_single_stack.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bomex_single_stack_nonequil\"\n key: \"gpu_bomex_single_stack_nonequil\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/bomex_single_stack.jl --moisture-model nonequilibrium --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_bomex_edmf\"\n key: \"gpu_bomex_edmf\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/bomex_edmf.jl --diagnostics=default --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_ekman\"\n key: \"gpu_ekman\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/ekman_layer.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_sbl_edmf\"\n key: \"gpu_sbl_edmf\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/stable_bl_edmf.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"cpu_sbl_edmf_fvm\"\n key: \"cpu_sbl_edmf_fvm\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/stable_bl_edmf_fvm.jl --fix-rng-seed\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"gpu_sbl_an1d\"\n key: \"gpu_sbl_an1d\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/stable_bl_anelastic1d.jl --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"cpu_sbl_edmf_coupled\"\n key: \"cpu_sbl_edmf_coupled\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/stable_bl_coupled_edmf_an1d.jl --fix-rng-seed\"\n agents:\n config: cpu\n queue: central\n slurm_ntasks: 1\n\n - label: \"gpu_bomex_bulk_sfc_flux\"\n key: \"gpu_bomex_bulk_sfc_flux\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/bomex_les.jl --surface-flux=bulk --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_taylor_green\"\n key: \"gpu_taylor_green\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/taylor_green.jl --diagnostics=default --fix-rng-seed\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_rising_bubble_theta_formulation\"\n key: \"gpu_rising_bubble_theta_formulation\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/rising_bubble_theta_formulation.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_rising_bubble_bryan_mrrk\"\n key: \"gpu_rising_bubble_bryan_mrrk\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/rising_bubble_bryan.jl --fast-method=MultirateRungeKutta \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_rising_bubble_bryan_ark\"\n key: \"gpu_rising_bubble_bryan_ark\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/rising_bubble_bryan.jl --fast-method=AdditiveRungeKutta \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_rising_bubble_bryan_mis\"\n key: \"gpu_rising_bubble_bryan_mis\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/rising_bubble_bryan.jl --fast-method=MultirateInfinitesimalStep \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_schar_scalar_advection\"\n key: \"gpu_schar_scalar_advection\"\n command:\n - \"mpiexec julia --color=yes --project experiments/AtmosLES/schar_scalar_advection.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_windstress_long\"\n key: \"gpu_test_windstress_long\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_windstress_long.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_test_ocean_gyre_long\"\n key: \"gpu_test_ocean_gyre_long\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/HydrostaticBoussinesq/test_ocean_gyre_long.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_simple_box_ivd\"\n key: \"gpu_simple_box_ivd\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/simple_box_ivd.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_simple_dbl_gyre\"\n key: \"gpu_simple_dbl_gyre\"\n command:\n - \"mpiexec julia --color=yes --project test/Ocean/SplitExplicit/simple_dbl_gyre.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_discrete_hydrostatic_balance\"\n key: \"gpu_discrete_hydrostatic_balance\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/Model/discrete_hydrostatic_balance.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 3\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_haverkamp_test\"\n key: \"gpu_haverkamp_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/haverkamp_test.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_soil_params\"\n key: \"gpu_soil_params\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/soil_heterogeneity.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_stable_bl_edmf_implicit_test\"\n key: \"gpu_stable_bl_edmf_implicit_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Atmos/EDMF/stable_bl_single_stack_implicit.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_heat_analytic_unit_test\"\n key: \"gpu_heat_analytic_unit_test\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/heat_analytic_unit_test.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_soil_test_bc\"\n key: \"gpu_soil_test_bc\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_bc.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_soil_test_bc_3d\"\n key: \"gpu_soil_test_bc_3d\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_bc_3d.jl \"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_land_overland_flow_vcatchment\"\n key: \"gpu_land_overland_flow_vcatchment\"\n command:\n - \"mpiexec julia --color=yes --project test/Land/Model/test_overland_flow_vcatchment.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n\n - label: \"gpu_unittests\"\n key: \"gpu_unittests\"\n command:\n - \"julia --color=yes --project test/runtests_gpu.jl\"\n agents:\n config: gpu\n queue: central\n slurm_ntasks: 1\n slurm_gres: \"gpu:1\"\n" }, { "path": ".codecov.yml", "content": "comment: off\n" }, { "path": ".dev/.gitignore", "content": "Manifest.toml\n" }, { "path": ".dev/Project.toml", "content": "[deps]\nJuliaFormatter = \"98e50ef6-434e-11e9-1051-2b60c6c9e899\"\n\n[compat]\nJuliaFormatter = \"0.10\"\n" }, { "path": ".dev/clima_formatter_default_image.jl", "content": "#!/usr/bin/env julia\n#! format: off\n#\n# Called with no arguments will replace the default system image with one that\n# includes the JuliaFormatter\n#\n# Called with a single argument the system image for the formatter will be\n# placed in the path specified by the argument (relative to the callers path)\n\nusing Pkg\nPkg.add(\"PackageCompiler\")\nusing PackageCompiler\n\nif isfile(PackageCompiler.backup_default_sysimg_path())\n @error \"\"\"\n A custom default system image already exists.\n Either restore default with:\n julia -e \"using PackageCompiler; PackageCompiler.restore_default_sysimage()\"\n or use the script\n $(abspath(joinpath(@__DIR__, \"..\", \".dev\", \"clima_formatter_image.jl\")))\n which will use a custom path for the system image\n \"\"\"\n exit(1)\nend\n\n# If a current Manifest exist for the formatter we remove it so that we have the\n# latest version\nrm(joinpath(@__DIR__, \"Manifest.toml\"); force = true)\n\nPkg.activate(joinpath(@__DIR__))\nPackageCompiler.create_sysimage(\n :JuliaFormatter;\n precompile_execution_file = joinpath(@__DIR__, \"precompile.jl\"),\n replace_default = true,\n)\n" }, { "path": ".dev/clima_formatter_image.jl", "content": "#!/usr/bin/env julia\n#\n# Called with no arguments will build the formattter system image \n# PATH_TO_CLIMATEMACHINE/.git/hooks/JuliaFormatterSysimage.so\n#\n# Called with a single argument the system image for the formatter will be\n# placed in the path specified by the argument (relative to the callers path)\n\nsysimage_path = abspath(\n isempty(ARGS) ?\n joinpath(@__DIR__, \"..\", \".git\", \"hooks\", \"JuliaFormatterSysimage.so\") :\n ARGS[1],\n)\n\n@info \"\"\"\nCreating system image object file at:\n $(sysimage_path)\n\"\"\"\n\nusing Pkg\nPkg.add(\"PackageCompiler\")\nusing PackageCompiler\n\n# If a current Manifest exist for the formatter we remove it so that we have the\n# latest version\nrm(joinpath(@__DIR__, \"Manifest.toml\"); force = true)\n\nPkg.activate(joinpath(@__DIR__))\nPackageCompiler.create_sysimage(\n :JuliaFormatter;\n precompile_execution_file = joinpath(@__DIR__, \"precompile.jl\"),\n sysimage_path = sysimage_path,\n)\n" }, { "path": ".dev/clima_formatter_options.jl", "content": "clima_formatter_options = (\n indent = 4,\n margin = 80,\n always_for_in = true,\n whitespace_typedefs = true,\n whitespace_ops_in_indices = true,\n remove_extra_newlines = false,\n)\n" }, { "path": ".dev/climaformat.jl", "content": "#!/usr/bin/env julia\n#\n# This is an adapted version of format.jl from JuliaFormatter with the\n# following license:\n#\n# MIT License\n# Copyright (c) 2019 Dominique Luna\n#\n# Permission is hereby granted, free of charge, to any person obtaining a\n# copy of this software and associated documentation files (the\n# \"Software\"), to deal in the Software without restriction, including\n# without limitation the rights to use, copy, modify, merge, publish,\n# distribute, sublicense, and/or sell copies of the Software, and to permit\n# persons to whom the Software is furnished to do so, subject to the\n# following conditions:\n#\n# The above copyright notice and this permission notice shall be included\n# in all copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS\n# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF\n# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN\n# NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,\n# DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR\n# OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE\n# USE OR OTHER DEALINGS IN THE SOFTWARE.\n#\nusing Pkg\nPkg.activate(@__DIR__)\nPkg.instantiate()\n\nusing JuliaFormatter\n\ninclude(\"clima_formatter_options.jl\")\n\nhelp = \"\"\"\nUsage: climaformat.jl [flags] [FILE/PATH]...\n\nFormats the given julia files using the CLIMA formatting options. If paths\nare given it will format the julia files in the paths. Otherwise, it will\nformat all changed julia files.\n\n -v, --verbose\n Print the name of the files being formatted with relevant details.\n\n -h, --help\n Print this message.\n\"\"\"\n\nfunction parse_opts!(args::Vector{String})\n i = 1\n opts = Dict{Symbol, Union{Int, Bool}}()\n while i ≤ length(args)\n arg = args[i]\n if arg[1] != '-'\n i += 1\n continue\n end\n if arg == \"-v\" || arg == \"--verbose\"\n opt = :verbose\n elseif arg == \"-h\" || arg == \"--help\"\n opt = :help\n else\n error(\"invalid option $arg\")\n end\n if opt in (:verbose, :help)\n opts[opt] = true\n deleteat!(args, i)\n end\n end\n return opts\nend\n\nopts = parse_opts!(ARGS)\nif haskey(opts, :help)\n write(stdout, help)\n exit(0)\nend\nif isempty(ARGS)\n filenames = readlines(`git ls-files \"*.jl\"`)\nelse\n filenames = ARGS\nend\n\nformat(filenames; clima_formatter_options..., opts...)\n" }, { "path": ".dev/hooks/pre-commit", "content": "#!/usr/bin/env julia\n#\n# Called by git-commit with no arguments. This checks to make sure that all\n# .jl files are indented correctly before a commit is made.\n#\n# To enable this hook, make this file executable and copy it in\n# $GIT_DIR/hooks.\n\ntoplevel_directory = chomp(read(`git rev-parse --show-toplevel`, String))\n\nusing Pkg\nPkg.activate(joinpath(toplevel_directory, \".dev\"))\nPkg.instantiate()\n\nusing JuliaFormatter\n\ninclude(joinpath(toplevel_directory, \".dev\", \"clima_formatter_options.jl\"))\n\nneeds_format = false\n\nfor diffoutput in split.(readlines(`git diff --name-status --cached`))\n status = diffoutput[1]\n filename = diffoutput[end]\n (!startswith(status, \"D\") && endswith(filename, \".jl\")) || continue\n\n a = read(`git show :$filename`, String)\n b = format_text(a; clima_formatter_options...)\n\n if a != b\n fullfilename = joinpath(toplevel_directory, filename)\n\n @error \"\"\"File $filename needs to be indented with:\n julia $(joinpath(toplevel_directory, \".dev\", \"climaformat.jl\")) $fullfilename\n and added to the git index via\n git add $fullfilename\n \"\"\"\n global needs_format = true\n end\nend\n\nexit(needs_format ? 1 : 0)\n" }, { "path": ".dev/hooks/pre-commit.sysimage", "content": "#!/usr/bin/env -S julia -J.git/hooks/JuliaFormatterSysimage.so\n#\n# Called by git-commit with no arguments. This checks to make sure that all\n# .jl files are indented correctly before a commit is made.\n#\n# To enable this hook, make this file executable and copy it in\n# $GIT_DIR/hooks.\n\ntoplevel_directory = chomp(read(`git rev-parse --show-toplevel`, String))\n\nusing Pkg\nPkg.activate(joinpath(toplevel_directory, \".dev\"))\nPkg.instantiate()\n\nusing JuliaFormatter\n\ninclude(joinpath(toplevel_directory, \".dev\", \"clima_formatter_options.jl\"))\n\nneeds_format = false\n\nfor diffoutput in split.(readlines(`git diff --name-status --cached`))\n status = diffoutput[1]\n filename = diffoutput[end]\n (!startswith(status, \"D\") && endswith(filename, \".jl\")) || continue\n\n a = read(`git show :$filename`, String)\n b = format_text(a; clima_formatter_options...)\n\n if a != b\n fullfilename = joinpath(toplevel_directory, filename)\n\n @error \"\"\"File $filename needs to be indented with:\n julia -J$(joinpath(toplevel_directory, \".git/hooks\", \"JuliaFormatterSysimage.so\")) $(joinpath(toplevel_directory, \".dev\", \"climaformat.jl\")) $fullfilename\n and added to the git index via\n git add $fullfilename\n \"\"\"\n global needs_format = true\n end\nend\n\nexit(needs_format ? 1 : 0)\n" }, { "path": ".dev/precompile.jl", "content": "using JuliaFormatter\n\ninclude(\"clima_formatter_options.jl\")\n\nformat(@__FILE__; clima_formatter_options...)\n" }, { "path": ".dev/systemimage/climate_machine_image.jl", "content": "#!/usr/bin/env julia\n#\n# Called with no arguments will create the system image\n# ClimateMachine.so\n# in the `@__DIR__` directory.\n#\n# Called with a single argument the system image will be placed in the path\n# specified by the argument (relative to the callers path)\n#\n# Called with a specified systemimg path and `true`, the system image will\n# compile the climate machine package module (useful for CI)\n\nsysimage_path =\n isempty(ARGS) ? joinpath(@__DIR__, \"ClimateMachine.so\") : abspath(ARGS[1])\n\nclimatemachine_pkg = get(ARGS, 2, \"false\") == \"true\"\n\n@info \"Creating system image object file at: '$(sysimage_path)'\"\n@info \"Building ClimateMachine into system image: $(climatemachine_pkg)\"\n\nstart_time = time()\n\nusing Pkg\nPkg.add(\"PackageCompiler\")\n\nPkg.activate(joinpath(@__DIR__, \"..\", \"..\"))\nPkg.instantiate(verbose = true)\n\npkgs = Symbol[]\nif climatemachine_pkg\n push!(pkgs, :ClimateMachine)\nelse\n append!(\n pkgs,\n [Symbol(v.name) for v in values(Pkg.dependencies()) if v.is_direct_dep],\n )\nend\n\n# use package compiler\nusing PackageCompiler\nPackageCompiler.create_sysimage(\n pkgs,\n sysimage_path = sysimage_path,\n precompile_execution_file = joinpath(\n @__DIR__,\n \"..\",\n \"..\",\n \"test\",\n \"Numerics\",\n \"DGMethods\",\n \"Euler\",\n \"isentropicvortex.jl\",\n ),\n)\n\ntot_secs = Int(floor(time() - start_time))\n@info \"Created system image object file at: $(sysimage_path)\"\n@info \"System image build time: $tot_secs sec\"\n" }, { "path": ".github/issue_template.md", "content": "### Description\n\n\n" }, { "path": ".github/pull_request_template.md", "content": "### Description\n\n\n\n\n\n- [ ] Code follows the [style guidelines](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) OR N/A.\n- [ ] Unit tests are included OR N/A.\n- [ ] Code is exercised in an integration test OR N/A.\n- [ ] Documentation has been added/updated OR N/A.\n" }, { "path": ".github/workflows/CompatHelper.yml", "content": "name: CompatHelper\n\non:\n schedule:\n - cron: '00 * * * *'\n\njobs:\n CompatHelper:\n runs-on: ubuntu-latest\n steps:\n - uses: julia-actions/setup-julia@latest\n with:\n version: 1.5.4\n - name: Pkg.add(\"CompatHelper\")\n run: julia -e 'using Pkg; Pkg.add(\"CompatHelper\")'\n - name: CompatHelper.main()\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n run: julia -e 'using CompatHelper; CompatHelper.main()'\n" }, { "path": ".github/workflows/Coverage.yaml", "content": "name: Coverage\n\non:\n schedule:\n # * is a special character in YAML so you have to quote this string\n # Run at 2am every day:\n - cron: '0 2 * * *'\n\njobs:\n coverage:\n runs-on: ubuntu-latest\n\n steps:\n - uses: actions/checkout@v2.2.0\n\n - name: Set up Julia\n uses: julia-actions/setup-julia@latest\n with:\n version: 1.5.4\n \n - name: Test with coverage\n run: |\n julia --project -e 'using Pkg; Pkg.instantiate()'\n julia --project -e 'using Pkg; Pkg.test(coverage=true)'\n\n - name: Generate coverage file\n run: julia --project -e 'using Pkg; Pkg.add(\"Coverage\");\n using Coverage;\n LCOV.writefile(\"coverage-lcov.info\", Codecov.process_folder())'\n if: success()\n\n - name: Submit coverage\n uses: codecov/codecov-action@v1.0.7\n with:\n token: ${{secrets.CODECOV_TOKEN}}\n if: success()" }, { "path": ".github/workflows/DocCleanup.yml", "content": "name: Doc Preview Cleanup\n\non:\n pull_request:\n types: [closed]\n\njobs:\n doc-preview-cleanup:\n runs-on: ubuntu-latest\n steps:\n - name: Checkout gh-pages branch\n uses: actions/checkout@v2\n with:\n ref: gh-pages\n\n - name: Delete preview and history\n run: |\n git config user.name \"Documenter.jl\"\n git config user.email \"documenter@juliadocs.github.io\"\n git rm -rf \"previews/PR$PRNUM\"\n git commit -m \"delete preview\"\n git branch gh-pages-new $(echo \"delete history\" | git commit-tree HEAD^{tree})\n env:\n PRNUM: ${{ github.event.number }}\n\n - name: Push changes\n run: |\n git push --force origin gh-pages-new:gh-pages\n" }, { "path": ".github/workflows/Documenter.yaml", "content": "name: Documentation\n\non:\n pull_request:\n paths:\n - 'docs/**'\n - 'tutorials/**'\n - 'src/**'\n - 'Project.toml'\n - 'Manifest.toml'\n\njobs:\n docs-build:\n runs-on: ubuntu-latest\n timeout-minutes: 90\n steps:\n - name: Cancel Previous Runs\n uses: styfle/cancel-workflow-action@0.4.0\n with:\n access_token: ${{ github.token }}\n - uses: actions/checkout@v2.2.0\n - name: Install System Dependencies\n run: |\n sudo apt-get update\n sudo apt-get -qq install libxt6 libxrender1 libxext6 libgl1-mesa-glx libqt5widgets5 xvfb\n\n - uses: julia-actions/setup-julia@latest\n with:\n version: 1.5.4\n\n # https://discourse.julialang.org/t/recommendation-cache-julia-artifacts-in-ci-services/35484\n - name: Cache artifacts\n uses: actions/cache@v1\n env:\n cache-name: cache-artifacts\n with:\n path: ~/.julia/artifacts\n key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}\n restore-keys: |\n ${{ runner.os }}-test-${{ env.cache-name }}-\n ${{ runner.os }}-test-\n ${{ runner.os }}-\n\n - name: Install Julia dependencies\n env:\n JULIA_PROJECT: \"docs/\"\n run: |\n julia --project -e 'using Pkg; Pkg.instantiate()'\n julia --project=docs/ -e 'using Pkg; Pkg.instantiate()'\n julia --project=docs/ -e 'using Pkg; Pkg.precompile()'\n - name: Build and deploy\n # Run with X virtual frame buffer as GR (default backend for Plots.jl) needs\n # an X session to run without warnings\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n XDG_RUNTIME_DIR: \"/home/runner\"\n JULIA_PROJECT: \"docs/\"\n CLIMATEMACHINE_DOCS_GENERATE_TUTORIALS: \"false\"\n ClIMATEMACHINE_SETTINGS_DISABLE_GPU: \"true\"\n CLIMATEMACHINE_SETTINGS_DISABLE_CUSTOM_LOGGER: \"true\"\n CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: \"true\"\n run: xvfb-run -- julia --project=docs/ --color=yes docs/make.jl\n - name: Help! Documenter Failed\n run: |\n cat .github/workflows/doc_build_common_error_messages.md\n if: failure()\n" }, { "path": ".github/workflows/JuliaFormatter.yml", "content": "name: JuliaFormatter\n\non:\n push:\n branches:\n - master\n - trying\n - staging\n tags: '*'\n pull_request:\n\njobs:\n format:\n runs-on: ubuntu-latest\n timeout-minutes: 30\n steps:\n - name: Cancel Previous Runs\n uses: styfle/cancel-workflow-action@0.4.0\n with:\n access_token: ${{ github.token }}\n\n - uses: actions/checkout@v2.2.0\n\n - uses: dorny/paths-filter@v2.9.1\n id: filter\n with:\n filters: |\n julia_file_change:\n - added|modified: '**.jl'\n \n - uses: julia-actions/setup-julia@latest\n if: steps.filter.outputs.julia_file_change == 'true'\n with:\n version: 1.5.4\n\n - name: Apply JuliaFormatter\n if: steps.filter.outputs.julia_file_change == 'true'\n run: |\n julia --project=.dev .dev/climaformat.jl .\n\n - name: Check formatting diff\n if: steps.filter.outputs.julia_file_change == 'true'\n run: |\n git diff --color=always --exit-code\n" }, { "path": ".github/workflows/Linux-UnitTests.yml", "content": "name: Unit Tests\n\non:\n pull_request:\n paths:\n - 'src/**'\n - 'test/**'\n - 'Project.toml'\n - 'Manifest.toml'\n\njobs:\n test:\n runs-on: ubuntu-latest\n timeout-minutes: 60\n strategy:\n fail-fast: true\n matrix:\n test-modules: [\"Atmos,Common,InputOutput,Utilities\",\n \"Driver\",\n \"Diagnostics\",\n \"Arrays,Numerics/ODESolvers,Numerics/SystemSolvers\",\n \"Ocean\",\n \"Land\",]\n\n env:\n CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: \"true\"\n\n steps:\n - name: Cancel Previous Runs\n uses: styfle/cancel-workflow-action@0.4.0\n with:\n access_token: ${{ github.token }}\n\n - name: Checkout\n uses: actions/checkout@v2.2.0\n\n - name: Set up Julia\n uses: julia-actions/setup-julia@latest\n with:\n version: 1.5.4\n\n # https://discourse.julialang.org/t/recommendation-cache-julia-artifacts-in-ci-services/35484\n - name: Cache artifacts\n uses: actions/cache@v1\n env:\n cache-name: cache-artifacts\n with:\n path: ~/.julia/artifacts \n key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}\n restore-keys: |\n ${{ runner.os }}-test-${{ env.cache-name }}-\n ${{ runner.os }}-test-\n ${{ runner.os }}-\n\n - name: Install Project Packages\n run: |\n julia --project=@. -e 'using Pkg; Pkg.instantiate()'\n julia --project=@. -e 'using Pkg; Pkg.precompile()'\n\n - name: Run Unit Tests\n env:\n TEST_MODULES: ${{ matrix.test-modules }}\n run: |\n julia --project=@. -e 'using Pkg; Pkg.test(test_args=map(String, split(ENV[\"TEST_MODULES\"], \",\")))'\n" }, { "path": ".github/workflows/OS-UnitTests.yml", "content": "name: OS Unit Tests\n\non:\n push:\n branches:\n - staging\n - trying\n\njobs:\n test-os:\n timeout-minutes: 210\n strategy:\n fail-fast: true\n matrix:\n os: [ubuntu-latest, windows-latest, macos-latest]\n\n runs-on: ${{ matrix.os }}\n\n # Workaround for OSX MPICH issue:\n # https://github.com/pmodels/mpich/issues/4710\n env:\n MPICH_INTERFACE_HOSTNAME: \"localhost\"\n CLIMATEMACHINE_TEST_RUNMPI_LOCALHOST: \"true\"\n CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: \"true\"\n\n steps:\n - name: Cancel Previous Runs\n uses: styfle/cancel-workflow-action@0.4.0\n with:\n access_token: ${{ github.token }}\n\n - name: Checkout\n uses: actions/checkout@v2.2.0\n \n # Setup a filter and only run if src/ test/ folder content changes\n # or project depedencies\n - uses: dorny/paths-filter@v2\n id: filter\n with:\n filters: |\n run_test:\n - 'src/**'\n - 'test/**'\n - 'Project.toml'\n - 'Manifest.toml'\n\n\n - name: Set up Julia\n uses: julia-actions/setup-julia@latest\n if: steps.filter.outputs.run_test == 'true'\n with:\n version: 1.5.4\n\n # https://discourse.julialang.org/t/recommendation-cache-julia-artifacts-in-ci-services/35484\n - name: Cache artifacts\n uses: actions/cache@v1\n if: steps.filter.outputs.run_test == 'true'\n env:\n cache-name: cache-artifacts\n with:\n path: ~/.julia/artifacts \n key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}\n restore-keys: |\n ${{ runner.os }}-test-${{ env.cache-name }}-\n ${{ runner.os }}-test-\n ${{ runner.os }}-\n\n - name: Install Project Packages\n if: steps.filter.outputs.run_test == 'true'\n run: |\n julia --project=@. -e 'using Pkg; Pkg.instantiate()'\n\n - name: Build System Image\n if: steps.filter.outputs.run_test == 'true'\n continue-on-error: true\n run: |\n julia --project .dev/systemimage/climate_machine_image.jl ClimateMachine.so true\n\n - name: Run Unit Tests\n if: steps.filter.outputs.run_test == 'true'\n run: |\n julia --project -J ClimateMachine.so -e 'using Pkg; Pkg.test()'\n" }, { "path": ".github/workflows/PR-Comment.yml", "content": "name: Trigger action on PR comment\non:\n issue_comment:\n types: [created]\n\njobs:\n trigger-doc-build:\n if: ${{ github.event.issue.pull_request &&\n ( github.event.comment.author_association == 'OWNER' ||\n github.event.comment.author_association == 'MEMBER' ||\n github.event.comment.author_association == 'COLLABORATOR' ) &&\n contains(github.event.comment.body, '@climabot build docs') }}\n runs-on: ubuntu-latest\n steps:\n - uses: octokit/request-action@v2.0.26\n id: get_pr # need sha of commit\n with:\n route: GET /repos/{repository}/pulls/{pull_number}\n repository: ${{ github.repository }}\n pull_number: ${{ github.event.issue.number }}\n env:\n GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}\n \n - name: Trigger Buildkite Pipeline\n id: buildkite\n uses: CliMA/buildkite-pipeline-action@master\n with:\n access_token: ${{ secrets.BUILDKITE }}\n pipeline: 'clima/climatemachine-docs'\n branch: ${{ fromJson(steps.get_pr.outputs.data).head.ref }}\n commit: ${{ fromJson(steps.get_pr.outputs.data).head.sha }}\n message: \":github: Triggered by comment on PR #${{ github.event.issue.number }}\"\n env: '{\"PULL_REQUEST\": ${{ github.event.issue.number }} }'\n async: true\n\n - name: Create comment\n uses: peter-evans/create-or-update-comment@v1\n with:\n issue-number: ${{ github.event.issue.number }}\n body: |\n Docs build created: ${{ steps.buildkite.outputs.web_url }}\n Preview link: https://clima.github.io/ClimateMachine.jl/previews/PR${{ github.event.issue.number}}\n \n" }, { "path": ".github/workflows/doc_build_common_error_messages.md", "content": "# Documenter common warning/error messages:\n\n## General notes\n - Changing order of API/HowToGuides does not fix unresolved path issues\n\n## no doc found for reference\n```\n┌ Warning: no doc found for reference '[`BatchedGeneralizedMinimalResidual`](@ref)' in src/HowToGuides/Numerics/SystemSolvers/IterativeSolvers.md.\n└ @ Documenter.CrossReferences ~/.julia/packages/Documenter/PLD7m/src/CrossReferences.jl:160\n```\n - Missing entry in ```@docs ``` in API\n\n\n## Reference could not be found\n```\n┌ Warning: reference for 'ClimateMachine.ODESolvers.solve!' could not be found in src\\APIs\\Driver\\index.md.\n└ @ Documenter.CrossReferences C:\\Users\\kawcz\\.julia\\packages\\Documenter\\PLD7m\\src\\CrossReferences.jl:104\n```\n - Missing doc string?\n\n## invalid local link: unresolved path\n```\n┌ Warning: invalid local link: unresolved path in APIs/Atmos/AtmosModel.md\n│ link.text =\n│ 1-element Array{Any,1}:\n│ Markdown.Code(\"\", \"FlatOrientation\")\n│ link.url = \"@ref\"\n```\n - Missing entry in ```@docs ``` for FlatOrientation in API OR\n - The \"code\" in the reference must be to actual code and not arbitrary text\n\n## unable to get the binding\n```\n┌ Warning: unable to get the binding for 'ClimateMachine.Atmos.AtmosModel.NoOrientation' in `@docs` block in src/APIs/Atmos/AtmosModel.md:9-14 from expression ':(ClimateMachine.Atmos.AtmosModel.NoOrientation)' in module ClimateMachine\n│ ```@docs\n│ ClimateMachine.Atmos.AtmosModel.NoOrientation\n...\n│ ```\n...\n```\n - `ClimateMachine.Atmos.AtmosModel.NoOrientation` should be `ClimateMachine.Atmos.NoOrientation`\n\n## Other useful tips\n- The syntax is white space sensitive: Do not leave any extra new line between the end of the doc string (denoted by triple double-quotes `\"\"\"`) and the code of the defined method / type / module name that you are describing.\n- In the doc string, indent the method / type / module signature and do not indent the descriptive text.\n- Any method name and the corresponding signature in the doc string have to match 1:1 (be careful of missing/extra exclamation points `!`)\n" }, { "path": ".gitignore", "content": "# Temporary\n*.DS_Store\n*.swp\n*.jl.cov\n*.jl.*.cov\n*.jl.mem\n*.DS_Store\n*~\nTAGS\n\n# Docs\ndocs/build/\ndocs/site/\ndocs/src/tutorials/\ndocs/src/generated/\ndocs/transient_dir/generated/\n!docs/src/assets/*.png\n!docs/src/assets/*.svg\n\n# Deps\nsrc/**/Manifest.toml\n\n# Data\n*.vtk\n*.dat\n*.vtu\n*.pvtu\n*.nc\n*.jld2\n\n# Figs\n*.png\n*.jpg\n*.jpeg\n*.svg\n\n# Movies\n*.gif\n*.mp4\n\n# Julia System Images\n*.so\n" }, { "path": "LICENSE.md", "content": "Copyright 2019 Climate Modeling Alliance\n\nLicensed under the Apache License, Version 2.0 (the \"License\");\nyou may not use this file except in compliance with the License.\nYou may obtain a copy of the License at\n\n http://www.apache.org/licenses/LICENSE-2.0\n\nUnless required by applicable law or agreed to in writing, software\ndistributed under the License is distributed on an \"AS IS\" BASIS,\nWITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\nSee the License for the specific language governing permissions and\nlimitations under the License.\n" }, { "path": "Manifest.toml", "content": "# This file is machine-generated - editing it directly is not advised\n\n[[AbstractFFTs]]\ndeps = [\"LinearAlgebra\"]\ngit-tree-sha1 = \"485ee0867925449198280d4af84bdb46a2a404d0\"\nuuid = \"621f4979-c628-5d54-868e-fcf4e3e8185c\"\nversion = \"1.0.1\"\n\n[[Adapt]]\ndeps = [\"LinearAlgebra\"]\ngit-tree-sha1 = \"ffcfa2d345aaee0ef3d8346a073d5dd03c983ebe\"\nuuid = \"79e6a3ab-5dfb-504d-930d-738a2a938a0e\"\nversion = \"3.2.0\"\n\n[[ArgParse]]\ndeps = [\"Logging\", \"TextWrap\"]\ngit-tree-sha1 = \"e928ca0a49f7b0564044b39108c70c160f03e05a\"\nuuid = \"c7e460c6-2fb9-53a9-8c5b-16f535851c63\"\nversion = \"1.1.2\"\n\n[[ArgTools]]\ngit-tree-sha1 = \"bdf73eec6a88885256f282d48eafcad25d7de494\"\nuuid = \"0dad84c5-d112-42e6-8d28-ef12dabb789f\"\nversion = \"1.1.1\"\n\n[[ArnoldiMethod]]\ndeps = [\"LinearAlgebra\", \"Random\", \"StaticArrays\"]\ngit-tree-sha1 = \"f87e559f87a45bece9c9ed97458d3afe98b1ebb9\"\nuuid = \"ec485272-7323-5ecc-a04f-4719b315124d\"\nversion = \"0.1.0\"\n\n[[ArrayInterface]]\ndeps = [\"IfElse\", \"LinearAlgebra\", \"Requires\", \"SparseArrays\", \"Static\"]\ngit-tree-sha1 = \"ce17bad65d0842b34a15fffc8879a9f68f08a67f\"\nuuid = \"4fba245c-0d91-5ea0-9b3e-6abc04ee57a9\"\nversion = \"3.1.6\"\n\n[[ArrayLayouts]]\ndeps = [\"FillArrays\", \"LinearAlgebra\", \"SparseArrays\"]\ngit-tree-sha1 = \"9aa9647b58147a81f7359eacc7d6249ac3a3e3d4\"\nuuid = \"4c555306-a7a7-4459-81d9-ec55ddd5c99a\"\nversion = \"0.6.4\"\n\n[[ArtifactWrappers]]\ndeps = [\"DocStringExtensions\", \"Downloads\", \"Pkg\"]\ngit-tree-sha1 = \"e9b52e63e3ea81a504412807c9426566e26c232d\"\nuuid = \"a14bc488-3040-4b00-9dc1-f6467924858a\"\nversion = \"0.1.1\"\n\n[[Artifacts]]\ndeps = [\"Pkg\"]\ngit-tree-sha1 = \"c30985d8821e0cd73870b17b0ed0ce6dc44cb744\"\nuuid = \"56f22d72-fd6d-98f1-02f0-08ddc0907c33\"\nversion = \"1.3.0\"\n\n[[BFloat16s]]\ndeps = [\"LinearAlgebra\", \"Test\"]\ngit-tree-sha1 = \"4af69e205efc343068dc8722b8dfec1ade89254a\"\nuuid = \"ab4f0b2a-ad5b-11e8-123f-65d77653426b\"\nversion = \"0.1.0\"\n\n[[Base64]]\nuuid = \"2a0f44e3-6c83-55bd-87e4-b1978d98bd5f\"\n\n[[BenchmarkTools]]\ndeps = [\"JSON\", \"Logging\", \"Printf\", \"Statistics\", \"UUIDs\"]\ngit-tree-sha1 = \"9e62e66db34540a0c919d72172cc2f642ac71260\"\nuuid = \"6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf\"\nversion = \"0.5.0\"\n\n[[CEnum]]\ngit-tree-sha1 = \"215a9aa4a1f23fbd05b92769fdd62559488d70e9\"\nuuid = \"fa961155-64e5-5f13-b03f-caf6b980ea82\"\nversion = \"0.4.1\"\n\n[[CFTime]]\ndeps = [\"Dates\", \"Printf\"]\ngit-tree-sha1 = \"bca6cb6ee746e6485ca4535f6cc29cf3579a0f20\"\nuuid = \"179af706-886a-5703-950a-314cd64e0468\"\nversion = \"0.1.1\"\n\n[[CLIMAParameters]]\ndeps = [\"Test\"]\ngit-tree-sha1 = \"0801216ee1670a1e5280cdb0e5fda60cc4b992ca\"\nuuid = \"6eacf6c3-8458-43b9-ae03-caf5306d3d53\"\nversion = \"0.2.0\"\n\n[[CUDA]]\ndeps = [\"AbstractFFTs\", \"Adapt\", \"BFloat16s\", \"CEnum\", \"CompilerSupportLibraries_jll\", \"DataStructures\", \"ExprTools\", \"GPUArrays\", \"GPUCompiler\", \"LLVM\", \"Libdl\", \"LinearAlgebra\", \"Logging\", \"MacroTools\", \"NNlib\", \"Pkg\", \"Printf\", \"Random\", \"Reexport\", \"Requires\", \"SparseArrays\", \"Statistics\", \"TimerOutputs\"]\ngit-tree-sha1 = \"6ccc73b2d8b671f7a65c92b5f08f81422ebb7547\"\nuuid = \"052768ef-5323-5732-b1bb-66c8b64840ba\"\nversion = \"2.4.1\"\n\n[[Cassette]]\ngit-tree-sha1 = \"742fbff99a2798f02bd37d25087efb5615b5a207\"\nuuid = \"7057c7e9-c182-5462-911a-8362d720325c\"\nversion = \"0.3.5\"\n\n[[ChainRulesCore]]\ndeps = [\"Compat\", \"LinearAlgebra\", \"SparseArrays\"]\ngit-tree-sha1 = \"644c24cd6344348f1c645efab24b707088be526a\"\nuuid = \"d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4\"\nversion = \"0.9.34\"\n\n[[CloudMicrophysics]]\ndeps = [\"CLIMAParameters\", \"DocStringExtensions\", \"SpecialFunctions\", \"Thermodynamics\"]\ngit-tree-sha1 = \"52c85d99a842a1d9ebf6f36e8dd34d55e903b0b5\"\nuuid = \"6a9e3e04-43cd-43ba-94b9-e8782df3c71b\"\nversion = \"0.3.1\"\n\n[[CodecZlib]]\ndeps = [\"TranscodingStreams\", \"Zlib_jll\"]\ngit-tree-sha1 = \"ded953804d019afa9a3f98981d99b33e3db7b6da\"\nuuid = \"944b1d66-785c-5afd-91f1-9de20f533193\"\nversion = \"0.7.0\"\n\n[[Combinatorics]]\ngit-tree-sha1 = \"08c8b6831dc00bfea825826be0bc8336fc369860\"\nuuid = \"861a8166-3701-5b0c-9a16-15d98fcdc6aa\"\nversion = \"1.0.2\"\n\n[[CommonSolve]]\ngit-tree-sha1 = \"68a0743f578349ada8bc911a5cbd5a2ef6ed6d1f\"\nuuid = \"38540f10-b2f7-11e9-35d8-d573e4eb0ff2\"\nversion = \"0.2.0\"\n\n[[CommonSubexpressions]]\ndeps = [\"MacroTools\", \"Test\"]\ngit-tree-sha1 = \"7b8a93dba8af7e3b42fecabf646260105ac373f7\"\nuuid = \"bbf7d656-a473-5ed7-a52c-81e309532950\"\nversion = \"0.3.0\"\n\n[[Compat]]\ndeps = [\"Base64\", \"Dates\", \"DelimitedFiles\", \"Distributed\", \"InteractiveUtils\", \"LibGit2\", \"Libdl\", \"LinearAlgebra\", \"Markdown\", \"Mmap\", \"Pkg\", \"Printf\", \"REPL\", \"Random\", \"SHA\", \"Serialization\", \"SharedArrays\", \"Sockets\", \"SparseArrays\", \"Statistics\", \"Test\", \"UUIDs\", \"Unicode\"]\ngit-tree-sha1 = \"919c7f3151e79ff196add81d7f4e45d91bbf420b\"\nuuid = \"34da2185-b29b-5c13-b0c7-acf172513d20\"\nversion = \"3.25.0\"\n\n[[CompilerSupportLibraries_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"8e695f735fca77e9708e795eda62afdb869cbb70\"\nuuid = \"e66e0078-7015-5450-92f7-15fbd957f2ae\"\nversion = \"0.3.4+0\"\n\n[[ConstructionBase]]\ndeps = [\"LinearAlgebra\"]\ngit-tree-sha1 = \"48920211c95a6da1914a06c44ec94be70e84ffff\"\nuuid = \"187b0558-2788-49d3-abe0-74a17ed4e7c9\"\nversion = \"1.1.0\"\n\n[[Coverage]]\ndeps = [\"CoverageTools\", \"HTTP\", \"JSON\", \"LibGit2\", \"MbedTLS\"]\ngit-tree-sha1 = \"a485b62b1ce53368e3a1a3a8dc0c39a61b9d3d9c\"\nuuid = \"a2441757-f6aa-5fb2-8edb-039e3f45d037\"\nversion = \"1.2.0\"\n\n[[CoverageTools]]\ndeps = [\"LibGit2\"]\ngit-tree-sha1 = \"08b72d2f2154e33dc2aeb1bfcd4a83cb283abd4f\"\nuuid = \"c36e975a-824b-4404-a568-ef97ca766997\"\nversion = \"1.2.2\"\n\n[[Crayons]]\ngit-tree-sha1 = \"3f71217b538d7aaee0b69ab47d9b7724ca8afa0d\"\nuuid = \"a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f\"\nversion = \"4.0.4\"\n\n[[CubedSphere]]\ndeps = [\"Elliptic\", \"Printf\", \"Rotations\", \"TaylorSeries\", \"Test\"]\ngit-tree-sha1 = \"f66fabd1ee5df59a7ba47c7873a6332c19e0c03f\"\nuuid = \"7445602f-e544-4518-8976-18f8e8ae6cdb\"\nversion = \"0.2.0\"\n\n[[DataAPI]]\ngit-tree-sha1 = \"dfb3b7e89e395be1e25c2ad6d7690dc29cc53b1d\"\nuuid = \"9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a\"\nversion = \"1.6.0\"\n\n[[DataStructures]]\ndeps = [\"Compat\", \"InteractiveUtils\", \"OrderedCollections\"]\ngit-tree-sha1 = \"4437b64df1e0adccc3e5d1adbc3ac741095e4677\"\nuuid = \"864edb3b-99cc-5e75-8d2d-829cb0a9cfe8\"\nversion = \"0.18.9\"\n\n[[DataValueInterfaces]]\ngit-tree-sha1 = \"bfc1187b79289637fa0ef6d4436ebdfe6905cbd6\"\nuuid = \"e2d170a0-9d28-54be-80f0-106bbe20a464\"\nversion = \"1.0.0\"\n\n[[Dates]]\ndeps = [\"Printf\"]\nuuid = \"ade2ca70-3891-5945-98fb-dc099432e06a\"\n\n[[DelimitedFiles]]\ndeps = [\"Mmap\"]\nuuid = \"8bb1440f-4735-579b-a4ab-409b98df4dab\"\n\n[[Dierckx]]\ndeps = [\"Dierckx_jll\"]\ngit-tree-sha1 = \"5fefbe52e9a6e55b8f87cb89352d469bd3a3a090\"\nuuid = \"39dd38d3-220a-591b-8e3c-4c3a8c710a94\"\nversion = \"0.5.1\"\n\n[[Dierckx_jll]]\ndeps = [\"CompilerSupportLibraries_jll\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"a580560f526f6fc6973e8bad2b036514a4e3b013\"\nuuid = \"cd4c43a9-7502-52ba-aa6d-59fb2a88580b\"\nversion = \"0.0.1+0\"\n\n[[DiffEqBase]]\ndeps = [\"ArrayInterface\", \"ChainRulesCore\", \"DataStructures\", \"DocStringExtensions\", \"FunctionWrappers\", \"IterativeSolvers\", \"LabelledArrays\", \"LinearAlgebra\", \"Logging\", \"MuladdMacro\", \"NonlinearSolve\", \"Parameters\", \"Printf\", \"RecursiveArrayTools\", \"RecursiveFactorization\", \"Reexport\", \"Requires\", \"SciMLBase\", \"SparseArrays\", \"StaticArrays\", \"Statistics\", \"SuiteSparse\", \"ZygoteRules\"]\ngit-tree-sha1 = \"c2ff625248a0967adff1dc1f79c6a41e2531f081\"\nuuid = \"2b5f629d-d688-5b77-993f-72d75c75574e\"\nversion = \"6.57.8\"\n\n[[DiffResults]]\ndeps = [\"StaticArrays\"]\ngit-tree-sha1 = \"c18e98cba888c6c25d1c3b048e4b3380ca956805\"\nuuid = \"163ba53b-c6d8-5494-b064-1a9d43ac40c5\"\nversion = \"1.0.3\"\n\n[[DiffRules]]\ndeps = [\"NaNMath\", \"Random\", \"SpecialFunctions\"]\ngit-tree-sha1 = \"214c3fcac57755cfda163d91c58893a8723f93e9\"\nuuid = \"b552c78f-8df3-52c6-915a-8e097449b14b\"\nversion = \"1.0.2\"\n\n[[DispatchedTuples]]\ngit-tree-sha1 = \"4ccc236f2e2f6e4a15b093d76184ffded23d211f\"\nuuid = \"508c55e1-51b4-41fd-a5ca-7eb0327d070d\"\nversion = \"0.2.0\"\n\n[[Distances]]\ndeps = [\"LinearAlgebra\", \"Statistics\"]\ngit-tree-sha1 = \"366715149014943abd71aa647a07a43314158b2d\"\nuuid = \"b4f34e82-e78d-54a5-968a-f98e89d6e8f7\"\nversion = \"0.10.2\"\n\n[[Distributed]]\ndeps = [\"Random\", \"Serialization\", \"Sockets\"]\nuuid = \"8ba89e20-285c-5b6f-9357-94700520ee1b\"\n\n[[Distributions]]\ndeps = [\"FillArrays\", \"LinearAlgebra\", \"PDMats\", \"Printf\", \"QuadGK\", \"Random\", \"SparseArrays\", \"SpecialFunctions\", \"Statistics\", \"StatsBase\", \"StatsFuns\"]\ngit-tree-sha1 = \"e64debe8cd174cc52d7dd617ebc5492c6f8b698c\"\nuuid = \"31c24e10-a181-5473-b8eb-7969acd0382f\"\nversion = \"0.24.15\"\n\n[[DocStringExtensions]]\ndeps = [\"LibGit2\", \"Markdown\", \"Pkg\", \"Test\"]\ngit-tree-sha1 = \"50ddf44c53698f5e784bbebb3f4b21c5807401b1\"\nuuid = \"ffbed154-4ef7-542d-bbb7-c09d3a79fcae\"\nversion = \"0.8.3\"\n\n[[DoubleFloats]]\ndeps = [\"GenericSVD\", \"GenericSchur\", \"LinearAlgebra\", \"Polynomials\", \"Printf\", \"Quadmath\", \"Random\", \"Requires\", \"SpecialFunctions\"]\ngit-tree-sha1 = \"d5cb090c9f59e5024e0c94be0714c5de8ff5dc99\"\nuuid = \"497a8b3b-efae-58df-a0af-a86822472b78\"\nversion = \"1.1.18\"\n\n[[Downloads]]\ndeps = [\"ArgTools\", \"LibCURL\", \"NetworkOptions\"]\ngit-tree-sha1 = \"5de8c54d269fd7ab430656c27de73e63eb07a979\"\nuuid = \"f43a241f-c20a-4ad4-852c-f6b1247861c6\"\nversion = \"1.4.0\"\n\n[[Elliptic]]\ngit-tree-sha1 = \"71c79e77221ab3a29918aaf6db4f217b89138608\"\nuuid = \"b305315f-e792-5b7a-8f41-49f472929428\"\nversion = \"1.0.1\"\n\n[[ExponentialUtilities]]\ndeps = [\"LinearAlgebra\", \"Printf\", \"Requires\", \"SparseArrays\"]\ngit-tree-sha1 = \"712cb5af8db62836913970ee035a5fa742986f00\"\nuuid = \"d4d017d3-3776-5f7e-afef-a10c40355c18\"\nversion = \"1.8.1\"\n\n[[ExprTools]]\ngit-tree-sha1 = \"10407a39b87f29d47ebaca8edbc75d7c302ff93e\"\nuuid = \"e2ba6199-217a-4e67-a87a-7c52f15ade04\"\nversion = \"0.1.3\"\n\n[[EzXML]]\ndeps = [\"Printf\", \"XML2_jll\"]\ngit-tree-sha1 = \"0fa3b52a04a4e210aeb1626def9c90df3ae65268\"\nuuid = \"8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615\"\nversion = \"1.1.0\"\n\n[[FFTW]]\ndeps = [\"AbstractFFTs\", \"FFTW_jll\", \"IntelOpenMP_jll\", \"Libdl\", \"LinearAlgebra\", \"MKL_jll\", \"Reexport\"]\ngit-tree-sha1 = \"1b48dbde42f307e48685fa9213d8b9f8c0d87594\"\nuuid = \"7a1cc6ca-52ef-59f5-83cd-3a7055c09341\"\nversion = \"1.3.2\"\n\n[[FFTW_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"3676abafff7e4ff07bbd2c42b3d8201f31653dcc\"\nuuid = \"f5851436-0d7a-5f13-b9de-f02708fd171a\"\nversion = \"3.3.9+8\"\n\n[[FastClosures]]\ngit-tree-sha1 = \"acebe244d53ee1b461970f8910c235b259e772ef\"\nuuid = \"9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a\"\nversion = \"0.3.2\"\n\n[[FileIO]]\ndeps = [\"Pkg\", \"Requires\", \"UUIDs\"]\ngit-tree-sha1 = \"9cdfbf5c0ed88ad0dcdb02544416c8e5a73addef\"\nuuid = \"5789e2e9-d7fb-5bc7-8068-2c6fae9b9549\"\nversion = \"1.6.4\"\n\n[[FillArrays]]\ndeps = [\"LinearAlgebra\", \"Random\", \"SparseArrays\"]\ngit-tree-sha1 = \"31939159aeb8ffad1d4d8ee44d07f8558273120a\"\nuuid = \"1a297f60-69ca-5386-bcde-b61e274b549b\"\nversion = \"0.11.7\"\n\n[[FiniteDiff]]\ndeps = [\"ArrayInterface\", \"LinearAlgebra\", \"Requires\", \"SparseArrays\", \"StaticArrays\"]\ngit-tree-sha1 = \"f6f80c8f934efd49a286bb5315360be66956dfc4\"\nuuid = \"6a86dc24-6348-571c-b903-95158fe2bd41\"\nversion = \"2.8.0\"\n\n[[Formatting]]\ndeps = [\"Printf\"]\ngit-tree-sha1 = \"8339d61043228fdd3eb658d86c926cb282ae72a8\"\nuuid = \"59287772-0a20-5a39-b81b-1366585eb4c0\"\nversion = \"0.4.2\"\n\n[[ForwardDiff]]\ndeps = [\"CommonSubexpressions\", \"DiffResults\", \"DiffRules\", \"NaNMath\", \"Random\", \"SpecialFunctions\", \"StaticArrays\"]\ngit-tree-sha1 = \"c68fb7481b71519d313114dca639b35262ff105f\"\nuuid = \"f6369f11-7733-5829-9624-2563aa707210\"\nversion = \"0.10.17\"\n\n[[FunctionWrappers]]\ngit-tree-sha1 = \"241552bc2209f0fa068b6415b1942cc0aa486bcc\"\nuuid = \"069b7b12-0de2-55c6-9aab-29f3d0a68a2e\"\nversion = \"1.1.2\"\n\n[[Future]]\ndeps = [\"Random\"]\nuuid = \"9fa8497b-333b-5362-9e8d-4d0656e87820\"\n\n[[GPUArrays]]\ndeps = [\"AbstractFFTs\", \"Adapt\", \"LinearAlgebra\", \"Printf\", \"Random\", \"Serialization\"]\ngit-tree-sha1 = \"f99a25fe0313121f2f9627002734c7d63b4dd3bd\"\nuuid = \"0c68f7d7-f131-5f86-a1c3-88cf8149b2d7\"\nversion = \"6.2.0\"\n\n[[GPUCompiler]]\ndeps = [\"DataStructures\", \"InteractiveUtils\", \"LLVM\", \"Libdl\", \"Scratch\", \"Serialization\", \"TimerOutputs\", \"UUIDs\"]\ngit-tree-sha1 = \"c853c810b52a80f9aad79ab109207889e57f41ef\"\nuuid = \"61eb1bfa-7361-4325-ad38-22787b887f55\"\nversion = \"0.8.3\"\n\n[[GaussQuadrature]]\ndeps = [\"SpecialFunctions\"]\ngit-tree-sha1 = \"ce3079d0172eaa258f31c30dec9ae045092447d9\"\nuuid = \"d54b0c1a-921d-58e0-8e36-89d8069c0969\"\nversion = \"0.5.5\"\n\n[[GenericSVD]]\ndeps = [\"LinearAlgebra\"]\ngit-tree-sha1 = \"62909c3eda8a25b5673a367d1ad2392ebb265211\"\nuuid = \"01680d73-4ee2-5a08-a1aa-533608c188bb\"\nversion = \"0.3.0\"\n\n[[GenericSchur]]\ndeps = [\"LinearAlgebra\", \"Printf\"]\ngit-tree-sha1 = \"372e48d7f3ced17fdc888a841bcce77be417ce57\"\nuuid = \"c145ed77-6b09-5dd9-b285-bf645a82121e\"\nversion = \"0.5.0\"\n\n[[HDF5_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"LibCURL_jll\", \"Libdl\", \"OpenSSL_jll\", \"Pkg\", \"Zlib_jll\"]\ngit-tree-sha1 = \"fd83fa0bde42e01952757f01149dd968c06c4dba\"\nuuid = \"0234f1f7-429e-5d53-9886-15a909be8d59\"\nversion = \"1.12.0+1\"\n\n[[HTTP]]\ndeps = [\"Base64\", \"Dates\", \"IniFile\", \"MbedTLS\", \"NetworkOptions\", \"Sockets\", \"URIs\"]\ngit-tree-sha1 = \"c9f380c76d8aaa1fa7ea9cf97bddbc0d5b15adc2\"\nuuid = \"cd3eb016-35fb-5094-929b-558a96fad6f3\"\nversion = \"0.9.5\"\n\n[[Hwloc]]\ndeps = [\"Hwloc_jll\"]\ngit-tree-sha1 = \"ffdcd4272a7cc36442007bca41aa07ca3cc5fda4\"\nuuid = \"0e44f5e4-bd66-52a0-8798-143a42290a1d\"\nversion = \"1.3.0\"\n\n[[Hwloc_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"aac91e34ef4c166e0857e3d6052a3467e5732ceb\"\nuuid = \"e33a78d0-f292-5ffc-b300-72abe9b543c8\"\nversion = \"2.4.1+0\"\n\n[[IOCapture]]\ndeps = [\"Logging\"]\ngit-tree-sha1 = \"1868e4e7ad2f93d8de0904d89368c527b46aa6a1\"\nuuid = \"b5f81e59-6552-4d32-b1f0-c071b021bf89\"\nversion = \"0.2.1\"\n\n[[IfElse]]\ngit-tree-sha1 = \"28e837ff3e7a6c3cdb252ce49fb412c8eb3caeef\"\nuuid = \"615f187c-cbe4-4ef1-ba3b-2fcf58d6d173\"\nversion = \"0.1.0\"\n\n[[Inflate]]\ngit-tree-sha1 = \"f5fc07d4e706b84f72d54eedcc1c13d92fb0871c\"\nuuid = \"d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9\"\nversion = \"0.1.2\"\n\n[[IniFile]]\ndeps = [\"Test\"]\ngit-tree-sha1 = \"098e4d2c533924c921f9f9847274f2ad89e018b8\"\nuuid = \"83e8ac13-25f8-5344-8a64-a9f2b223428f\"\nversion = \"0.5.0\"\n\n[[IntelOpenMP_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"d979e54b71da82f3a65b62553da4fc3d18c9004c\"\nuuid = \"1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0\"\nversion = \"2018.0.3+2\"\n\n[[InteractiveUtils]]\ndeps = [\"Markdown\"]\nuuid = \"b77e0a4c-d291-57a0-90e8-8db25a27a240\"\n\n[[Intervals]]\ndeps = [\"Dates\", \"Printf\", \"RecipesBase\", \"Serialization\", \"TimeZones\"]\ngit-tree-sha1 = \"323a38ed1952d30586d0fe03412cde9399d3618b\"\nuuid = \"d8418881-c3e1-53bb-8760-2df7ec849ed5\"\nversion = \"1.5.0\"\n\n[[IterativeSolvers]]\ndeps = [\"LinearAlgebra\", \"Printf\", \"Random\", \"RecipesBase\", \"SparseArrays\"]\ngit-tree-sha1 = \"6f5ef3206d9dc6510a8b8e2334b96454a2ade590\"\nuuid = \"42fd0dbc-a981-5370-80f2-aaf504508153\"\nversion = \"0.9.0\"\n\n[[IteratorInterfaceExtensions]]\ngit-tree-sha1 = \"a3f24677c21f5bbe9d2a714f95dcd58337fb2856\"\nuuid = \"82899510-4779-5014-852e-03e436cf321d\"\nversion = \"1.0.0\"\n\n[[JLD2]]\ndeps = [\"CodecZlib\", \"DataStructures\", \"MacroTools\", \"Mmap\", \"Pkg\", \"Printf\", \"Requires\", \"UUIDs\"]\ngit-tree-sha1 = \"b8343a7f96591404ade118b3a7014e1a52062465\"\nuuid = \"033835bb-8acc-5ee8-8aae-3f567f8a3819\"\nversion = \"0.4.2\"\n\n[[JLLWrappers]]\ngit-tree-sha1 = \"a431f5f2ca3f4feef3bd7a5e94b8b8d4f2f647a0\"\nuuid = \"692b3bcd-3c85-4b1f-b108-f13ce0eb3210\"\nversion = \"1.2.0\"\n\n[[JSON]]\ndeps = [\"Dates\", \"Mmap\", \"Parsers\", \"Unicode\"]\ngit-tree-sha1 = \"81690084b6198a2e1da36fcfda16eeca9f9f24e4\"\nuuid = \"682c06a0-de6a-54ab-a142-c8b1cf79cde6\"\nversion = \"0.21.1\"\n\n[[KernelAbstractions]]\ndeps = [\"Adapt\", \"CUDA\", \"Cassette\", \"InteractiveUtils\", \"MacroTools\", \"SpecialFunctions\", \"StaticArrays\", \"UUIDs\"]\ngit-tree-sha1 = \"ee7f03c23d874c8353813a44315daf82a1e82046\"\nuuid = \"63c18a36-062a-441e-b654-da1e3ab1ce7c\"\nversion = \"0.5.3\"\n\n[[LLVM]]\ndeps = [\"CEnum\", \"Libdl\", \"Printf\", \"Unicode\"]\ngit-tree-sha1 = \"b616937c31337576360cb9fb872ec7633af7b194\"\nuuid = \"929cbde3-209d-540e-8aea-75f648917ca0\"\nversion = \"3.6.0\"\n\n[[LabelledArrays]]\ndeps = [\"ArrayInterface\", \"LinearAlgebra\", \"MacroTools\", \"StaticArrays\"]\ngit-tree-sha1 = \"5e288800819c323de5897fa6d5a002bdad54baf7\"\nuuid = \"2ee39098-c373-598a-b85f-a56591580800\"\nversion = \"1.5.0\"\n\n[[LambertW]]\ngit-tree-sha1 = \"2d9f4009c486ef676646bca06419ac02061c088e\"\nuuid = \"984bce1d-4616-540c-a9ee-88d1112d94c9\"\nversion = \"0.4.5\"\n\n[[LazyArrays]]\ndeps = [\"ArrayLayouts\", \"FillArrays\", \"LinearAlgebra\", \"MacroTools\", \"MatrixFactorizations\", \"SparseArrays\", \"StaticArrays\"]\ngit-tree-sha1 = \"91abe45baaaf05f855215b3221d02f06f96734e1\"\nuuid = \"5078a376-72f3-5289-bfd5-ec5146d43c02\"\nversion = \"0.20.9\"\n\n[[LazyArtifacts]]\ndeps = [\"Pkg\"]\ngit-tree-sha1 = \"4bb5499a1fc437342ea9ab7e319ede5a457c0968\"\nuuid = \"4af54fe1-eca0-43a8-85a7-787d91b784e3\"\nversion = \"1.3.0\"\n\n[[LibCURL]]\ndeps = [\"LibCURL_jll\", \"MozillaCACerts_jll\"]\ngit-tree-sha1 = \"cdbe7465ab7b52358804713a53c7fe1dac3f8a3f\"\nuuid = \"b27032c2-a3e7-50c8-80cd-2d36dbcbfd21\"\nversion = \"0.6.3\"\n\n[[LibCURL_jll]]\ndeps = [\"LibSSH2_jll\", \"Libdl\", \"MbedTLS_jll\", \"Pkg\", \"Zlib_jll\", \"nghttp2_jll\"]\ngit-tree-sha1 = \"897d962c20031e6012bba7b3dcb7a667170dad17\"\nuuid = \"deac9b47-8bc7-5906-a0fe-35ac56dc84c0\"\nversion = \"7.70.0+2\"\n\n[[LibGit2]]\ndeps = [\"Printf\"]\nuuid = \"76f85450-5226-5b5a-8eaa-529ad045b433\"\n\n[[LibSSH2_jll]]\ndeps = [\"Libdl\", \"MbedTLS_jll\", \"Pkg\"]\ngit-tree-sha1 = \"717705533148132e5466f2924b9a3657b16158e8\"\nuuid = \"29816b5a-b9ab-546f-933c-edad1886dfa8\"\nversion = \"1.9.0+3\"\n\n[[Libdl]]\nuuid = \"8f399da3-3557-5675-b5ff-fb832c97cbdb\"\n\n[[Libiconv_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"cba7b560fcc00f8cd770fa85a498cbc1d63ff618\"\nuuid = \"94ce4f54-9a6c-5748-9c1c-f9c7231a4531\"\nversion = \"1.16.0+8\"\n\n[[LightGraphs]]\ndeps = [\"ArnoldiMethod\", \"DataStructures\", \"Distributed\", \"Inflate\", \"LinearAlgebra\", \"Random\", \"SharedArrays\", \"SimpleTraits\", \"SparseArrays\", \"Statistics\"]\ngit-tree-sha1 = \"432428df5f360964040ed60418dd5601ecd240b6\"\nuuid = \"093fc24a-ae57-5d10-9952-331d41423f4d\"\nversion = \"1.3.5\"\n\n[[LightXML]]\ndeps = [\"Libdl\", \"XML2_jll\"]\ngit-tree-sha1 = \"e129d9391168c677cd4800f5c0abb1ed8cb3794f\"\nuuid = \"9c8b4983-aa76-5018-a973-4c85ecc9e179\"\nversion = \"0.9.0\"\n\n[[LineSearches]]\ndeps = [\"LinearAlgebra\", \"NLSolversBase\", \"NaNMath\", \"Parameters\", \"Printf\"]\ngit-tree-sha1 = \"f27132e551e959b3667d8c93eae90973225032dd\"\nuuid = \"d3d80556-e9d4-5f37-9878-2ab0fcc64255\"\nversion = \"7.1.1\"\n\n[[LinearAlgebra]]\ndeps = [\"Libdl\"]\nuuid = \"37e2e46d-f89d-539d-b4ee-838fcccc9c8e\"\n\n[[Literate]]\ndeps = [\"Base64\", \"IOCapture\", \"JSON\", \"REPL\"]\ngit-tree-sha1 = \"501a1a74a0c825037860d36d87d703e987d39dbc\"\nuuid = \"98b081ad-f1c9-55d3-8b20-4c87d4299306\"\nversion = \"2.8.1\"\n\n[[Logging]]\nuuid = \"56ddb016-857b-54e1-b83d-db4d58db5568\"\n\n[[LoopVectorization]]\ndeps = [\"ArrayInterface\", \"DocStringExtensions\", \"IfElse\", \"LinearAlgebra\", \"OffsetArrays\", \"Requires\", \"SLEEFPirates\", \"ThreadingUtilities\", \"UnPack\", \"VectorizationBase\"]\ngit-tree-sha1 = \"5684e4aafadaf668dce27f12d67df4888fa58181\"\nuuid = \"bdcacae8-1622-11e9-2a5c-532679323890\"\nversion = \"0.11.2\"\n\n[[MKL_jll]]\ndeps = [\"Artifacts\", \"IntelOpenMP_jll\", \"JLLWrappers\", \"LazyArtifacts\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"5455aef09b40e5020e1520f551fa3135040d4ed0\"\nuuid = \"856f044c-d86e-5d09-b602-aeab76dc8ba7\"\nversion = \"2021.1.1+2\"\n\n[[MPI]]\ndeps = [\"Distributed\", \"DocStringExtensions\", \"Libdl\", \"MPICH_jll\", \"MicrosoftMPI_jll\", \"OpenMPI_jll\", \"Pkg\", \"Random\", \"Requires\", \"Serialization\", \"Sockets\"]\ngit-tree-sha1 = \"38d0d0255db2316077f7d5dcf8f40c3940e8d534\"\nuuid = \"da04e1cc-30fd-572f-bb4f-1f8673147195\"\nversion = \"0.17.0\"\n\n[[MPICH_jll]]\ndeps = [\"CompilerSupportLibraries_jll\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"4d37f1e07b4e2a74462eebf9ee48c626d15ffdac\"\nuuid = \"7cb0a576-ebde-5e09-9194-50597f1243b4\"\nversion = \"3.3.2+10\"\n\n[[MacroTools]]\ndeps = [\"Markdown\", \"Random\"]\ngit-tree-sha1 = \"6a8a2a625ab0dea913aba95c11370589e0239ff0\"\nuuid = \"1914dd2f-81c6-5fcd-8719-6d5c9610ff09\"\nversion = \"0.5.6\"\n\n[[Markdown]]\ndeps = [\"Base64\"]\nuuid = \"d6f4376e-aef5-505a-96c1-9c027394607a\"\n\n[[MatrixFactorizations]]\ndeps = [\"ArrayLayouts\", \"LinearAlgebra\", \"Printf\", \"Random\"]\ngit-tree-sha1 = \"4154951579535cfba4d716b96dedd9d0beaefcb9\"\nuuid = \"a3b82374-2e81-5b9e-98ce-41277c0e4c87\"\nversion = \"0.8.2\"\n\n[[MbedTLS]]\ndeps = [\"Dates\", \"MbedTLS_jll\", \"Random\", \"Sockets\"]\ngit-tree-sha1 = \"1c38e51c3d08ef2278062ebceade0e46cefc96fe\"\nuuid = \"739be429-bea8-5141-9913-cc70e7f3736d\"\nversion = \"1.0.3\"\n\n[[MbedTLS_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"0eef589dd1c26a3ac9d753fe1a8bcad63f956fa6\"\nuuid = \"c8ffd9c3-330d-5841-b78e-0817d7145fa1\"\nversion = \"2.16.8+1\"\n\n[[MicrosoftMPI_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"e5c90234b3967684c9c6f87b4a54549b4ce21836\"\nuuid = \"9237b28f-5490-5468-be7b-bb81f5f5e6cf\"\nversion = \"10.1.3+0\"\n\n[[Missings]]\ndeps = [\"DataAPI\"]\ngit-tree-sha1 = \"f8c673ccc215eb50fcadb285f522420e29e69e1c\"\nuuid = \"e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28\"\nversion = \"0.4.5\"\n\n[[Mmap]]\nuuid = \"a63ad114-7e13-5084-954f-fe012c677804\"\n\n[[Mocking]]\ndeps = [\"ExprTools\"]\ngit-tree-sha1 = \"916b850daad0d46b8c71f65f719c49957e9513ed\"\nuuid = \"78c3b35d-d492-501b-9361-3d52fe80e533\"\nversion = \"0.7.1\"\n\n[[MozillaCACerts_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"bbcac5afd9049834366c3b68d792971e3d981799\"\nuuid = \"14a3606d-f60d-562e-9121-12d972cd8159\"\nversion = \"2020.10.14+0\"\n\n[[MuladdMacro]]\ngit-tree-sha1 = \"c6190f9a7fc5d9d5915ab29f2134421b12d24a68\"\nuuid = \"46d2c3a1-f734-5fdb-9937-b9b9aeba4221\"\nversion = \"0.2.2\"\n\n[[NCDatasets]]\ndeps = [\"CFTime\", \"DataStructures\", \"Dates\", \"NetCDF_jll\", \"Printf\"]\ngit-tree-sha1 = \"b71d83c87d80f5c54c55a7a9a3aa42bf931c72aa\"\nuuid = \"85f8d34a-cbdd-5861-8df4-14fed0d494ab\"\nversion = \"0.11.3\"\n\n[[NLSolversBase]]\ndeps = [\"DiffResults\", \"Distributed\", \"FiniteDiff\", \"ForwardDiff\"]\ngit-tree-sha1 = \"50608f411a1e178e0129eab4110bd56efd08816f\"\nuuid = \"d41bc354-129a-5804-8e4c-c37616107c6c\"\nversion = \"7.8.0\"\n\n[[NLsolve]]\ndeps = [\"Distances\", \"LineSearches\", \"LinearAlgebra\", \"NLSolversBase\", \"Printf\", \"Reexport\"]\ngit-tree-sha1 = \"019f12e9a1a7880459d0173c182e6a99365d7ac1\"\nuuid = \"2774e3e8-f4cf-5e23-947b-6d7e65073b56\"\nversion = \"4.5.1\"\n\n[[NNlib]]\ndeps = [\"ChainRulesCore\", \"Compat\", \"LinearAlgebra\", \"Pkg\", \"Requires\", \"Statistics\"]\ngit-tree-sha1 = \"ab1d43fead2ecb9aa5ae460d3d547c2cf8d89461\"\nuuid = \"872c559c-99b0-510c-b3b7-b6c96a88d5cd\"\nversion = \"0.7.17\"\n\n[[NaNMath]]\ngit-tree-sha1 = \"bfe47e760d60b82b66b61d2d44128b62e3a369fb\"\nuuid = \"77ba4419-2d1f-58cd-9bb1-8ffee604a2e3\"\nversion = \"0.3.5\"\n\n[[NetCDF_jll]]\ndeps = [\"Artifacts\", \"HDF5_jll\", \"JLLWrappers\", \"LibCURL_jll\", \"LibSSH2_jll\", \"Libdl\", \"MbedTLS_jll\", \"Pkg\", \"Zlib_jll\", \"nghttp2_jll\"]\ngit-tree-sha1 = \"d5835f95aea3b93965a1a7c06de9aace8cb82d99\"\nuuid = \"7243133f-43d8-5620-bbf4-c2c921802cf3\"\nversion = \"400.701.400+0\"\n\n[[NetworkOptions]]\ngit-tree-sha1 = \"ed3157f48a05543cce9b241e1f2815f7e843d96e\"\nuuid = \"ca575930-c2e3-43a9-ace4-1e988b2c1908\"\nversion = \"1.2.0\"\n\n[[NonlinearSolve]]\ndeps = [\"ArrayInterface\", \"FiniteDiff\", \"ForwardDiff\", \"IterativeSolvers\", \"LinearAlgebra\", \"RecursiveArrayTools\", \"RecursiveFactorization\", \"Reexport\", \"SciMLBase\", \"Setfield\", \"StaticArrays\", \"UnPack\"]\ngit-tree-sha1 = \"ef18e47df4f3917af35be5e5d7f5d97e8a83b0ec\"\nuuid = \"8913a72c-1f9b-4ce2-8d82-65094dcecaec\"\nversion = \"0.3.8\"\n\n[[NonlinearSolvers]]\ndeps = [\"CUDA\", \"DocStringExtensions\", \"ForwardDiff\"]\ngit-tree-sha1 = \"13de2bfa3716485129b59d2fdc1c48904c2cfa15\"\nuuid = \"f4b8ab15-8e73-4e04-9661-b5912071d22b\"\nversion = \"0.1.0\"\n\n[[OffsetArrays]]\ndeps = [\"Adapt\"]\ngit-tree-sha1 = \"b3dfef5f2be7d7eb0e782ba9146a5271ee426e90\"\nuuid = \"6fe1bfb0-de20-5000-8ca7-80f57d26f881\"\nversion = \"1.6.2\"\n\n[[OpenMPI_jll]]\ndeps = [\"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"41b983e26a7ab8c9bf05f7d70c274b817d541b46\"\nuuid = \"fe0851c0-eecd-5654-98d4-656369965a5c\"\nversion = \"4.0.2+2\"\n\n[[OpenSSL_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"71bbbc616a1d710879f5a1021bcba65ffba6ce58\"\nuuid = \"458c3c95-2e84-50aa-8efc-19380b2a3a95\"\nversion = \"1.1.1+6\"\n\n[[OpenSpecFun_jll]]\ndeps = [\"Artifacts\", \"CompilerSupportLibraries_jll\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"9db77584158d0ab52307f8c04f8e7c08ca76b5b3\"\nuuid = \"efe28fd5-8261-553b-a9e1-b2916fc3738e\"\nversion = \"0.5.3+4\"\n\n[[OrderedCollections]]\ngit-tree-sha1 = \"4fa2ba51070ec13fcc7517db714445b4ab986bdf\"\nuuid = \"bac558e1-5e72-5ebc-8fee-abe8a469f55d\"\nversion = \"1.4.0\"\n\n[[OrdinaryDiffEq]]\ndeps = [\"Adapt\", \"ArrayInterface\", \"DataStructures\", \"DiffEqBase\", \"ExponentialUtilities\", \"FastClosures\", \"FiniteDiff\", \"ForwardDiff\", \"LinearAlgebra\", \"Logging\", \"MacroTools\", \"MuladdMacro\", \"NLsolve\", \"RecursiveArrayTools\", \"Reexport\", \"SparseArrays\", \"SparseDiffTools\", \"StaticArrays\", \"UnPack\"]\ngit-tree-sha1 = \"d22a75b8ae5b77543c4e1f8eae1ff01ce1f64453\"\nuuid = \"1dea7af3-3e70-54e6-95c3-0bf5283fa5ed\"\nversion = \"5.52.2\"\n\n[[PDMats]]\ndeps = [\"LinearAlgebra\", \"SparseArrays\", \"SuiteSparse\"]\ngit-tree-sha1 = \"f82a0e71f222199de8e9eb9a09977bd0767d52a0\"\nuuid = \"90014a1f-27ba-587c-ab20-58faa44d9150\"\nversion = \"0.11.0\"\n\n[[PackageCompiler]]\ndeps = [\"Libdl\", \"Pkg\", \"UUIDs\"]\ngit-tree-sha1 = \"d448727c4b86be81b225b738c88d30334fda6779\"\nuuid = \"9b87118b-4619-50d2-8e1e-99f35a4d4d9d\"\nversion = \"1.2.5\"\n\n[[Parameters]]\ndeps = [\"OrderedCollections\", \"UnPack\"]\ngit-tree-sha1 = \"2276ac65f1e236e0a6ea70baff3f62ad4c625345\"\nuuid = \"d96e819e-fc66-5662-9728-84c9c7592b0a\"\nversion = \"0.12.2\"\n\n[[Parsers]]\ndeps = [\"Dates\"]\ngit-tree-sha1 = \"c8abc88faa3f7a3950832ac5d6e690881590d6dc\"\nuuid = \"69de0a69-1ddd-5017-9359-2bf0b02dc9f0\"\nversion = \"1.1.0\"\n\n[[Pkg]]\ndeps = [\"Dates\", \"LibGit2\", \"Libdl\", \"Logging\", \"Markdown\", \"Printf\", \"REPL\", \"Random\", \"SHA\", \"UUIDs\"]\nuuid = \"44cfe95a-1eb2-52ea-b672-e2afdf69b78f\"\n\n[[Polynomials]]\ndeps = [\"Intervals\", \"LinearAlgebra\", \"RecipesBase\"]\ngit-tree-sha1 = \"c6b8b87670b9e765db3001ffe640e0583a5ec317\"\nuuid = \"f27b6e38-b328-58d1-80ce-0feddd5e7a45\"\nversion = \"2.0.3\"\n\n[[PrettyTables]]\ndeps = [\"Crayons\", \"Formatting\", \"Markdown\", \"Reexport\", \"Tables\"]\ngit-tree-sha1 = \"574a6b3ea95f04e8757c0280bb9c29f1a5e35138\"\nuuid = \"08abe8d2-0d0c-5749-adfa-8a2ac140af0d\"\nversion = \"0.11.1\"\n\n[[Printf]]\ndeps = [\"Unicode\"]\nuuid = \"de0858da-6303-5e67-8744-51eddeeeb8d7\"\n\n[[QuadGK]]\ndeps = [\"DataStructures\", \"LinearAlgebra\"]\ngit-tree-sha1 = \"12fbe86da16df6679be7521dfb39fbc861e1dc7b\"\nuuid = \"1fd47b50-473d-5c70-9696-f719f8f3bcdc\"\nversion = \"2.4.1\"\n\n[[Quadmath]]\ndeps = [\"Printf\", \"Random\", \"Requires\"]\ngit-tree-sha1 = \"5a8f74af8eae654086a1d058b4ec94ff192e3de0\"\nuuid = \"be4d8f0f-7fa4-5f49-b795-2f01399ab2dd\"\nversion = \"0.5.5\"\n\n[[REPL]]\ndeps = [\"InteractiveUtils\", \"Markdown\", \"Sockets\"]\nuuid = \"3fa0cd96-eef1-5676-8a61-b3b8758bbffb\"\n\n[[Random]]\ndeps = [\"Serialization\"]\nuuid = \"9a3f8284-a2c9-5f02-9a11-845980a1fd5c\"\n\n[[RecipesBase]]\ngit-tree-sha1 = \"b3fb709f3c97bfc6e948be68beeecb55a0b340ae\"\nuuid = \"3cdcf5f2-1ef4-517c-9805-6587b60abb01\"\nversion = \"1.1.1\"\n\n[[RecursiveArrayTools]]\ndeps = [\"ArrayInterface\", \"LinearAlgebra\", \"RecipesBase\", \"Requires\", \"StaticArrays\", \"Statistics\", \"ZygoteRules\"]\ngit-tree-sha1 = \"271a36e18c8806332b7bd0f57e50fcff0d428b11\"\nuuid = \"731186ca-8d62-57ce-b412-fbd966d074cd\"\nversion = \"2.11.0\"\n\n[[RecursiveFactorization]]\ndeps = [\"LinearAlgebra\", \"LoopVectorization\"]\ngit-tree-sha1 = \"20f0ad1b2760da770d31be71f777740d25807631\"\nuuid = \"f2c3362d-daeb-58d1-803e-2bc74f2840b4\"\nversion = \"0.1.11\"\n\n[[Reexport]]\ngit-tree-sha1 = \"57d8440b0c7d98fc4f889e478e80f268d534c9d5\"\nuuid = \"189a3867-3050-52da-a836-e630ba90ab69\"\nversion = \"1.0.0\"\n\n[[Requires]]\ndeps = [\"UUIDs\"]\ngit-tree-sha1 = \"4036a3bd08ac7e968e27c203d45f5fff15020621\"\nuuid = \"ae029012-a4dd-5104-9daa-d747884805df\"\nversion = \"1.1.3\"\n\n[[Rmath]]\ndeps = [\"Random\", \"Rmath_jll\"]\ngit-tree-sha1 = \"86c5647b565873641538d8f812c04e4c9dbeb370\"\nuuid = \"79098fc4-a85e-5d69-aa6a-4863f24498fa\"\nversion = \"0.6.1\"\n\n[[Rmath_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"1b7bf41258f6c5c9c31df8c1ba34c1fc88674957\"\nuuid = \"f50d1b31-88e8-58de-be2c-1cc44531875f\"\nversion = \"0.2.2+2\"\n\n[[RootSolvers]]\ndeps = [\"DocStringExtensions\", \"ForwardDiff\", \"Test\"]\ngit-tree-sha1 = \"0e5b394adc5c6fb39b3964bce2a259a44cc312d3\"\nuuid = \"7181ea78-2dcb-4de3-ab41-2b8ab5a31e74\"\nversion = \"0.2.0\"\n\n[[Rotations]]\ndeps = [\"LinearAlgebra\", \"StaticArrays\", \"Statistics\"]\ngit-tree-sha1 = \"2ed8d8a16d703f900168822d83699b8c3c1a5cd8\"\nuuid = \"6038ab10-8711-5258-84ad-4b1120ba62dc\"\nversion = \"1.0.2\"\n\n[[SHA]]\nuuid = \"ea8e919c-243c-51af-8825-aaa63cd721ce\"\n\n[[SLEEFPirates]]\ndeps = [\"IfElse\", \"Libdl\", \"VectorizationBase\"]\ngit-tree-sha1 = \"ab6194c92dcf38036cd9513e4ab12cd76a613da1\"\nuuid = \"476501e8-09a2-5ece-8869-fb82de89a1fa\"\nversion = \"0.6.10\"\n\n[[SciMLBase]]\ndeps = [\"ArrayInterface\", \"CommonSolve\", \"Distributed\", \"DocStringExtensions\", \"IteratorInterfaceExtensions\", \"LinearAlgebra\", \"Logging\", \"RecipesBase\", \"RecursiveArrayTools\", \"StaticArrays\", \"Statistics\", \"Tables\", \"TreeViews\"]\ngit-tree-sha1 = \"617d5ade740dc628884b6a33e1b02b9bb950e9b3\"\nuuid = \"0bca4576-84f4-4d90-8ffe-ffa030f20462\"\nversion = \"1.9.1\"\n\n[[Scratch]]\ndeps = [\"Dates\"]\ngit-tree-sha1 = \"ad4b278adb62d185bbcb6864dc24959ab0627bf6\"\nuuid = \"6c6a2e73-6563-6170-7368-637461726353\"\nversion = \"1.0.3\"\n\n[[Serialization]]\nuuid = \"9e88b42a-f829-5b0c-bbe9-9e923198166b\"\n\n[[Setfield]]\ndeps = [\"ConstructionBase\", \"Future\", \"MacroTools\", \"Requires\"]\ngit-tree-sha1 = \"d5640fc570fb1b6c54512f0bd3853866bd298b3e\"\nuuid = \"efcf1570-3423-57d1-acb7-fd33fddbac46\"\nversion = \"0.7.0\"\n\n[[SharedArrays]]\ndeps = [\"Distributed\", \"Mmap\", \"Random\", \"Serialization\"]\nuuid = \"1a1011a3-84de-559e-8e89-a11a2f7dc383\"\n\n[[SimpleTraits]]\ndeps = [\"InteractiveUtils\", \"MacroTools\"]\ngit-tree-sha1 = \"daf7aec3fe3acb2131388f93a4c409b8c7f62226\"\nuuid = \"699a6c99-e7fa-54fc-8d76-47d257e15c1d\"\nversion = \"0.9.3\"\n\n[[Sockets]]\nuuid = \"6462fe0b-24de-5631-8697-dd941f90decc\"\n\n[[SortingAlgorithms]]\ndeps = [\"DataStructures\", \"Random\", \"Test\"]\ngit-tree-sha1 = \"03f5898c9959f8115e30bc7226ada7d0df554ddd\"\nuuid = \"a2af1166-a08f-5f64-846c-94a0d3cef48c\"\nversion = \"0.3.1\"\n\n[[SparseArrays]]\ndeps = [\"LinearAlgebra\", \"Random\"]\nuuid = \"2f01184e-e22b-5df5-ae63-d93ebab69eaf\"\n\n[[SparseDiffTools]]\ndeps = [\"Adapt\", \"ArrayInterface\", \"Compat\", \"DataStructures\", \"FiniteDiff\", \"ForwardDiff\", \"LightGraphs\", \"LinearAlgebra\", \"Requires\", \"SparseArrays\", \"VertexSafeGraphs\"]\ngit-tree-sha1 = \"d05bc362e3fa1b0e2361594a706fc63ffbd140f3\"\nuuid = \"47a9eef4-7e08-11e9-0b38-333d64bd3804\"\nversion = \"1.13.0\"\n\n[[SpecialFunctions]]\ndeps = [\"ChainRulesCore\", \"OpenSpecFun_jll\"]\ngit-tree-sha1 = \"5919936c0e92cff40e57d0ddf0ceb667d42e5902\"\nuuid = \"276daf66-3868-5448-9aa4-cd146d93841b\"\nversion = \"1.3.0\"\n\n[[Static]]\ndeps = [\"IfElse\"]\ngit-tree-sha1 = \"ddec5466a1d2d7e58adf9a427ba69763661aacf6\"\nuuid = \"aedffcd0-7271-4cad-89d0-dc628f76c6d3\"\nversion = \"0.2.4\"\n\n[[StaticArrays]]\ndeps = [\"LinearAlgebra\", \"Random\", \"Statistics\"]\ngit-tree-sha1 = \"9da72ed50e94dbff92036da395275ed114e04d49\"\nuuid = \"90137ffa-7385-5640-81b9-e52037218182\"\nversion = \"1.0.1\"\n\n[[StaticNumbers]]\ndeps = [\"Requires\"]\ngit-tree-sha1 = \"a0df7d5ade3fd0f0e6c93ad63facc05b12c40e6a\"\nuuid = \"c5e4b96a-f99f-5557-8ed2-dc63ef9b5131\"\nversion = \"0.3.3\"\n\n[[Statistics]]\ndeps = [\"LinearAlgebra\", \"SparseArrays\"]\nuuid = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\n\n[[StatsBase]]\ndeps = [\"DataAPI\", \"DataStructures\", \"LinearAlgebra\", \"Missings\", \"Printf\", \"Random\", \"SortingAlgorithms\", \"SparseArrays\", \"Statistics\"]\ngit-tree-sha1 = \"a83fa3021ac4c5a918582ec4721bc0cf70b495a9\"\nuuid = \"2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91\"\nversion = \"0.33.4\"\n\n[[StatsFuns]]\ndeps = [\"Rmath\", \"SpecialFunctions\"]\ngit-tree-sha1 = \"ced55fd4bae008a8ea12508314e725df61f0ba45\"\nuuid = \"4c63d2b9-4356-54db-8cca-17b64c39e42c\"\nversion = \"0.9.7\"\n\n[[SuiteSparse]]\ndeps = [\"Libdl\", \"LinearAlgebra\", \"Serialization\", \"SparseArrays\"]\nuuid = \"4607b0f0-06f3-5cda-b6b1-a6196a1729e9\"\n\n[[SurfaceFluxes]]\ndeps = [\"CLIMAParameters\", \"DocStringExtensions\", \"KernelAbstractions\", \"NonlinearSolvers\", \"StaticArrays\"]\ngit-tree-sha1 = \"0f2685b633ebf751e50842ee3525880f7d043162\"\nuuid = \"49b00bb7-8bd4-4f2b-b78c-51cd0450215f\"\nversion = \"0.1.1\"\n\n[[TableTraits]]\ndeps = [\"IteratorInterfaceExtensions\"]\ngit-tree-sha1 = \"b1ad568ba658d8cbb3b892ed5380a6f3e781a81e\"\nuuid = \"3783bdb8-4a98-5b6b-af9a-565f29a5fe9c\"\nversion = \"1.0.0\"\n\n[[Tables]]\ndeps = [\"DataAPI\", \"DataValueInterfaces\", \"IteratorInterfaceExtensions\", \"LinearAlgebra\", \"TableTraits\", \"Test\"]\ngit-tree-sha1 = \"a9ff3dfec713c6677af435d6a6d65f9744feef67\"\nuuid = \"bd369af6-aec1-5ad0-b16a-f7cc5008161c\"\nversion = \"1.4.1\"\n\n[[TaylorSeries]]\ndeps = [\"InteractiveUtils\", \"LinearAlgebra\", \"Markdown\", \"Requires\", \"SparseArrays\"]\ngit-tree-sha1 = \"66f4d1993bae49eeba21a1634b5f65782585a42c\"\nuuid = \"6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea\"\nversion = \"0.10.13\"\n\n[[Test]]\ndeps = [\"Distributed\", \"InteractiveUtils\", \"Logging\", \"Random\"]\nuuid = \"8dfed614-e22c-5e08-85e1-65c5234f0b40\"\n\n[[TextWrap]]\ngit-tree-sha1 = \"9250ef9b01b66667380cf3275b3f7488d0e25faf\"\nuuid = \"b718987f-49a8-5099-9789-dcd902bef87d\"\nversion = \"1.0.1\"\n\n[[Thermodynamics]]\ndeps = [\"CLIMAParameters\", \"DocStringExtensions\", \"ExprTools\", \"KernelAbstractions\", \"Random\", \"RootSolvers\"]\ngit-tree-sha1 = \"c7d73eae235caffdab85778d8b6c371691ad1b3e\"\nuuid = \"b60c26fb-14c3-4610-9d3e-2d17fe7ff00c\"\nversion = \"0.5.1\"\n\n[[ThreadingUtilities]]\ndeps = [\"VectorizationBase\"]\ngit-tree-sha1 = \"e3032c97b183e6e2baf4d2cc4fe60c4292a4a707\"\nuuid = \"8290d209-cae3-49c0-8002-c8c24d57dab5\"\nversion = \"0.2.5\"\n\n[[TimeZones]]\ndeps = [\"Dates\", \"EzXML\", \"Mocking\", \"Pkg\", \"Printf\", \"RecipesBase\", \"Serialization\", \"Unicode\"]\ngit-tree-sha1 = \"4ba8a9579a243400db412b50300cd61d7447e583\"\nuuid = \"f269a46b-ccf7-5d73-abea-4c690281aa53\"\nversion = \"1.5.3\"\n\n[[TimerOutputs]]\ndeps = [\"Printf\"]\ngit-tree-sha1 = \"32cdbe6cd2d214c25a0b88f985c9e0092877c236\"\nuuid = \"a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f\"\nversion = \"0.5.8\"\n\n[[TranscodingStreams]]\ndeps = [\"Random\", \"Test\"]\ngit-tree-sha1 = \"7c53c35547de1c5b9d46a4797cf6d8253807108c\"\nuuid = \"3bb67fe8-82b1-5028-8e26-92a6c54297fa\"\nversion = \"0.9.5\"\n\n[[TreeViews]]\ndeps = [\"Test\"]\ngit-tree-sha1 = \"8d0d7a3fe2f30d6a7f833a5f19f7c7a5b396eae6\"\nuuid = \"a2a6695c-b41b-5b7d-aed9-dbfdeacea5d7\"\nversion = \"0.3.0\"\n\n[[URIs]]\ngit-tree-sha1 = \"7855809b88d7b16e9b029afd17880930626f54a2\"\nuuid = \"5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4\"\nversion = \"1.2.0\"\n\n[[UUIDs]]\ndeps = [\"Random\", \"SHA\"]\nuuid = \"cf7118a7-6976-5b1a-9a39-7adc72f591a4\"\n\n[[UnPack]]\ngit-tree-sha1 = \"387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b\"\nuuid = \"3a884ed6-31ef-47d7-9d2a-63182c4928ed\"\nversion = \"1.0.2\"\n\n[[Unicode]]\nuuid = \"4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5\"\n\n[[VectorizationBase]]\ndeps = [\"ArrayInterface\", \"Hwloc\", \"IfElse\", \"Libdl\", \"LinearAlgebra\"]\ngit-tree-sha1 = \"486842a62c4a1bc23f7c8457d64e683a00d6d0e9\"\nuuid = \"3d5dd08c-fd9d-11e8-17fa-ed2836048c2f\"\nversion = \"0.18.14\"\n\n[[VertexSafeGraphs]]\ndeps = [\"LightGraphs\"]\ngit-tree-sha1 = \"b9b450c99a3ca1cc1c6836f560d8d887bcbe356e\"\nuuid = \"19fa3120-7c27-5ec5-8db8-b0b0aa330d6f\"\nversion = \"0.1.2\"\n\n[[WriteVTK]]\ndeps = [\"Base64\", \"CodecZlib\", \"FillArrays\", \"LightXML\", \"Random\", \"TranscodingStreams\"]\ngit-tree-sha1 = \"37eef911a9c5211e0ae4362dc0477cfe6c537ffa\"\nuuid = \"64499a7a-5c06-52f2-abe2-ccb03c286192\"\nversion = \"1.9.1\"\n\n[[XML2_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Libiconv_jll\", \"Pkg\", \"Zlib_jll\"]\ngit-tree-sha1 = \"be0db24f70aae7e2b89f2f3092e93b8606d659a6\"\nuuid = \"02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a\"\nversion = \"2.9.10+3\"\n\n[[Zlib_jll]]\ndeps = [\"Artifacts\", \"JLLWrappers\", \"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"320228915c8debb12cb434c59057290f0834dbf6\"\nuuid = \"83775a58-1f1d-513f-b197-d71354ab007a\"\nversion = \"1.2.11+18\"\n\n[[ZygoteRules]]\ndeps = [\"MacroTools\"]\ngit-tree-sha1 = \"9e7a1e8ca60b742e508a315c17eef5211e7fbfd7\"\nuuid = \"700de1a5-db45-46bc-99cf-38207098b444\"\nversion = \"0.2.1\"\n\n[[nghttp2_jll]]\ndeps = [\"Libdl\", \"Pkg\"]\ngit-tree-sha1 = \"8e2c44ab4d49ad9518f359ed8b62f83ba8beede4\"\nuuid = \"8e850ede-7688-5339-a07c-302acd2aaf8d\"\nversion = \"1.40.0+2\"\n" }, { "path": "Project.toml", "content": "name = \"ClimateMachine\"\nuuid = \"777c4786-024f-11e9-21a3-85d5d4106250\"\nauthors = [\"Climate Modeling Alliance\"]\nversion = \"0.3.0-DEV\"\n\n[deps]\nAdapt = \"79e6a3ab-5dfb-504d-930d-738a2a938a0e\"\nArgParse = \"c7e460c6-2fb9-53a9-8c5b-16f535851c63\"\nArtifactWrappers = \"a14bc488-3040-4b00-9dc1-f6467924858a\"\nBenchmarkTools = \"6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf\"\nCLIMAParameters = \"6eacf6c3-8458-43b9-ae03-caf5306d3d53\"\nCUDA = \"052768ef-5323-5732-b1bb-66c8b64840ba\"\nCloudMicrophysics = \"6a9e3e04-43cd-43ba-94b9-e8782df3c71b\"\nCombinatorics = \"861a8166-3701-5b0c-9a16-15d98fcdc6aa\"\nCoverage = \"a2441757-f6aa-5fb2-8edb-039e3f45d037\"\nCubedSphere = \"7445602f-e544-4518-8976-18f8e8ae6cdb\"\nDates = \"ade2ca70-3891-5945-98fb-dc099432e06a\"\nDelimitedFiles = \"8bb1440f-4735-579b-a4ab-409b98df4dab\"\nDierckx = \"39dd38d3-220a-591b-8e3c-4c3a8c710a94\"\nDiffEqBase = \"2b5f629d-d688-5b77-993f-72d75c75574e\"\nDispatchedTuples = \"508c55e1-51b4-41fd-a5ca-7eb0327d070d\"\nDistributions = \"31c24e10-a181-5473-b8eb-7969acd0382f\"\nDocStringExtensions = \"ffbed154-4ef7-542d-bbb7-c09d3a79fcae\"\nDoubleFloats = \"497a8b3b-efae-58df-a0af-a86822472b78\"\nDownloads = \"f43a241f-c20a-4ad4-852c-f6b1247861c6\"\nFFTW = \"7a1cc6ca-52ef-59f5-83cd-3a7055c09341\"\nFileIO = \"5789e2e9-d7fb-5bc7-8068-2c6fae9b9549\"\nFormatting = \"59287772-0a20-5a39-b81b-1366585eb4c0\"\nForwardDiff = \"f6369f11-7733-5829-9624-2563aa707210\"\nGaussQuadrature = \"d54b0c1a-921d-58e0-8e36-89d8069c0969\"\nInteractiveUtils = \"b77e0a4c-d291-57a0-90e8-8db25a27a240\"\nJLD2 = \"033835bb-8acc-5ee8-8aae-3f567f8a3819\"\nKernelAbstractions = \"63c18a36-062a-441e-b654-da1e3ab1ce7c\"\nLambertW = \"984bce1d-4616-540c-a9ee-88d1112d94c9\"\nLazyArrays = \"5078a376-72f3-5289-bfd5-ec5146d43c02\"\nLinearAlgebra = \"37e2e46d-f89d-539d-b4ee-838fcccc9c8e\"\nLiterate = \"98b081ad-f1c9-55d3-8b20-4c87d4299306\"\nLogging = \"56ddb016-857b-54e1-b83d-db4d58db5568\"\nMPI = \"da04e1cc-30fd-572f-bb4f-1f8673147195\"\nMacroTools = \"1914dd2f-81c6-5fcd-8719-6d5c9610ff09\"\nNCDatasets = \"85f8d34a-cbdd-5861-8df4-14fed0d494ab\"\nNLsolve = \"2774e3e8-f4cf-5e23-947b-6d7e65073b56\"\nNonlinearSolvers = \"f4b8ab15-8e73-4e04-9661-b5912071d22b\"\nOrderedCollections = \"bac558e1-5e72-5ebc-8fee-abe8a469f55d\"\nOrdinaryDiffEq = \"1dea7af3-3e70-54e6-95c3-0bf5283fa5ed\"\nPackageCompiler = \"9b87118b-4619-50d2-8e1e-99f35a4d4d9d\"\nPkg = \"44cfe95a-1eb2-52ea-b672-e2afdf69b78f\"\nPrettyTables = \"08abe8d2-0d0c-5749-adfa-8a2ac140af0d\"\nPrintf = \"de0858da-6303-5e67-8744-51eddeeeb8d7\"\nRandom = \"9a3f8284-a2c9-5f02-9a11-845980a1fd5c\"\nRootSolvers = \"7181ea78-2dcb-4de3-ab41-2b8ab5a31e74\"\nRotations = \"6038ab10-8711-5258-84ad-4b1120ba62dc\"\nSpecialFunctions = \"276daf66-3868-5448-9aa4-cd146d93841b\"\nStaticArrays = \"90137ffa-7385-5640-81b9-e52037218182\"\nStaticNumbers = \"c5e4b96a-f99f-5557-8ed2-dc63ef9b5131\"\nStatistics = \"10745b16-79ce-11e8-11f9-7d13ad32a3b2\"\nSurfaceFluxes = \"49b00bb7-8bd4-4f2b-b78c-51cd0450215f\"\nThermodynamics = \"b60c26fb-14c3-4610-9d3e-2d17fe7ff00c\"\nUnPack = \"3a884ed6-31ef-47d7-9d2a-63182c4928ed\"\nWriteVTK = \"64499a7a-5c06-52f2-abe2-ccb03c286192\"\n\n[compat]\nAdapt = \"2.0.2, 3.2\"\nArgParse = \"1.1\"\nArtifactWrappers = \"0.1.1\"\nBenchmarkTools = \"0.5\"\nCLIMAParameters = \"0.2\"\nCUDA = \"2.0\"\nCloudMicrophysics = \"0.3\"\nCombinatorics = \"1.0\"\nCoverage = \"1.0\"\nDierckx = \"0.4, 0.5\"\nDiffEqBase = \"6.47\"\nDispatchedTuples = \"0.2\"\nDistributions = \"0.22, 0.23, 0.24\"\nDocStringExtensions = \"0.8\"\nDoubleFloats = \"1.1\"\nFFTW = \"1.2\"\nFileIO = \"1.2\"\nFormatting = \"0.4\"\nForwardDiff = \"0.10\"\nGaussQuadrature = \"0.5\"\nJLD2 = \"0.1, 0.2, 0.3, 0.4\"\nKernelAbstractions = \"0.4.1, 0.5\"\nLambertW = \"0.4\"\nLazyArrays = \"0.15, 0.16, 0.18, 0.19, 0.20\"\nLiterate = \"2.2\"\nMPI = \"0.16, 0.17\"\nNCDatasets = \"0.10, 0.11\"\nNLsolve = \"4.4\"\nNonlinearSolvers = \"0.1\"\nOrderedCollections = \"1.1\"\nOrdinaryDiffEq = \"5.41.0\"\nPackageCompiler = \"1.2\"\nPrettyTables = \"0.9, 0.10, 0.11\"\nRootSolvers = \"0.1, 0.2\"\nSpecialFunctions = \"0.10, 1.0\"\nStaticArrays = \"0.12, 1.0\"\nStaticNumbers = \"0.3.2\"\nSurfaceFluxes = \"0.1\"\nThermodynamics = \"0.5.1\"\nUnPack = \"1.0\"\nWriteVTK = \"1.7\"\njulia = \"1.5\"\n\n[extras]\nTest = \"8dfed614-e22c-5e08-85e1-65c5234f0b40\"\n\n[targets]\ntest = [\"Test\"]\n" }, { "path": "README.md", "content": "# ClimateMachine.jl\n\n***NOTE THAT THIS REPO IS NOT CURRENTLY BEING MAINTAINED. PLEASE SEE THE FOLLOWING REPOSITORIES:***\n\n - [CliMA/ClimaCore.jl](http://github.com/CliMA/ClimaCore.jl)\n - [CliMA/ClimaAtmos.jl](http://github.com/CliMA/ClimaAtmos.jl)\n - [CliMA/Thermodynamics.jl](http://github.com/CliMA/Thermodynamics.jl)\n - [CliMA/CloudMicrophysics.jl](http://github.com/CliMA/CloudMicrophysics.jl)\n - [CliMA/SurfaceFluxes.jl](http://github.com/CliMA/SurfaceFluxes.jl)\n - [CliMA/ClimaLSM.jl](http://github.com/CliMA/ClimaLSM.jl)\n - [CliMA/Oceananigans.jl](http://github.com/CliMA/Oceananigans.jl)\n - [CliMA/ClimaCoupler.jl](http://github.com/CliMA/ClimaCoupler.jl)\n\nThe Climate Machine is a new Earth system model that leverages recent advances in the computational and data sciences to learn directly from a wealth of Earth observations from space and the ground. The Climate Machine will harness more data than ever before, providing a new level of accuracy to predictions of droughts, heat waves, and rainfall extremes.\n\n| **Documentation** | [![dev][docs-latest-img]][docs-latest-url] |\n|----------------------|--------------------------------------------------|\n| **Docs Build** | [![docs build][docs-bld-img]][docs-bld-url] |\n| **Unit tests** | [![unit tests][unit-tests-img]][unit-tests-url] |\n| **Code Coverage** | [![codecov][codecov-img]][codecov-url] |\n| **Bors** | [![Bors enabled][bors-img]][bors-url] |\n| **DOI** | [![Zenodo][zenodo-img]][zenodo-url] |\n\n[docs-bld-img]: https://github.com/CliMA/ClimateMachine.jl/workflows/Documentation/badge.svg\n[docs-bld-url]: https://github.com/CliMA/ClimateMachine.jl/actions?query=workflow%3ADocumentation\n\n[docs-latest-img]: https://img.shields.io/badge/docs-latest-blue.svg\n[docs-latest-url]: https://CliMA.github.io/ClimateMachine.jl/latest/\n\n[unit-tests-img]: https://github.com/CliMA/ClimateMachine.jl/workflows/OS%20Unit%20Tests/badge.svg\n[unit-tests-url]: https://github.com/CliMA/ClimateMachine.jl/actions?query=workflow%3A%22OS+Unit+Tests%22\n\n[codecov-img]: https://codecov.io/gh/CliMA/ClimateMachine.jl/branch/master/graph/badge.svg\n[codecov-url]: https://codecov.io/gh/CliMA/ClimateMachine.jl\n\n[bors-img]: https://bors.tech/images/badge_small.svg\n[bors-url]: https://app.bors.tech/repositories/11521\n\n[zenodo-img]: https://zenodo.org/badge/162166244.svg\n[zenodo-url]: https://zenodo.org/badge/latestdoi/162166244\n\nFor installation instructions and explanations on how to use the Climate Machine, please look at the [Documentation](https://clima.github.io/ClimateMachine.jl/latest/GettingStarted/Installation/).\n" }, { "path": "bors.toml", "content": "status = [\n \"test-os (ubuntu-latest)\",\n \"test-os (windows-latest)\",\n \"test-os (macos-latest)\",\n \"buildkite/climatemachine-ci\",\n \"buildkite/climatemachine-docs\",\n \"format\"\n]\ndelete_merged_branches = true\ntimeout_sec = 12600\nblock_labels = [ \"do-not-merge-yet\" ]\ncut_body_after = \" H[1] = F G H I J\n axes |/ H[2] = K L M N O\n 0------ X[0] high low\n\nwhere the 15-bit Hilbert integer = `A B C D E F G H I J K L M N O` is stored\nin `H`\n\nThis function is based on public domain code from John Skilling which can be\nfound in [Skilling2004](@cite).\n\"\"\"\nfunction hilbertcode(Y::AbstractArray{T}; bits = 8 * sizeof(T)) where {T}\n # Below is Skilling's AxestoTranspose\n X = deepcopy(Y)\n n = length(X)\n M = one(T) << (bits - 1)\n\n Q = M\n for j in 1:(bits - 1)\n P = Q - one(T)\n for i in 1:n\n if X[i] & Q != zero(T)\n X[1] ⊻= P\n else\n t = (X[1] ⊻ X[i]) & P\n X[1] ⊻= t\n X[i] ⊻= t\n end\n end\n Q >>>= one(T)\n end\n\n for i in 2:n\n X[i] ⊻= X[i - 1]\n end\n\n t = zero(T)\n Q = M\n for j in 1:(bits - 1)\n if X[n] & Q != zero(T)\n t ⊻= Q - one(T)\n end\n Q >>>= one(T)\n end\n\n for i in 1:n\n X[i] ⊻= t\n end\n\n # Below we transpose X and store it in H, i.e.:\n #\n # X[0] = A D G J M H[0] = A B C D E\n # X[1] = B E H K N <-------> H[1] = F G H I J\n # X[2] = C F I L O H[2] = K L M N O\n #\n # The 15-bit Hilbert integer is then = A B C D E F G H I J K L M N O\n H = zero(X)\n for i in 0:(n - 1), j in 0:(bits - 1)\n k = i * bits + j\n bit = (X[n - mod(k, n)] >>> div(k, n)) & one(T)\n H[n - i] |= (bit << j)\n end\n\n return H\nend\n\n\"\"\"\n centroidtocode(comm::MPI.Comm, elemtocorner; coortocode, CT)\n\nReturns a code for each element based on its centroid.\n\nThese element codes can be used to determine a linear ordering for the\npartition function.\n\nThe communicator `comm` is used to calculate the bounding box for representing\nthe centroids in coordinates of type `CT`, defaulting to `CT=UInt64`. These\ninteger coordinates are converted to a code using the function `coortocode`,\nwhich defaults to `hilbertcode`.\n\nThe array containing the element corner coordinates, `elemtocorner`, is used\nto compute the centroids. `elemtocorner` is a dimension by number of corners\nby number of elements array.\n\"\"\"\nfunction centroidtocode(\n comm::MPI.Comm,\n elemtocorner;\n coortocode = hilbertcode,\n CT = UInt64,\n)\n (d, nvert, nelem) = size(elemtocorner)\n\n centroids = sum(elemtocorner, dims = 2) ./ nvert\n T = eltype(centroids)\n\n centroidmin =\n (nelem > 0) ? minimum(centroids, dims = 3) : fill(typemax(T), d)\n centroidmax =\n (nelem > 0) ? maximum(centroids, dims = 3) : fill(typemin(T), d)\n\n centroidmin = MPI.Allreduce(centroidmin, min, comm)\n centroidmax = MPI.Allreduce(centroidmax, max, comm)\n centroidsize = centroidmax - centroidmin\n\n # Fix centroidsize to be nonzero. It can be zero for a couple of reasons.\n # For example, it will be zero if we have just one element.\n if iszero(centroidsize)\n centroidsize = ones(T, d)\n else\n for i in 1:d\n if iszero(centroidsize[i])\n centroidsize[i] = maximum(centroidsize)\n end\n end\n end\n\n code = Array{CT}(undef, d, nelem)\n for e in 1:nelem\n c = (centroids[:, 1, e] .- centroidmin) ./ centroidsize\n X =\n CT.(floor.(\n typemax(CT) .* BigFloat.(c; precision = 16 * sizeof(CT)),\n ))\n code[:, e] = coortocode(X)\n end\n\n code\nend\n\n\"\"\"\n brickmesh(x, periodic; part=1, numparts=1; boundary)\n\nGenerate a brick mesh with coordinates given by the tuple `x` and the\nperiodic dimensions given by the `periodic` tuple.\n\nThe brick can optionally be partitioned into `numparts` and this returns\npartition `part`. This is a simple Cartesian partition and further\npartitioning (e.g, based on a space-filling curve) should be done before the\nmesh is used for computation.\n\nBy default boundary faces will be marked with a one and other faces with a\nzero. Specific boundary numbers can also be passed for each face of the brick\nin `boundary`. This will mark the nonperiodic brick faces with the given\nboundary number.\n\n# Examples\n\nWe can build a 3 by 2 element two-dimensional mesh that is periodic in the\n\\$x_2\\$-direction with\n```jldoctest brickmesh\njulia> (elemtovert, elemtocoord, elemtobndy, faceconnections) =\n brickmesh((2:5,4:6), (false,true); boundary=((1,2), (3,4)));\n```\nThis returns the mesh structure for\n\n x_2\n\n ^\n |\n 6- 9----10----11----12\n | | | | |\n | | 4 | 5 | 6 |\n | | | | |\n 5- 5-----6-----7-----8\n | | | | |\n | | 1 | 2 | 3 |\n | | | | |\n 4- 1-----2-----3-----4\n |\n +--|-----|-----|-----|--> x_1\n 2 3 4 5\n\nThe (number of corners by number of elements) array `elemtovert` gives the\nglobal vertex number for the corners of each element.\n```jldoctest brickmesh\njulia> elemtovert\n4×6 Array{Int64,2}:\n 1 2 3 5 6 7\n 2 3 4 6 7 8\n 5 6 7 9 10 11\n 6 7 8 10 11 12\n```\nNote that the vertices are listed in Cartesian order.\n\nThe (dimension by number of corners by number of elements) array `elemtocoord`\ngives the coordinates of the corners of each element.\n```jldoctes brickmesh\njulia> elemtocoord\n2×4×6 Array{Int64,3}:\n[:, :, 1] =\n 2 3 2 3\n 4 4 5 5\n\n[:, :, 2] =\n 3 4 3 4\n 4 4 5 5\n\n[:, :, 3] =\n 4 5 4 5\n 4 4 5 5\n\n[:, :, 4] =\n 2 3 2 3\n 5 5 6 6\n\n[:, :, 5] =\n 3 4 3 4\n 5 5 6 6\n\n[:, :, 6] =\n 4 5 4 5\n 5 5 6 6\n```\n\nThe (number of faces by number of elements) array `elemtobndy` gives the\nboundary number for each face of each element. A zero will be given for\nconnected faces.\n```jldoctest brickmesh\njulia> elemtobndy\n4×6 Array{Int64,2}:\n 1 0 0 1 0 0\n 0 0 2 0 0 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n```\nNote that the faces are listed in Cartesian order.\n\nFinally, the periodic face connections are given in `faceconnections` which is a\nlist of arrays, one for each connection.\nEach array in the list is given in the format `[e, f, vs...]` where\n - `e` is the element number;\n - `f` is the face number; and\n - `vs` is the global vertices that face associated with.\nI the example\n```jldoctest brickmesh\njulia> faceconnections\n3-element Array{Array{Int64,1},1}:\n [4, 4, 1, 2]\n [5, 4, 2, 3]\n [6, 4, 3, 4]\n```\nwe see that face `4` of element `5` is associated with vertices `[2 3]` (the\nvertices for face `1` of element `2`).\n\"\"\"\nfunction brickmesh(\n x,\n periodic;\n part = 1,\n numparts = 1,\n boundary = ntuple(j -> (1, 1), length(x)),\n)\n if boundary isa Matrix\n boundary = tuple(mapslices(x -> tuple(x...), boundary, dims = 1)...)\n end\n\n @assert length(x) == length(periodic)\n @assert length(x) >= 1\n @assert 1 <= part <= numparts\n\n T = promote_type(eltype.(x)...)\n d = length(x)\n nvert = 2^d\n nface = 2d\n\n nelemdim = length.(x) .- 1\n elemlocal = linearpartition(prod(nelemdim), part, numparts)\n\n elemtovert = Array{Int}(undef, nvert, length(elemlocal))\n elemtocoord = Array{T}(undef, d, nvert, length(elemlocal))\n elemtobndy = zeros(Int, nface, length(elemlocal))\n faceconnections = Array{Array{Int, 1}}(undef, 0)\n\n verts = LinearIndices(ntuple(j -> 1:length(x[j]), d))\n elems = CartesianIndices(ntuple(j -> 1:(length(x[j]) - 1), d))\n\n p = reshape(1:nvert, ntuple(j -> 2, d))\n fmask = hcat((\n p[ntuple(\n j ->\n (j == div(f - 1, 2) + 1) ?\n ((mod(f - 1, 2) + 1):(mod(f - 1, 2) + 1)) : (:),\n d,\n )...,][:] for f in 1:nface\n )...)\n\n\n for (e, ec) in enumerate(elems[elemlocal])\n corners = CartesianIndices(ntuple(j -> ec[j]:(ec[j] + 1), d))\n for (v, vc) in enumerate(corners)\n elemtovert[v, e] = verts[vc]\n\n for j in 1:d\n elemtocoord[j, v, e] = x[j][vc[j]]\n end\n end\n\n for i in 1:d\n if !periodic[i] && ec[i] == 1\n elemtobndy[2 * (i - 1) + 1, e] = boundary[i][1]\n end\n if !periodic[i] && ec[i] == nelemdim[i]\n elemtobndy[2 * (i - 1) + 2, e] = boundary[i][2]\n end\n end\n\n for i in 1:d\n if periodic[i] && ec[i] == nelemdim[i]\n neighcorners = CartesianIndices(ntuple(\n j -> (i == j) ? (1:2) : (ec[j]:(ec[j] + 1)),\n d,\n ))\n push!(\n faceconnections,\n vcat(e, 2i, verts[neighcorners[fmask[:, 2i - 1]]]),\n )\n end\n end\n end\n\n (elemtovert, elemtocoord, elemtobndy, faceconnections)\nend\n\n\"\"\"\n parallelsortcolumns(comm::MPI.Comm, A;\n alg::Base.Sort.Algorithm=Base.Sort.DEFAULT_UNSTABLE,\n lt=isless,\n by=identity,\n rev::Union{Bool,Nothing}=nothing)\n\nSorts the columns of the distributed matrix `A`.\n\nSee the documentation of `sort!` for a description of the keyword arguments.\n\nThis function assumes `A` has the same number of rows on each MPI rank but can\nhave a different number of columns.\n\"\"\"\nfunction parallelsortcolumns(\n comm::MPI.Comm,\n A;\n alg::Base.Sort.Algorithm = Base.Sort.DEFAULT_UNSTABLE,\n lt = isless,\n by = identity,\n rev::Union{Bool, Nothing} = nothing,\n)\n\n m, n = size(A)\n T = eltype(A)\n\n csize = MPI.Comm_size(comm)\n crank = MPI.Comm_rank(comm)\n croot = 0\n\n A = sortslices(A, dims = 2, alg = alg, lt = lt, by = by, rev = rev)\n\n npivots = clamp(n, 0, csize)\n pivots = T[A[i, div(n * p, npivots) + 1] for i in 1:m, p in 0:(npivots - 1)]\n pivotcounts = MPI.Allgather(Cint(length(pivots)), comm)\n pivots = MPI.Allgatherv!(\n pivots,\n VBuffer(similar(pivots, sum(pivotcounts)), pivotcounts),\n comm,\n )\n pivots = reshape(pivots, m, div(length(pivots), m))\n pivots =\n sortslices(pivots, dims = 2, alg = alg, lt = lt, by = by, rev = rev)\n\n # if we don't have any pivots then we must have zero columns\n if size(pivots) == (m, 0)\n return A\n end\n\n pivots = [\n pivots[i, div(div(length(pivots), m) * r, csize) + 1]\n for i in 1:m, r in 0:(csize - 1)\n ]\n\n cols = map(i -> view(A, :, i), 1:n)\n senddispls = Cint[\n (\n searchsortedfirst(cols, pivots[:, i], lt = lt, by = by, rev = rev) - 1\n ) * m for i in 1:csize\n ]\n sendcounts = Cint[\n (i == csize ? n * m : senddispls[i + 1]) - senddispls[i]\n for i in 1:csize\n ]\n\n recvcounts = similar(sendcounts)\n MPI.Alltoall!(UBuffer(sendcounts, 1), MPI.UBuffer(recvcounts, 1), comm)\n B = similar(A, sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(A, sendcounts, senddispls),\n VBuffer(B, recvcounts),\n comm,\n )\n B = reshape(B, m, div(length(B), m))\n sortslices(B, dims = 2, alg = alg, lt = lt, by = by, rev = rev)\nend\n\n\"\"\"\n getpartition(comm::MPI.Comm, elemtocode)\n\nReturns an equally weighted partition of a distributed set of elements by\nsorting their codes given in `elemtocode`.\n\nThe codes for each element, `elemtocode`, are given as an array with a single\nentry per local element or as a matrix with a column for each local element.\n\nThe partition is returned as a tuple three parts:\n\n - `partsendorder`: permutation of elements into sending order\n - `partsendstarts`: start entries in the send array for each rank\n - `partrecvstarts`: start entries in the receive array for each rank\n\nNote that both `partsendstarts` and `partrecvstarts` are of length\n`MPI.Comm_size(comm)+1` where the last entry has the total number of elements\nto send or receive, respectively.\n\"\"\"\ngetpartition(comm::MPI.Comm, elemtocode::AbstractVector) =\n getpartition(comm, reshape(elemtocode, 1, length(elemtocode)))\n\nfunction getpartition(comm::MPI.Comm, elemtocode::AbstractMatrix)\n (ncode, nelem) = size(elemtocode)\n\n csize = MPI.Comm_size(comm)\n crank = MPI.Comm_rank(comm)\n\n CT = eltype(elemtocode)\n\n A = CT[\n elemtocode # code\n collect(CT, 1:nelem)' # original element number\n fill(CT(MPI.Comm_rank(comm)), (1, nelem)) # original rank\n fill(typemax(CT), (1, nelem))\n ] # new rank\n m, n = size(A)\n\n # sort by just code\n A = parallelsortcolumns(comm, A)\n\n # count the distribution of A\n counts = MPI.Allgather(last(size(A)), comm)\n starts = ones(Int, csize + 1)\n for i in 1:csize\n starts[i + 1] = counts[i] + starts[i]\n end\n\n # loop to determine new rank\n j = range(starts[crank + 1], stop = starts[crank + 2] - 1)\n for r in 0:(csize - 1)\n k = linearpartition(starts[end] - 1, r + 1, csize)\n o = intersect(k, j) .- (starts[crank + 1] - 1)\n A[ncode + 3, o] .= r\n end\n\n # sort by original rank and code\n A = sortslices(A, dims = 2, by = x -> x[[ncode + 2, (1:ncode)...]])\n\n # count number of elements that are going to be sent\n sendcounts = zeros(Cint, csize)\n for i in 1:last(size(A))\n sendcounts[A[ncode + 2, i] + 1] += m\n end\n\n # communicate columns of A to original rank\n recvcounts = similar(sendcounts)\n MPI.Alltoall!(UBuffer(sendcounts, 1), UBuffer(recvcounts, 1), comm)\n B = similar(A, sum(recvcounts))\n MPI.Alltoallv!(VBuffer(A, sendcounts), VBuffer(B, recvcounts), comm)\n B = reshape(B, m, div(length(B), m))\n\n # check to make sure we didn't drop any elements\n @assert nelem == n == size(B)[2]\n\n partsendcounts = zeros(Cint, csize)\n for i in 1:last(size(B))\n partsendcounts[B[ncode + 3, i] + 1] += 1\n end\n partsendstarts = ones(Int, csize + 1)\n for i in 1:csize\n partsendstarts[i + 1] = partsendcounts[i] + partsendstarts[i]\n end\n\n partsendorder = Int.(B[ncode + 1, :])\n\n partrecvcounts = similar(partsendcounts)\n MPI.Alltoall!(UBuffer(partsendcounts, 1), UBuffer(partrecvcounts, 1), comm)\n\n partrecvstarts = ones(Int, csize + 1)\n for i in 1:csize\n partrecvstarts[i + 1] = partrecvcounts[i] + partrecvstarts[i]\n end\n\n partsendorder, partsendstarts, partrecvstarts\nend\n\n\"\"\"\n partition(comm::MPI.Comm, elemtovert, elemtocoord, elemtobndy,\n faceconnections)\n\nThis function takes in a mesh (as returned for example by `brickmesh`) and\nreturns a Hilbert curve based partitioned mesh.\n\"\"\"\nfunction partition(\n comm::MPI.Comm,\n elemtovert,\n elemtocoord,\n elemtobndy,\n faceconnections,\n globord = [],\n)\n (d, nvert, nelem) = size(elemtocoord)\n\n csize = MPI.Comm_size(comm)\n crank = MPI.Comm_rank(comm)\n\n nface = 2d\n nfacevert = 2^(d - 1)\n\n # Here we expand the list of face connections into a structure that is easy\n # to partition. The cost is extra memory transfer. If this becomes a\n # bottleneck something more efficient may be implemented.\n #\n elemtofaceconnect =\n zeros(eltype(eltype(faceconnections)), nfacevert, nface, nelem)\n for fc in faceconnections\n elemtofaceconnect[:, fc[2], fc[1]] = fc[3:end]\n end\n\n elemtocode = centroidtocode(comm, elemtocoord; CT = UInt64)\n sendorder, sendstarts, recvstarts = getpartition(comm, elemtocode)\n\n elemtovert = elemtovert[:, sendorder]\n elemtocoord = elemtocoord[:, :, sendorder]\n elemtobndy = elemtobndy[:, sendorder]\n elemtofaceconnect = elemtofaceconnect[:, :, sendorder]\n\n if !isempty(globord)\n globord = globord[sendorder]\n end\n\n\n sendcounts = diff(sendstarts)\n recvcounts = similar(sendcounts)\n MPI.Alltoall!(UBuffer(sendcounts, 1), UBuffer(recvcounts, 1), comm)\n\n newelemtovert = similar(elemtovert, nvert * sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(elemtovert, Cint(nvert) .* sendcounts),\n VBuffer(newelemtovert, Cint(nvert) .* recvcounts),\n comm,\n )\n\n newelemtocoord = similar(elemtocoord, d * nvert * sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(elemtocoord, Cint(d * nvert) .* sendcounts),\n VBuffer(newelemtocoord, Cint(d * nvert) .* recvcounts),\n comm,\n )\n\n newelemtobndy = similar(elemtobndy, nface * sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(elemtobndy, Cint(nface) .* sendcounts),\n VBuffer(newelemtobndy, Cint(nface) .* recvcounts),\n comm,\n )\n\n newelemtofaceconnect =\n similar(elemtofaceconnect, nfacevert * nface * sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(elemtofaceconnect, Cint(nfacevert * nface) .* sendcounts),\n VBuffer(newelemtofaceconnect, Cint(nfacevert * nface) .* recvcounts),\n comm,\n )\n\n if !isempty(globord)\n newglobord = similar(globord, sum(recvcounts))\n MPI.Alltoallv!(\n VBuffer(globord, sendcounts),\n VBuffer(newglobord, recvcounts),\n comm,\n )\n else\n newglobord = similar(globord)\n end\n\n newnelem = recvstarts[end] - 1\n newelemtovert = reshape(newelemtovert, nvert, newnelem)\n newelemtocoord = reshape(newelemtocoord, d, nvert, newnelem)\n newelemtobndy = reshape(newelemtobndy, nface, newnelem)\n newelemtofaceconnect =\n reshape(newelemtofaceconnect, nfacevert, nface, newnelem)\n\n # reorder local elements based on code of new elements\n A = UInt64[\n centroidtocode(comm, newelemtocoord; CT = UInt64)\n collect(1:newnelem)'\n ]\n A = sortslices(A, dims = 2)\n newsortorder = view(A, d + 1, :)\n\n newelemtovert = newelemtovert[:, newsortorder]\n newelemtocoord = newelemtocoord[:, :, newsortorder]\n newelemtobndy = newelemtobndy[:, newsortorder]\n newelemtofaceconnect = newelemtofaceconnect[:, :, newsortorder]\n\n newfaceconnections = similar(faceconnections, 0)\n for e in 1:newnelem, f in 1:nface\n if newelemtofaceconnect[1, f, e] > 0\n push!(newfaceconnections, vcat(e, f, newelemtofaceconnect[:, f, e]))\n end\n end\n\n if !isempty(globord)\n newglobord = newglobord[newsortorder]\n end\n\n (\n newelemtovert,\n newelemtocoord,\n newelemtobndy,\n newfaceconnections,\n newglobord,\n )#sendorder)\nend\n\n\"\"\"\n minmaxflip(x, y)\n\nReturns `x, y` sorted lowest to highest and a bool that indicates if a swap\nwas needed.\n\"\"\"\nminmaxflip(x, y) = y < x ? (y, x, true) : (x, y, false)\n\n\"\"\"\n vertsortandorder(a)\n\nReturns `(a)` and an ordering `o==0`.\n\"\"\"\nvertsortandorder(a) = ((a,), 1)\n\n\"\"\"\n vertsortandorder(a, b)\n\nReturns sorted vertex numbers `(a,b)` and an ordering `o` depending on the\norder needed to sort the elements. This ordering is given below including the\nvetex ordering for faces.\n\n o= 0 1\n\n (a,b) (b,a)\n\n a b\n | |\n | |\n b a\n\"\"\"\nfunction vertsortandorder(a, b)\n a, b, s1 = minmaxflip(a, b)\n o = s1 ? 2 : 1\n ((a, b), o)\nend\n\n\"\"\"\n vertsortandorder(a, b, c)\n\nReturns sorted vertex numbers `(a,b,c)` and an ordering `o` depending on the\norder needed to sort the elements. This ordering is given below including the\nvetex ordering for faces.\n\n o= 1 2 3 4 5 6\n\n (a,b,c) (c,a,b) (b,c,a) (b,a,c) (c,b,a) (a,c,b)\n\n /c\\\\ /b\\\\ /a\\\\ /c\\\\ /a\\\\ /b\\\\\n / \\\\ / \\\\ / \\\\ / \\\\ / \\\\ / \\\\\n /a___b\\\\ /c___a\\\\ /b___c\\\\ /b___a\\\\ /c___b\\\\ /a___c\\\\\n\"\"\"\nfunction vertsortandorder(a, b, c)\n # Use a (Bose-Nelson Algorithm based) sorting network from\n # to sort the vertices.\n b, c, s1 = minmaxflip(b, c)\n a, c, s2 = minmaxflip(a, c)\n a, b, s3 = minmaxflip(a, b)\n\n if !s1 && !s2 && !s3\n o = 1\n elseif !s1 && s2 && s3\n o = 2\n elseif s1 && !s2 && s3\n o = 3\n elseif !s1 && !s2 && s3\n o = 4\n elseif s1 && s2 && s3\n o = 5\n elseif s1 && !s2 && !s3\n o = 6\n else\n error(\"Problem finding vertex ordering $((a,b,c)) with flips\n $((s1,s2,s3))\")\n end\n\n ((a, b, c), o)\nend\n\n\"\"\"\n vertsortandorder(a, b, c, d)\n\nReturns sorted vertex numbers `(a,b,c,d)` and an ordering `o` depending on the\norder needed to sort the elements. This ordering is given below including the\nvetex ordering for faces.\n\n o= 1 2 3 4 5 6 7 8\n\n (a,b, (a,c, (b,a, (b,d, (c,a, (c,d, (d,b, (d,c,\n c,d) b,d) c,d) a,c) d,b) a,b) c,a) b,a)\n\n c---d b---d c---d a---c d---b a---b c---a b---a\n | | | | | | | | | | | | | | | |\n a---b a---c b---a b---d c---a c---d d---b d---c\n\"\"\"\nfunction vertsortandorder(a, b, c, d)\n # Use a (Bose-Nelson Algorithm based) sorting network from\n # to sort the vertices.\n a, b, s1 = minmaxflip(a, b)\n c, d, s2 = minmaxflip(c, d)\n a, c, s3 = minmaxflip(a, c)\n b, d, s4 = minmaxflip(b, d)\n b, c, s5 = minmaxflip(b, c)\n\n if !s1 && !s2 && !s3 && !s4 && !s5\n o = 1\n elseif !s1 && !s2 && !s3 && !s4 && s5\n o = 2\n elseif s1 && !s2 && !s3 && !s4 && !s5\n o = 3\n elseif !s1 && !s2 && s3 && s4 && s5\n o = 4\n elseif s1 && s2 && !s3 && !s4 && s5\n o = 5\n elseif !s1 && !s2 && s3 && s4 && !s5\n o = 6\n elseif s1 && s2 && s3 && s4 && s5\n o = 7\n elseif s1 && s2 && s3 && s4 && !s5\n o = 8\n else\n # FIXME: some possible orientations are missing since there are a total of\n # 24. Missing orientations:\n #=\n d---c d---c b---c a---d c---b d---a a---b b---a\n | | | | | | | | | | | | | | | |\n a---b b---a a---d b---c d---a c---b d---c c---d\n\n c---b d---b b---d b---c c---a d---a a---d a---c\n | | | | | | | | | | | | | | | |\n a---d a---c c---a d---a b---d b---c c---b d---b\n =#\n error(\"Problem finding vertex ordering $((a,b,c,d))\n with flips $((s1,s2,s3,s4,s5))\")\n end\n\n ((a, b, c, d), o)\nend\n\n\"\"\"\n connectmesh(comm::MPI.Comm, elemtovert, elemtocoord, elemtobndy,\n faceconnections)\n\nThis function takes in a mesh (as returned for example by `brickmesh`) and\nreturns a connected mesh. This returns a `NamedTuple` of:\n\n - `elems` the range of element indices\n - `realelems` the range of real (aka nonghost) element indices\n - `ghostelems` the range of ghost element indices\n - `ghostfaces` ghost element to face is received;\n `ghostfaces[f,ge] == true` if face `f` of ghost element `ge` is received.\n - `sendelems` an array of send element indices\n - `sendfaces` send element to face is sent;\n `sendfaces[f,se] == true` if face `f` of send element `se` is sent.\n - `elemtocoord` element to vertex coordinates; `elemtocoord[d,i,e]` is the\n `d`th coordinate of corner `i` of element `e`\n - `elemtoelem` element to neighboring element; `elemtoelem[f,e]` is the\n number of the element neighboring element `e` across face `f`. If there is\n no neighboring element then `elemtoelem[f,e] == e`.\n - `elemtoface` element to neighboring element face; `elemtoface[f,e]` is the\n face number of the element neighboring element `e` across face `f`. If\n there is no neighboring element then `elemtoface[f,e] == f`.\n - `elemtoordr` element to neighboring element order; `elemtoordr[f,e]` is the\n ordering number of the element neighboring element `e` across face `f`. If\n there is no neighboring element then `elemtoordr[f,e] == 1`.\n - `elemtobndy` element to bounday number; `elemtobndy[f,e]` is the\n boundary number of face `f` of element `e`. If there is a neighboring\n element then `elemtobndy[f,e] == 0`.\n - `nabrtorank` a list of the MPI ranks for the neighboring processes\n - `nabrtorecv` a range in ghost elements to receive for each neighbor\n - `nabrtosend` a range in `sendelems` to send for each neighbor\n\n\"\"\"\nfunction connectmesh(\n comm::MPI.Comm,\n elemtovert,\n elemtocoord,\n elemtobndy,\n faceconnections;\n dim = size(elemtocoord, 1),\n)\n d = dim\n (coorddim, nvert, nelem) = size(elemtocoord)\n nface, nfacevert = 2d, 2^(d - 1)\n\n p = reshape(1:nvert, ntuple(j -> 2, d))\n fmask = hcat((\n p[ntuple(\n j ->\n (j == div(f - 1, 2) + 1) ?\n ((mod(f - 1, 2) + 1):(mod(f - 1, 2) + 1)) : (:),\n d,\n )...][:] for f in 1:nface\n )...)\n\n csize = MPI.Comm_size(comm)\n crank = MPI.Comm_rank(comm)\n\n VT = eltype(elemtovert)\n A = Array{VT}(undef, nfacevert + 8, nface * nelem)\n\n MR, ME, MF, MO, NR, NE, NF, NO = nfacevert .+ (1:8)\n for e in 1:nelem\n v = reshape(elemtovert[:, e], ntuple(j -> 2, d))\n for f in 1:nface\n j = (e - 1) * nface + f\n fv, o = vertsortandorder(v[fmask[:, f]]...)\n A[1:nfacevert, j] .= fv\n A[MR, j] = crank\n A[ME, j] = e\n A[MF, j] = f\n A[MO, j] = o\n A[NR, j] = typemax(VT)\n A[NE, j] = typemax(VT)\n A[NF, j] = typemax(VT)\n A[NO, j] = typemax(VT)\n end\n end\n\n # use neighboring vertices for connected faces\n for fc in faceconnections\n e = fc[1]\n f = fc[2]\n v = fc[3:end]\n j = (e - 1) * nface + f\n fv, o = vertsortandorder(v...)\n A[1:nfacevert, j] .= fv\n A[MO, j] = o\n end\n\n A = parallelsortcolumns(comm, A, by = x -> x[1:nfacevert])\n m, n = size(A)\n\n # match faces\n j = 1\n while j <= n\n if j + 1 <= n && A[1:nfacevert, j] == A[1:nfacevert, j + 1]\n # found connected face\n A[NR:NO, j] = A[MR:MO, j + 1]\n A[NR:NO, j + 1] = A[MR:MO, j]\n j += 2\n else\n # found unconnect face\n A[NR:NO, j] = A[MR:MO, j]\n j += 1\n end\n end\n\n A = sortslices(A, dims = 2, by = x -> (x[MR], x[NR], x[ME], x[MF]))\n\n # count number of elements that are going to be sent\n sendcounts = zeros(Cint, csize)\n for i in 1:last(size(A))\n sendcounts[A[MR, i] + 1] += m\n end\n sendstarts = ones(Int, csize + 1)\n for i in 1:csize\n sendstarts[i + 1] = sendcounts[i] + sendstarts[i]\n end\n\n # communicate columns of A to original rank\n recvcounts = similar(sendcounts)\n MPI.Alltoall!(UBuffer(sendcounts, 1), UBuffer(recvcounts, 1), comm)\n\n B = similar(A, sum(recvcounts))\n MPI.Alltoallv!(VBuffer(A, sendcounts), VBuffer(B, recvcounts), comm)\n B = reshape(B, m, nface * nelem)\n\n # get element sending information\n B = sortslices(B, dims = 2, by = x -> (x[NR], x[ME]))\n sendelems = Int[]\n counts = zeros(Int, csize + 1)\n counts[1] = (last(size(B)) > 0) ? 1 : 0\n sr, se = -1, 0\n for i in 1:last(size(B))\n r, e = B[NR, i], B[ME, i]\n # See if we need to send element `e` to rank `r` and make sure that we\n # didn't already mark it for sending.\n if r != crank && !(sr == r && se == e)\n counts[r + 2] += 1\n append!(sendelems, e)\n sr, se = r, e\n end\n end\n\n # Mark which faces need to be sent\n sendfaces = BitArray(undef, nface, length(sendelems))\n sendfaces .= false\n sr, se, n = -1, 0, 0\n for i in 1:last(size(B))\n r, e, f = B[NR, i], B[ME, i], B[MF, i]\n if r != crank\n if !(sr == r && se == e)\n n += 1\n sr, se = r, e\n end\n sendfaces[f, n] = true\n end\n end\n\n sendstarts = cumsum(counts)\n nabrtosendrank = Int[\n r for r = 0:(csize - 1) if sendstarts[r + 2] - sendstarts[r + 1] > 0\n ]\n nabrtosend = UnitRange{Int}[\n (sendstarts[r + 1]:(sendstarts[r + 2] - 1))\n for r = 0:(csize - 1) if sendstarts[r + 2] - sendstarts[r + 1] > 0\n ]\n\n # get element receiving information\n B = sortslices(B, dims = 2, by = x -> (x[NR], x[NE]))\n counts = zeros(Int, csize + 1)\n counts[1] = (last(size(B)) > 0) ? 1 : 0\n sr, se = -1, 0\n nghost = 0\n for i in 1:last(size(B))\n r, e = B[NR, i], B[NE, i]\n if r != crank\n # Check to make sure we have not already marked the element for\n # receiving since we could be connected to the receiving element across\n # multiple faces.\n if !(sr == r && se == e)\n nghost += 1\n counts[r + 2] += 1\n sr, se = r, e\n end\n B[NE, i] = nelem + nghost\n end\n end\n\n # Mark which faces will be received\n ghostfaces = BitArray(undef, nface, nghost)\n ghostfaces .= false\n sr, se, ge = -1, 0, 0\n for i in 1:last(size(B))\n r, e, f = B[NR, i], B[NE, i], B[NF, i]\n if r != crank\n if !(sr == r && se == e)\n ge += 1\n sr, se = r, e\n end\n B[NR, i] = crank\n ghostfaces[f, ge] = true\n end\n end\n\n recvstarts = cumsum(counts)\n nabrtorecvrank = Int[\n r for r = 0:(csize - 1) if recvstarts[r + 2] - recvstarts[r + 1] > 0\n ]\n nabrtorecv = UnitRange{Int}[\n (recvstarts[r + 1]:(recvstarts[r + 2] - 1))\n for r = 0:(csize - 1) if recvstarts[r + 2] - recvstarts[r + 1] > 0\n ]\n\n @assert nabrtorecvrank == nabrtosendrank\n nabrtorank = nabrtorecvrank\n\n elemtoelem = repeat((1:(nelem + nghost))', nface, 1)\n elemtoface = repeat(1:nface, 1, nelem + nghost)\n elemtoordr = ones(Int, nface, nelem + nghost)\n\n if d == 2\n for i in 1:last(size(B))\n me, mf, mo = B[ME, i], B[MF, i], B[MO, i]\n ne, nf, no = B[NE, i], B[NF, i], B[NO, i]\n\n elemtoelem[mf, me] = ne\n elemtoface[mf, me] = nf\n elemtoordr[mf, me] = (no == mo ? 1 : 2)\n end\n else\n for i in 1:last(size(B))\n me, mf, mo = B[ME, i], B[MF, i], B[MO, i]\n ne, nf, no = B[NE, i], B[NF, i], B[NO, i]\n\n elemtoelem[mf, me] = ne\n elemtoface[mf, me] = nf\n if no != 1 || mo != 1\n error(\"TODO add support for other orientations\")\n end\n elemtoordr[mf, me] = 1\n end\n end\n\n # fill the ghost values in elemtocoord\n newelemtocoord = similar(elemtocoord, coorddim, nvert, nelem + nghost)\n newelemtobndy = similar(elemtobndy, nface, nelem + nghost)\n\n sendelemtocoord = elemtocoord[:, :, sendelems]\n sendelemtobndy = elemtobndy[:, sendelems]\n\n crreq = [\n MPI.Irecv!(view(newelemtocoord, :, :, nelem .+ er), r, 666, comm)\n for (r, er) in zip(nabrtorank, nabrtorecv)\n ]\n\n brreq = [\n MPI.Irecv!(view(newelemtobndy, :, nelem .+ er), r, 666, comm)\n for (r, er) in zip(nabrtorank, nabrtorecv)\n ]\n\n csreq = [\n MPI.Isend(view(sendelemtocoord, :, :, es), r, 666, comm)\n for (r, es) in zip(nabrtorank, nabrtosend)\n ]\n\n bsreq = [\n MPI.Isend(view(sendelemtobndy, :, es), r, 666, comm)\n for (r, es) in zip(nabrtorank, nabrtosend)\n ]\n\n newelemtocoord[:, :, 1:nelem] .= elemtocoord\n newelemtobndy[:, 1:nelem] .= elemtobndy\n\n MPI.Waitall!([csreq; crreq; bsreq; brreq])\n\n (\n elems = 1:(nelem + nghost), # range of element indices\n realelems = 1:nelem, # range of real element indices\n ghostelems = nelem .+ (1:nghost), # range of ghost element indices\n ghostfaces = ghostfaces,\n sendelems = sendelems, # array of send element indices\n sendfaces = sendfaces,\n elemtocoord = newelemtocoord, # element to vertex coordinates\n elemtovert = nothing, # element to locally unique global vertex number (for direct stiffness summation)\n elemtoelem = elemtoelem, # element to neighboring element\n elemtoface = elemtoface, # element to neighboring element face\n elemtoordr = elemtoordr, # element to neighboring element order\n elemtobndy = newelemtobndy, # element to boundary number\n nabrtorank = nabrtorank, # list of neighboring processes MPI ranks\n nabrtorecv = nabrtorecv, # neighbor receive ranges into `ghostelems`\n nabrtosend = nabrtosend,\n ) # neighbor send ranges into `sendelems`\nend\n\n\"\"\"\n (bndytoelem, bndytoface) = enumerateboundaryfaces!(elemtoelem, elemtobndy, periodicity, boundary)\n\nUpdate the `elemtoelem` array based on the boundary faces specified in\n`elemtobndy`. Builds the `bndytoelem` and `bndytoface` tuples.\n\"\"\"\nfunction enumerateboundaryfaces!(elemtoelem, elemtobndy, periodicity, boundary)\n nb = 0\n for i in 1:length(periodicity)\n if !periodicity[i]\n nb = max(nb, boundary[i]...)\n end\n end\n # Limitation of the boundary condition unrolling in the DG kernels\n # (should never be violated unless more general unstructured meshes are used\n # since cube meshes only have 6 faces in 3D, and only 1 bcs is currently allowed\n # per face)\n\n @assert nb <= 6\n\n bndytoelem = ntuple(b -> Vector{Int64}(), nb)\n bndytoface = ntuple(b -> Vector{Int64}(), nb)\n\n nface, nelem = size(elemtoelem)\n\n N = zeros(Int, nb)\n for e in 1:nelem\n for f in 1:nface\n d = elemtobndy[f, e]\n @assert 0 <= d <= nb\n if d != 0\n elemtoelem[f, e] = N[d] += 1\n push!(bndytoelem[d], e)\n push!(bndytoface[d], f)\n end\n end\n end\n return (bndytoelem, bndytoface)\nend\n\n\"\"\"\n connectmeshfull(comm::MPI.Comm, elemtovert, elemtocoord, elemtobndy,\n faceconnections)\n\nThis function takes in a mesh (as returned for example by `brickmesh`) and\nreturns a corner connected mesh. It is similar to [`connectmesh`](@ref) but returns \na corner connected mesh, i.e. `ghostelems` will also include remote elements \nthat are only connected by vertices. This returns a `NamedTuple` of:\n\n - `elems` the range of element indices\n - `realelems` the range of real (aka nonghost) element indices\n - `ghostelems` the range of ghost element indices\n - `ghostfaces` ghost element to face is received;\n `ghostfaces[f,ge] == true` if face `f` of ghost element `ge` is received.\n - `sendelems` an array of send element indices\n - `sendfaces` send element to face is sent;\n `sendfaces[f,se] == true` if face `f` of send element `se` is sent.\n - `elemtocoord` element to vertex coordinates; `elemtocoord[d,i,e]` is the\n `d`th coordinate of corner `i` of element `e`\n - `elemtoelem` element to neighboring element; `elemtoelem[f,e]` is the\n number of the element neighboring element `e` across face `f`. If there is\n no neighboring element then `elemtoelem[f,e] == e`.\n - `elemtoface` element to neighboring element face; `elemtoface[f,e]` is the\n face number of the element neighboring element `e` across face `f`. If\n there is no neighboring element then `elemtoface[f,e] == f`.\n - `elemtoordr` element to neighboring element order; `elemtoordr[f,e]` is the\n ordering number of the element neighboring element `e` across face `f`. If\n there is no neighboring element then `elemtoordr[f,e] == 1`.\n - `elemtobndy` element to bounday number; `elemtobndy[f,e]` is the\n boundary number of face `f` of element `e`. If there is a neighboring\n element then `elemtobndy[f,e] == 0`.\n - `nabrtorank` a list of the MPI ranks for the neighboring processes\n - `nabrtorecv` a range in ghost elements to receive for each neighbor\n - `nabrtosend` a range in `sendelems` to send for each neighbor\n\nThe algorithm assumes that the 2-D mesh can be gathered on single process,\nwhich is reasonable for the 2-D mesh.\n\"\"\"\nfunction connectmeshfull(\n comm::MPI.Comm,\n elemtovert,\n elemtocoord,\n elemtobndy,\n faceconnections;\n dim = size(elemtocoord, 1),\n)\n @assert dim == 2\n dim_coord = size(elemtocoord, 1)\n csize = MPI.Comm_size(comm)\n crank = MPI.Comm_rank(comm)\n root = 0\n I = eltype(elemtovert)\n FT = eltype(elemtocoord)\n nvert, nelem = size(elemtovert)\n nfaces = 2 * dim\n nelemv = zeros(I, csize)\n fmask = build_fmask(dim)\n nfvert = size(fmask, 1)\n # collecting # of realelems on each process\n MPI.Allgather!(nelem, UBuffer(nelemv, Cint(1)), comm)\n offset = cumsum(nelemv) .+ 1\n pushfirst!(offset, 1)\n nelemg = sum(nelemv)\n # collecting elemtovert on each process\n elemtovertg = zeros(I, nvert + 2, nelemg)\n MPI.Allgatherv!(\n [elemtovert; ones(I, 1, nelem) .* crank; reshape(Array(1:nelem), 1, :)],\n VBuffer(elemtovertg, nelemv .* (nvert + 2)),\n comm,\n )\n\n nvertg = maximum(elemtovertg[1:nvert, :])\n # collecting elemtocoordg on each process\n elemtocoordg = Array{FT}(undef, dim_coord, nvert, nelemg)\n MPI.Allgatherv!(\n elemtocoord,\n VBuffer(elemtocoordg, nelemv .* (dim_coord * nvert)),\n comm,\n )\n # collecting elemtobndy on each process\n elemtobndyg = zeros(I, nfaces, nelemg)\n MPI.Allgatherv!(elemtobndy, VBuffer(elemtobndyg, nelemv .* nfaces), comm)\n # accounting for vertices on periodic faces\n nper = length(faceconnections) * nfvert\n vconn = Array{Int}(undef, 2, nper) # stores the periodic and corresponding shadow vertex\n ctr = 1\n for fc in faceconnections\n e, f, v = fc[1], fc[2], fc[3:end]\n fv = elemtovert[fmask[:, f], e]\n fv, o = vertsortandorder(fv)\n v, o = vertsortandorder(v)\n for i in 1:nfvert\n vconn[1, ctr] = fv[1][i]\n vconn[2, ctr] = v[1][i]\n ctr += 1\n end\n end\n vconn = unique(vconn, dims = 2)\n nper = size(vconn, 2)\n nperv = zeros(I, csize)\n MPI.Allgather!(nper, UBuffer(nperv, Cint(1)), comm)\n nperg = sum(nperv)\n vconng = Array{Int}(undef, 2, nperg)\n MPI.Allgatherv!(vconn, VBuffer(vconng, nperv .* 2), comm)\n\n gldofv = -ones(Int, nvertg)\n pmarker = -ones(Int, nperg)\n # conflict-free periodic nodes\n for i in 1:nperg\n v1, v2 = vconng[1, i], vconng[2, i]\n if gldofv[v1] == -1 && gldofv[v2] == -1\n id = min(v1, v2)\n gldofv[v1] = id\n gldofv[v2] = id\n pmarker[i] = 1\n end\n end\n # dealing with doubly periodic nodes (whenever applicable)\n for i in 1:nperg\n if pmarker[i] == -1\n v1, v2 = vconng[1, i], vconng[2, i]\n id = min(gldofv[v1], gldofv[v2])\n gldofv[v1] = id\n gldofv[v2] = id\n end\n end\n # labeling non-periodic vertices\n for i in 1:nvertg\n if gldofv[i] == -1\n gldofv[i] = i\n end\n end\n #-------------------------------------------------------------\n elemtovertg_orig = deepcopy(elemtovertg)\n for i in 1:nelemg, j in 1:nvert\n elemtovertg[j, i] = gldofv[elemtovertg[j, i]]\n end\n #-------------------------------------------------\n nsend, nghost = 0, 0\n interiorelems, exteriorelems = Int[], Int[]\n vertgtoprocs = map(i -> zeros(Int, i), zeros(Int, nvertg)) # procs for each vertex\n vertgtolelem = map(i -> zeros(Int, i), zeros(Int, nvertg)) # local cell # for each vertex\n vertgtolvert = map(i -> zeros(Int, i), zeros(Int, nvertg)) # local vertex # in local cell for each vertex\n sendelems = map(i -> zeros(Int, i), zeros(Int, csize))\n recvelems = map(i -> zeros(Int, i), zeros(Int, csize))\n\n for icls in 1:nelemg # gathering connectivity info for each global vertex\n for ivt in 1:nvert\n gvt = elemtovertg[ivt, icls]\n prc = elemtovertg[nvert + 1, icls]\n lcl = elemtovertg[nvert + 2, icls]\n push!(vertgtoprocs[gvt], prc)\n push!(vertgtolelem[gvt], lcl)\n push!(vertgtolvert[gvt], ivt)\n end\n end\n\n for icls in 1:nelem # building sendelems and recvelems\n interior = true\n for ivt in 1:nvert\n vt = elemtovert[ivt, icls]\n gvt = gldofv[vt]\n for ip in 1:length(vertgtoprocs[gvt])\n proc = vertgtoprocs[gvt][ip]\n if proc ≠ crank\n lcell = vertgtolelem[gvt][ip]\n if findfirst(recvelems[proc + 1] .== lcell) == nothing\n push!(recvelems[proc + 1], lcell)\n nghost += 1\n end\n if findfirst(sendelems[proc + 1] .== icls) == nothing\n push!(sendelems[proc + 1], icls)\n nsend += 1\n end\n interior = false\n end\n end\n end\n interior ? push!(interiorelems, icls) : push!(exteriorelems, icls)\n end\n\n nabrtorank = Int[]\n newsendelems = Array{Int}(undef, nsend)\n nabrtosend = Array{UnitRange{Int64}}(undef, 0)\n nabrtorecv = Array{UnitRange{Int64}}(undef, 0)\n st_recv, en_recv = 1, 1\n st_send, en_send = 1, 1\n\n for ipr in 0:(csize - 1)\n l_send = length(sendelems[ipr + 1])\n l_recv = length(recvelems[ipr + 1])\n if l_send > 0\n en_send = st_send + l_send - 1\n sort!(sendelems[ipr + 1])\n newsendelems[st_send:en_send] .= sendelems[ipr + 1][:]\n push!(nabrtosend, st_send:en_send)\n st_send = en_send + 1\n push!(nabrtorank, ipr)\n\n end\n if l_recv > 0\n en_recv = st_recv + l_recv - 1\n sort!(recvelems[ipr + 1])\n push!(nabrtorecv, st_recv:en_recv)\n st_recv = en_recv + 1\n end\n end\n\n sendfaces = BitArray(zeros(nfaces, nsend))\n ghostfaces = BitArray(zeros(nfaces, nghost))\n\n newelemtovert = similar(elemtovert, nvert, (nelem + nghost))\n newelemtocoord = similar(elemtocoord, dim_coord, nvert, (nelem + nghost))\n newelemtobndy = similar(elemtobndy, nfaces, (nelem + nghost))\n newelemtovert[1:nvert, 1:nelem] .= elemtovert\n newelemtocoord[:, :, 1:nelem] .= elemtocoord\n newelemtobndy[:, 1:nelem] .= elemtobndy\n vmarker = BitArray(undef, nvert)\n ctrg, ctrs = 1, 1\n\n for ipr in 0:(csize - 1)\n # building ghost faces\n for icls in recvelems[ipr + 1]\n vmarker .= 0\n off = offset[ipr + 1]\n newelemtovert[1:nvert, nelem + ctrg] .=\n elemtovertg_orig[1:nvert, off + icls - 1]\n newelemtocoord[:, :, nelem + ctrg] .=\n elemtocoordg[:, :, off + icls - 1]\n newelemtobndy[:, nelem + ctrg] .= elemtobndyg[:, off + icls - 1]\n for ivt in 1:nvert\n gvt = elemtovertg[ivt, icls + off - 1]\n if findfirst(vertgtoprocs[gvt] .== crank) ≠ nothing\n vmarker[ivt] = 1\n end\n end\n for fc in 1:nfaces\n if findfirst(vmarker[fmask[:, fc]]) ≠ nothing\n ghostfaces[fc, ctrg] = 1\n end\n end\n ctrg += 1\n end\n # building send faces\n for icls in sendelems[ipr + 1]\n vmarker .= 0\n for ivt in 1:nvert\n vt = elemtovert[ivt, icls]\n gvt = gldofv[vt]\n if findfirst(vertgtoprocs[gvt] .== ipr) ≠ nothing\n vmarker[ivt] = 1\n end\n end\n for fc in 1:nfaces\n if findfirst(vmarker[fmask[:, fc]]) ≠ nothing\n sendfaces[fc, ctrs] = 1\n end\n end\n ctrs += 1\n end\n end\n\n A = zeros(Int, (nfvert + 3), nfaces * (nelem + nghost))\n for e in 1:(nelem + nghost)\n v = reshape(newelemtovert[:, e], ntuple(j -> 2, dim))\n for f in 1:nfaces\n j = (e - 1) * nfaces + f\n fv, o = vertsortandorder(v[fmask[:, f]]...) # faces vertices, B-N orientation\n for fvt in 1:nfvert\n A[fvt, j] = gldofv[fv[fvt]]\n end\n A[nfvert + 1, j] = o # orientation\n A[nfvert + 2, j] = e # local element number\n A[nfvert + 3, j] = f # local face number for element e\n end\n end\n\n A = sortslices(A, dims = 2, by = x -> x[1:nfvert])\n elemtoelem = Array{Int}(undef, nfaces, (nelem + nghost))\n elemtoface = Array{Int}(undef, nfaces, (nelem + nghost))\n elemtoordr = Array{Int}(undef, nfaces, (nelem + nghost))\n # match faces\n n = size(A, 2)\n j = 1\n while j ≤ n\n lel = A[nfvert + 2, j]\n lfc = A[nfvert + 3, j]\n if j + 1 ≤ n && A[1:nfvert, j] == A[1:nfvert, j + 1]\n # found connected face\n nel = A[nfvert + 2, j + 1] # local neighboring element\n nfc = A[nfvert + 3, j + 1] # local face # of local neighboring element\n elemtoelem[lfc, lel] = nel\n elemtoface[lfc, lel] = nfc\n elemtoelem[nfc, nel] = lel\n elemtoface[nfc, nel] = lfc\n if A[nfvert + 1, j] == A[nfvert + 1, j + 1]\n elemtoordr[lfc, lel] = 1\n elemtoordr[nfc, nel] = 1\n else\n elemtoordr[lfc, lel] = 2\n elemtoordr[nfc, nel] = 2\n end\n j += 2\n else\n # found unconnected face\n elemtoelem[lfc, lel] = lel\n elemtoface[lfc, lel] = lfc\n elemtoordr[lfc, lel] = 1\n j += 1\n end\n end\n # provide locally unique elemtovert for DSS\n uvert = -ones(Int, nvertg)\n for el in 1:last(size(newelemtovert)), i in 1:nvert\n newelemtovert[i, el] = gldofv[newelemtovert[i, el]]\n uvert[newelemtovert[i, el]] = 0\n end\n ctr = 1\n for i in 1:length(uvert)\n if uvert[i] == 0\n uvert[i] = ctr\n ctr += 1\n end\n end\n for el in 1:last(size(newelemtovert)), i in 1:nvert\n newelemtovert[i, el] = uvert[newelemtovert[i, el]]\n end\n #---------------------------------------------------------------\n (\n elems = 1:(nelem + nghost), # range of element indices\n realelems = 1:nelem, # range of real element indices\n ghostelems = nelem .+ (1:nghost), # range of ghost element indices\n ghostfaces = ghostfaces, # bit array of recv faces for ghost elems\n sendelems = newsendelems, # array of send element indices\n sendfaces = sendfaces, # bit array of send faces for send elems\n elemtocoord = newelemtocoord, # element to vertex coordinates\n elemtovert = newelemtovert, # element to locally unique global vertex number (for direct stiffness summation)\n elemtoelem = elemtoelem, # element to neighboring element\n elemtoface = elemtoface, # element to neighboring element face\n elemtoordr = elemtoordr, # element to neighboring element order\n elemtobndy = newelemtobndy, # element to boundary number\n nabrtorank = nabrtorank, # list of neighboring processes MPI ranks\n nabrtorecv = nabrtorecv, # neighbor receive ranges into `ghostelems`\n nabrtosend = nabrtosend, # neighbor send ranges into `sendelems`\n )\nend\n\n\"\"\"\n build_fmask(dim)\n\nReturns the face mask for mapping element vertices to face vertices.\n\nex: for 2D element with vertices (1, 2, 3, 4)\n\n 3---4 \n | | \n 1---2 \n\nthe function returns the face mask\nf1 | f2 | f3 | f4 \n=================\n 1 | 2 | 1 | 3\n 3 | 4 | 2 | 4\n================= \n\"\"\"\nfunction build_fmask(dim)\n nvert = 2^dim\n nfaces = 2 * dim\n p = reshape(1:nvert, ntuple(j -> 2, dim))\n fmask = Array{Int64}(undef, 2^(dim - 1), nfaces)\n f = 0\n for d in 1:dim\n for slice in eachslice(p, dims = d)\n fmask[:, f += 1] = vec(slice)\n end\n end\n return fmask\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/DSS.jl", "content": "module DSS\n\nusing ..Grids\nusing ClimateMachine.MPIStateArrays\nusing CUDA\nusing KernelAbstractions\nusing DocStringExtensions\n\nexport dss!\n\n\"\"\"\n dss3d(Q::MPIStateArray,\n grid::DiscontinuousSpectralElementGrid)\n\nThis function computes the 3D direct stiffness summation for all variables in the MPIStateArray.\n\n# Fields\n - `Q`: MPIStateArray\n - `grid`: Discontinuous Spectral Element Grid\n\"\"\"\nfunction dss!(\n Q::MPIStateArray,\n grid::DiscontinuousSpectralElementGrid{FT, 3};\n max_threads = 256,\n) where {FT}\n DA = arraytype(grid) # device array\n device = arraytype(grid) <: Array ? CPU() : CUDADevice() # device\n #----communication--------------------------\n event = MPIStateArrays.begin_ghost_exchange!(Q)\n event = MPIStateArrays.end_ghost_exchange!(Q, dependencies = event)\n wait(event)\n #----Direct Stiffness Summation-------------\n vertmap = grid.vertmap\n edgemap = grid.edgemap\n facemap = grid.facemap\n\n vtconn = grid.topology.vtconn\n fcconn = grid.topology.fcconn\n edgconn = grid.topology.edgconn\n\n vtconnoff = grid.topology.vtconnoff\n edgconnoff = grid.topology.edgconnoff\n\n nvt, nfc, nedg =\n length(vtconnoff) - 1, size(fcconn, 1), length(edgconnoff) - 1\n\n Nq = polynomialorders(grid) .+ 1\n Nqmax = maximum(Nq)\n Nemax = Nqmax - 2\n Nfmax = size(facemap, 1)\n args = (\n Q.data,\n vertmap,\n edgemap,\n facemap,\n vtconn,\n vtconnoff,\n edgconn,\n edgconnoff,\n fcconn,\n nvt,\n nedg,\n nfc,\n Nemax,\n Nfmax,\n )\n if device == CPU()\n dss3d_CPU!(args...)\n else\n n_items = nvt + nfc + nedg\n tx = max(n_items, max_threads)\n bx = cld(n_items, tx)\n @cuda threads = (tx) blocks = (bx) dss3d_CUDA!(args...)\n end\n return nothing\nend\n\nfunction dss3d_CUDA!(\n data,\n vertmap,\n edgemap,\n facemap,\n vtconn,\n vtconnoff,\n edgconn,\n edgconnoff,\n fcconn,\n nvt,\n nedg,\n nfc,\n Nemax,\n Nfmax,\n)\n I = eltype(nvt)\n FT = eltype(data)\n\n tx = threadIdx().x # threaid id\n bx = blockIdx().x # block id\n bxdim = blockDim().x # block dimension\n glx = tx + (bx - 1) * bxdim # global id\n nvars = size(data, 2)\n\n # A mesh node is either\n # - interior (no dss required)\n # - on an element corner / vertex,\n # - edge (excluding end point element corners / vertices)\n # - face (excluding edges and corners)\n if glx ≤ nvt #vertex DSS\n vx = glx\n for ivar in 1:nvars\n dss_vertex!(vtconn, vtconnoff, vertmap, data, ivar, vx, FT)\n end\n elseif glx > nvt && glx ≤ (nvt + nedg) # edge DSS\n ex = glx - nvt\n for ivar in 1:nvars\n dss_edge!(edgconn, edgconnoff, edgemap, Nemax, data, ivar, ex, FT)\n end\n elseif glx > (nvt + nedg) && glx ≤ (nvt + nedg + nfc) # face DSS\n fx = glx - (nvt + nedg)\n for ivar in 1:nvars\n dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)\n end\n end\n\n return nothing\nend\n\nfunction dss3d_CPU!(\n data,\n vertmap,\n edgemap,\n facemap,\n vtconn,\n vtconnoff,\n edgconn,\n edgconnoff,\n fcconn,\n nvt,\n nedg,\n nfc,\n Nemax,\n Nfmax,\n)\n I = eltype(nvt)\n FT = eltype(data)\n nvars = size(data, 2)\n\n # A mesh node is either\n # - interior (no dss required)\n # - on an element corner / vertex,\n # - edge (excluding end point element corners / vertices)\n # - face (excluding edges and corners)\n for ivar in 1:nvars\n for vx in 1:nvt # vertex DSS\n dss_vertex!(vtconn, vtconnoff, vertmap, data, ivar, vx, FT)\n end\n for ex in 1:nedg # edge DSS\n dss_edge!(edgconn, edgconnoff, edgemap, Nemax, data, ivar, ex, FT)\n end\n for fx in 1:nfc # face DSS\n dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)\n end\n end\n return nothing\nend\n\n\"\"\"\n dss_vertex!(\n vtconn,\n vtconnoff,\n vertmap,\n data,\n ivar,\n vx,\n ::Type{FT},\n) where {FT}\n\nThis function computes the direct stiffness summation for the vertex `vx`.\n\n# Fields\n - `vtconn`: vertex connectivity array\n - `vtconnoff`: offsets for vertex connectivity array\n - `vertmap`: map to vertex degrees of freedom: `vertmap[vx]` contains the \n degree of freedom located at vertex `vx`.\n - `data`: data field of MPIStateArray\n - `ivar`: variable # in the MPIStateArray\n - `vx`: unique edge number\n - `::Type{FT}`: Floating point type\n\"\"\"\nfunction dss_vertex!(\n vtconn,\n vtconnoff,\n vertmap,\n data,\n ivar,\n vx,\n ::Type{FT},\n) where {FT}\n @inbounds st = vtconnoff[vx]\n @inbounds nlvt = Int((vtconnoff[vx + 1] - vtconnoff[vx]) / 2)\n sumv = -FT(0)\n @inbounds for j in 1:nlvt\n lvt = vtconn[st + (j - 1) * 2]\n lelem = vtconn[st + (j - 1) * 2 + 1]\n loc = vertmap[lvt]\n sumv += data[loc, ivar, lelem]\n end\n @inbounds for j in 1:nlvt\n lvt = vtconn[st + (j - 1) * 2]\n lelem = vtconn[st + (j - 1) * 2 + 1]\n loc = vertmap[lvt]\n data[loc, ivar, lelem] = sumv\n end\nend\n\n\"\"\"\n dss_edge!(\n edgconn,\n edgconnoff,\n edgemap,\n Nemax,\n data,\n ivar,\n ex,\n ::Type{FT},\n) where {FT}\n\nThis function computes the direct stiffness summation for \nall degrees of freedom corresponding to edge `ex`. dss_edge!\napplies only to interior (non-vertex) edge nodes.\n\n# Fields\n - `edgconn`: edge connectivity array\n - `edgconnoff`: offsets for edge connectivity array\n - `edgemap`: map to edge degrees of freedom: \n `edgemap[i, edgno, orient]` contains the element node index of \n the `i`th interior node on edge `edgno`, under orientation `orient`. \n - `Nemax`: # of relevant degrees of freedom per edge (other dof are marked as -1)\n - `data`: data field of MPIStateArray\n - `ivar`: variable # in the MPIStateArray\n - `ex`: unique edge number\n - `::Type{FT}`: Floating point type\n\"\"\"\nfunction dss_edge!(\n edgconn,\n edgconnoff,\n edgemap,\n Nemax,\n data,\n ivar,\n ex,\n ::Type{FT},\n) where {FT}\n @inbounds st = edgconnoff[ex]\n @inbounds nledg = Int((edgconnoff[ex + 1] - edgconnoff[ex]) / 3)\n\n @inbounds for k in 1:Nemax\n sume = -FT(0)\n @inbounds for j in 1:nledg\n ledg = edgconn[st + (j - 1) * 3]\n lor = edgconn[st + (j - 1) * 3 + 1]\n lelem = edgconn[st + (j - 1) * 3 + 2]\n loc = edgemap[k, ledg, lor]\n if loc ≠ -1\n sume += data[loc, ivar, lelem]\n end\n end\n @inbounds for j in 1:nledg\n ledg = edgconn[st + (j - 1) * 3]\n lor = edgconn[st + (j - 1) * 3 + 1]\n lelem = edgconn[st + (j - 1) * 3 + 2]\n loc = edgemap[k, ledg, lor]\n if loc ≠ -1\n data[loc, ivar, lelem] = sume\n end\n end\n end\nend\n\"\"\"\n dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)\n\nThis function computes the direct stiffness summation for \nall degrees of freedom corresponding to face `fx`. dss_face!\napplies only to interior (non-vertex and non-edge) face nodes.\n\n# Fields\n - `fcconn`: face connectivity array\n - `facemap`: map to face degrees of freedom: `facemap[ij, fcno, orient]` \n contains the element node index of the `ij`th \n interior node on face `fcno` under orientation `orient`\n - `Nfmax`: # of relevant degrees of freedom per face (other dof are marked as -1)\n - `data`: data field of MPIStateArray\n - `ivar`: variable # in the MPIStateArray\n - `fx`: unique face number\n\"\"\"\nfunction dss_face!(fcconn, facemap, Nfmax, data, ivar, fx)\n @inbounds lfc = fcconn[fx, 1]\n @inbounds lel = fcconn[fx, 2]\n @inbounds nabrlfc = fcconn[fx, 3]\n @inbounds nabrlel = fcconn[fx, 4]\n @inbounds ordr = fcconn[fx, 5]\n ordr = ordr == 3 ? 2 : 1\n # mesh orientation 3 is a flip along the horizontal edges.\n # This is the only orientation currently support by the mesh generator\n # see vertsortandorder in\n # src/Numerics/Mesh/BrickMesh.jl \n @inbounds for j in 1:Nfmax\n loc1 = facemap[j, lfc, ordr]\n loc2 = facemap[j, nabrlfc, 1]\n if loc1 ≠ -1 && loc2 ≠ -1\n sumf = data[loc1, ivar, lel] + data[loc2, ivar, nabrlel]\n data[loc1, ivar, lel] = sumf\n data[loc2, ivar, nabrlel] = sumf\n end\n end\nend\n\nend\n" }, { "path": "src/Numerics/Mesh/Elements.jl", "content": "module Elements\nimport GaussQuadrature\n\n\"\"\"\n lglpoints(::Type{T}, N::Integer) where T <: AbstractFloat\n\nreturns the points `r` and weights `w` associated with the `N+1`-point\nGauss-Legendre-Lobatto quadrature rule of type `T`\n\n\"\"\"\nfunction lglpoints(::Type{T}, N::Integer) where {T <: AbstractFloat}\n @assert N ≥ 1\n GaussQuadrature.legendre(T, N + 1, GaussQuadrature.both)\nend\n\n\"\"\"\n glpoints(::Type{T}, N::Integer) where T <: AbstractFloat\n\nreturns the points `r` and weights `w` associated with the `N+1`-point\nGauss-Legendre quadrature rule of type `T`\n\"\"\"\nfunction glpoints(::Type{T}, N::Integer) where {T <: AbstractFloat}\n GaussQuadrature.legendre(T, N + 1, GaussQuadrature.neither)\nend\n\n\"\"\"\n baryweights(r)\n\nreturns the barycentric weights associated with the array of points `r`\n\nReference:\n [Berrut2004](@cite)\n\"\"\"\nfunction baryweights(r::AbstractVector{T}) where {T}\n Np = length(r)\n wb = ones(T, Np)\n\n for j in 1:Np\n for i in 1:Np\n if i != j\n wb[j] = wb[j] * (r[j] - r[i])\n end\n end\n wb[j] = T(1) / wb[j]\n end\n wb\nend\n\n\n\"\"\"\n spectralderivative(r::AbstractVector{T},\n wb=baryweights(r)::AbstractVector{T}) where T\n\nreturns the spectral differentiation matrix for a polynomial defined on the\npoints `r` with associated barycentric weights `wb`\n\nReference:\n - [Berrut2004](@cite)\n\"\"\"\nfunction spectralderivative(\n r::AbstractVector{T},\n wb = baryweights(r)::AbstractVector{T},\n) where {T}\n Np = length(r)\n @assert Np == length(wb)\n D = zeros(T, Np, Np)\n\n for k in 1:Np\n for j in 1:Np\n if k == j\n for l in 1:Np\n if l != k\n D[j, k] = D[j, k] + T(1) / (r[k] - r[l])\n end\n end\n else\n D[j, k] = (wb[k] / wb[j]) / (r[j] - r[k])\n end\n end\n end\n D\nend\n\n\"\"\"\n interpolationmatrix(rsrc::AbstractVector{T}, rdst::AbstractVector{T},\n wbsrc=baryweights(rsrc)::AbstractVector{T}) where T\n\nreturns the polynomial interpolation matrix for interpolating between the points\n`rsrc` (with associated barycentric weights `wbsrc`) and `rdst`\n\nReference:\n - [Berrut2004](@cite)\n\"\"\"\nfunction interpolationmatrix(\n rsrc::AbstractVector{T},\n rdst::AbstractVector{T},\n wbsrc = baryweights(rsrc)::AbstractVector{T},\n) where {T}\n Npdst = length(rdst)\n Npsrc = length(rsrc)\n @assert Npsrc == length(wbsrc)\n I = zeros(T, Npdst, Npsrc)\n for k in 1:Npdst\n for j in 1:Npsrc\n I[k, j] = wbsrc[j] / (rdst[k] - rsrc[j])\n if !isfinite(I[k, j])\n I[k, :] .= T(0)\n I[k, j] = T(1)\n break\n end\n end\n d = sum(I[k, :])\n I[k, :] = I[k, :] / d\n end\n I\nend\n\n\"\"\"\n jacobip(α, β, N, x)\n\nReturns a `(nx, N+1)` array containing the `N+1` Jacobi polynomials of order `N`, \nwith parameter `(α, β)`, evaluated on 1D grid `x`.\n\"\"\"\nfunction jacobip(α, β, N, x)\n a, b = GaussQuadrature.jacobi_coefs(N, α, β)\n V = GaussQuadrature.orthonormal_poly(x, a, b)\n return V\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/Filters.jl", "content": "module Filters\n\nusing SpecialFunctions\nusing LinearAlgebra, GaussQuadrature, KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\nusing StaticArrays\nusing ..Grids\nusing ..Grids: Direction, EveryDirection, HorizontalDirection, VerticalDirection\n\nusing ...MPIStateArrays\nusing ...VariableTemplates: @vars, varsize, Vars, varsindices\n\nexport AbstractSpectralFilter, AbstractFilter\nexport ExponentialFilter,\n CutoffFilter, MassPreservingCutoffFilter, TMARFilter, BoydVandevenFilter\n\nabstract type AbstractFilter end\nabstract type AbstractSpectralFilter <: AbstractFilter end\n\n\"\"\"\n AbstractFilterTarget\n\nAn abstract type representing variables that the filter\nwill act on\n\"\"\"\nabstract type AbstractFilterTarget end\n\n\"\"\"\n vars_state_filtered(::AbstractFilterTarget, FT)\n\nA tuple of symbols containing variables that the filter\nwill act on given a float type `FT`\n\"\"\"\nfunction vars_state_filtered end\n\n\"\"\"\n compute_filter_argument!(::AbstractFilterTarget,\n state_filter::Vars,\n state::Vars,\n state_auxiliary::Vars)\n\nCompute filter argument `state_filter` based on `state`\nand `state_auxiliary`\n\"\"\"\nfunction compute_filter_argument! end\n\"\"\"\n compute_filter_result!(::AbstractFilterTarget,\n state::Vars,\n state_filter::Vars,\n state_auxiliary::Vars)\n\nCompute filter result `state` based on the filtered state\n`state_filter` and `state_auxiliary`\n\"\"\"\nfunction compute_filter_result! end\n\nnumber_state_filtered(t::AbstractFilterTarget, FT) =\n varsize(vars_state_filtered(t, FT))\n\n\"\"\"\n FilterIndices(I)\n\nFilter variables based on their indices `I` where `I` can\nbe a range or a list of indices\n\n## Examples\n```julia\nFiltersIndices(1:3)\nFiltersIndices(1, 3, 5)\n```\n\"\"\"\nstruct FilterIndices{I} <: AbstractFilterTarget\n FilterIndices(I::Integer...) = new{I}()\n FilterIndices(I::AbstractRange) = new{I}()\nend\nvars_state_filtered(::FilterIndices{I}, FT) where {I} =\n @vars(_::SVector{length(I), FT})\n\nfunction compute_filter_argument!(\n ::FilterIndices{I},\n filter_state::Vars,\n state::Vars,\n aux::Vars,\n) where {I}\n @unroll for s in 1:length(I)\n @inbounds parent(filter_state)[s] = parent(state)[I[s]]\n end\nend\n\nfunction compute_filter_result!(\n ::FilterIndices{I},\n state::Vars,\n filter_state::Vars,\n aux::Vars,\n) where {I}\n @unroll for s in 1:length(I)\n @inbounds parent(state)[I[s]] = parent(filter_state)[s]\n end\nend\n\n\n\"\"\"\n spectral_filter_matrix(r, Nc, σ)\n\nReturns the filter matrix that takes function values at the interpolation\n`N+1` points, `r`, converts them into Legendre polynomial basis coefficients,\nmultiplies\n```math\nσ((n-N_c)/(N-N_c))\n```\nagainst coefficients `n=Nc:N` and evaluates the resulting polynomial at the\npoints `r`.\n\"\"\"\nfunction spectral_filter_matrix(r, Nc, σ)\n N = length(r) - 1\n T = eltype(r)\n\n @assert N >= 0\n @assert 0 <= Nc <= N\n\n a, b = GaussQuadrature.legendre_coefs(T, N)\n V = (N == 0 ? ones(T, 1, 1) : GaussQuadrature.orthonormal_poly(r, a, b))\n\n Σ = ones(T, N + 1)\n Σ[(Nc:N) .+ 1] .= σ.(((Nc:N) .- Nc) ./ (N - Nc))\n\n V * Diagonal(Σ) / V\nend\n\n\"\"\"\n modified_filter_matrix(r, Nc, σ)\n\nReturns the filter matrix that takes function values at the interpolation\n`N+1` points, `r`, converts them into Legendre polynomial basis coefficients,\nmultiplies\n```math\nσ((n-N_c)/(N-N_c))\n```\nagainst coefficients `n=Nc:N` and evaluates the resulting polynomial at the\npoints `r`. Unlike spectral_filter_matrix, this allows for the identity matrix,\nto be applied.\n\"\"\"\nfunction modified_filter_matrix(r, Nc, σ)\n N = length(r) - 1\n T = eltype(r)\n\n @assert N >= 0\n @assert 0 <= Nc\n\n Nc > N && return Array{T}(I, N + 1, N + 1)\n\n a, b = GaussQuadrature.legendre_coefs(T, N)\n V = (N == 0 ? ones(T, 1, 1) : GaussQuadrature.orthonormal_poly(r, a, b))\n\n Σ = ones(T, N + 1)\n Σ[(Nc:N) .+ 1] .= σ.(((Nc:N) .- Nc) ./ (N - Nc))\n\n V * Diagonal(Σ) / V\nend\n\n\"\"\"\n ExponentialFilter(grid, Nc=0, s=32, α=-log(eps(eltype(grid))))\n\nReturns the spectral filter with the filter function\n```math\nσ(η) = \\exp(-α η^s)\n```\nwhere `s` is the filter order (must be even), the filter starts with\npolynomial order `Nc`, and `alpha` is a parameter controlling the smallest\nvalue of the filter function.\n\"\"\"\nstruct ExponentialFilter{FM} <: AbstractSpectralFilter\n \"filter matrices in all directions (tuple of filter matrices)\"\n filter_matrices::FM\n\n function ExponentialFilter(\n grid,\n Nc = 0,\n s = 32,\n α = -log(eps(eltype(grid))),\n )\n dim = dimensionality(grid)\n\n # Support different filtering thresholds in different\n # directions (default behavior is to apply the same threshold\n # uniformly in all directions)\n if Nc isa Integer\n Nc = ntuple(i -> Nc, dim)\n elseif Nc isa NTuple{2} && dim == 3\n Nc = (Nc[1], Nc[1], Nc[2])\n end\n @assert length(Nc) == dim\n\n # Tuple of polynomial degrees (N₁, N₂, N₃)\n N = polynomialorders(grid)\n # In 2D, we assume same polynomial order in the horizontal\n @assert dim == 2 || N[1] == N[2]\n @assert iseven(s)\n @assert all(0 .<= Nc .<= N)\n\n σ(η) = exp(-α * η^s)\n\n AT = arraytype(grid)\n ξ = referencepoints(grid)\n filter_matrices =\n ntuple(i -> AT(spectral_filter_matrix(ξ[i], Nc[i], σ)), dim)\n new{typeof(filter_matrices)}(filter_matrices)\n end\nend\n\n\"\"\"\n BoydVandevenFilter(grid, Nc=0, s=32)\n\nReturns the spectral filter using the logarithmic error function of\nthe form:\n```math\nσ(η) = 1/2 erfc(2*sqrt(s)*χ(η)*(abs(η)-0.5))\n```\nwhenever s ≤ i ≤ N, and 1 otherwise. The function `χ(η)` is defined\nas\n```math\nχ(η) = sqrt(-log(1-4*(abs(η)-0.5)^2)/(4*(abs(η)-0.5)^2))\n```\nif `x != 0.5` and `1` otherwise. Here, `s` is the filter order,\nthe filter starts with polynomial order `Nc`, and `alpha` is a parameter\ncontrolling the smallest value of the filter function.\n\n### References\n - [Boyd1996](@cite)\n\"\"\"\nstruct BoydVandevenFilter{FM} <: AbstractSpectralFilter\n \"filter matrices in all directions (tuple of filter matrices)\"\n filter_matrices::FM\n\n function BoydVandevenFilter(grid, Nc = 0, s = 32)\n dim = dimensionality(grid)\n\n # Support different filtering thresholds in different\n # directions (default behavior is to apply the same threshold\n # uniformly in all directions)\n if Nc isa Integer\n Nc = ntuple(i -> Nc, dim)\n elseif Nc isa NTuple{2} && dim == 3\n Nc = (Nc[1], Nc[1], Nc[2])\n end\n @assert length(Nc) == dim\n\n # Tuple of polynomial degrees (N₁, N₂, N₃)\n N = polynomialorders(grid)\n # In 2D, we assume same polynomial order in the horizontal\n @assert dim == 2 || N[1] == N[2]\n @assert iseven(s)\n @assert all(0 .<= Nc .<= N)\n\n function σ(η)\n a = 2 * abs(η) - 1\n χ = iszero(a) ? one(a) : sqrt(-log1p(-a^2) / a^2)\n return erfc(sqrt(s) * χ * a) / 2\n end\n\n AT = arraytype(grid)\n ξ = referencepoints(grid)\n filter_matrices =\n ntuple(i -> AT(spectral_filter_matrix(ξ[i], Nc[i], σ)), dim)\n new{typeof(filter_matrices)}(filter_matrices)\n end\nend\n\n\"\"\"\n CutoffFilter(grid, Nc=polynomialorders(grid))\n\nReturns the spectral filter that zeros out polynomial modes greater than or\nequal to `Nc`.\n\"\"\"\nstruct CutoffFilter{FM} <: AbstractSpectralFilter\n \"filter matrices in all directions (tuple of filter matrices)\"\n filter_matrices::FM\n\n function CutoffFilter(grid, Nc = polynomialorders(grid))\n dim = dimensionality(grid)\n\n # Support different filtering thresholds in different\n # directions (default behavior is to apply the same threshold\n # uniformly in all directions)\n if Nc isa Integer\n Nc = ntuple(i -> Nc, dim)\n elseif Nc isa NTuple{2} && dim == 3\n Nc = (Nc[1], Nc[1], Nc[2])\n end\n @assert length(Nc) == dim\n\n # Tuple of polynomial degrees (N₁, N₂, N₃)\n N = polynomialorders(grid)\n # In 2D, we assume same polynomial order in the horizontal\n @assert dim == 2 || N[1] == N[2]\n @assert all(0 .<= Nc .<= N)\n\n σ(η) = 0\n\n AT = arraytype(grid)\n ξ = referencepoints(grid)\n filter_matrices =\n ntuple(i -> AT(spectral_filter_matrix(ξ[i], Nc[i], σ)), dim)\n new{typeof(filter_matrices)}(filter_matrices)\n end\nend\n\n\n\"\"\"\n MassPreservingCutoffFilter(grid, Nc=polynomialorders(grid))\n \nReturns the spectral filter that zeros out polynomial modes greater than or\nequal to `Nc` while preserving the cell average value. Use this filter if the\njacobian is nonconstant.\n\"\"\"\nstruct MassPreservingCutoffFilter{FM} <: AbstractSpectralFilter\n \"filter matrices in all directions (tuple of filter matrices)\"\n filter_matrices::FM\n\n function MassPreservingCutoffFilter(grid, Nc = polynomialorders(grid))\n dim = dimensionality(grid)\n\n # Support different filtering thresholds in different\n # directions (default behavior is to apply the same threshold\n # uniformly in all directions)\n if Nc isa Integer\n Nc = ntuple(i -> Nc, dim)\n elseif Nc isa NTuple{2} && dim == 3\n Nc = (Nc[1], Nc[1], Nc[2])\n end\n @assert length(Nc) == dim\n\n # Tuple of polynomial degrees (N₁, N₂, N₃)\n N = polynomialorders(grid)\n # In 2D, we assume same polynomial order in the horizontal\n @assert dim == 2 || N[1] == N[2]\n @assert all(0 .<= Nc)\n\n σ(η) = 0\n\n AT = arraytype(grid)\n ξ = referencepoints(grid)\n filter_matrices =\n ntuple(i -> AT(modified_filter_matrix(ξ[i], Nc[i], σ)), dim)\n new{typeof(filter_matrices)}(filter_matrices)\n end\nend\n\n\"\"\"\n TMARFilter()\n\nReturns the truncation and mass aware rescaling nonnegativity preservation\nfilter. The details of this filter are described in [Light2016](@cite)\n\nNote this needs to be used with a restrictive time step or a flux correction\nto ensure that grid integral is conserved.\n\n## Examples\n\nThis filter can be applied to the 3rd and 4th fields of an `MPIStateArray` `Q`\nwith the code\n\n```julia\nFilters.apply!(Q, (3, 4), grid, TMARFilter())\n```\n\nwhere `grid` is the associated `DiscontinuousSpectralElementGrid`.\n\"\"\"\nstruct TMARFilter <: AbstractFilter end\n\n\"\"\"\n Filters.apply!(Q::MPIStateArray,\n target,\n grid::DiscontinuousSpectralElementGrid,\n filter::AbstractSpectralFilter;\n kwargs...)\n\nApplies `filter` to `Q` given a `grid` and a custom `target`.\n\nA `target` can be any of the following:\n - a tuple or range of indices\n - a tuple of symbols or strings of variable names\n - a colon (`:`) to apply to all variables\n - a custom [`AbstractFilterTarget`]\n\nThe following keyword arguments are supported for some filters:\n- `direction`: for `AbstractSpectralFilter` controls if the filter is\n applied in the horizontal and/or vertical directions. It is assumed that the\n trailing dimension on the reference element is the vertical dimension and the\n rest are horizontal.\n- `state_auxiliary`: if `target` requires auxiliary state to compute its argument or results.\n\n# Examples\n\nSpecifying the `target` via indices:\n```julia\nFilters.apply!(Q, :, grid, TMARFilter())\nFilters.apply!(Q, (1, 3), grid, CutoffFilter(grid); direction=VerticalDirection())\n```\n\nSpeciying `target` via symbols or strings:\n```julia\nFilters.apply!(Q, (:ρ, \"energy.ρe\"), grid, TMARFilter())\nFilters.apply!(Q, (\"moisture.ρq_tot\",), grid, CutoffFilter(grid);\n direction=VerticalDirection())\n```\n\"\"\"\nfunction apply!(\n Q,\n target,\n grid::DiscontinuousSpectralElementGrid,\n filter::AbstractFilter;\n kwargs...,\n)\n device = typeof(Q.data) <: Array ? CPU() : CUDADevice()\n event = Event(device)\n event =\n apply_async!(Q, target, grid, filter; dependencies = event, kwargs...)\n wait(device, event)\nend\n\n\n\"\"\"\n Filters.apply_async!(Q, target, grid::DiscontinuousSpectralElementGrid,\n filter::AbstractFilter;\n dependencies,\n kwargs...)\n\nAn asynchronous version of [`Filters.apply!`](@ref), returning an `Event`\nobject. `dependencies` should be an `Event` or tuple of `Event`s which need to\nfinish before applying the filter.\n\n```julia\ncompstream = Filters.apply_async!(Q, :, grid, CutoffFilter(grid); dependencies=compstream)\nwait(compstream)\n```\n\"\"\"\nfunction apply_async! end\n\nfunction apply_async!(\n Q,\n target::AbstractFilterTarget,\n grid::DiscontinuousSpectralElementGrid,\n filter::AbstractSpectralFilter;\n dependencies,\n state_auxiliary = nothing,\n direction = EveryDirection(),\n)\n topology = grid.topology\n\n # Tuple of polynomial degrees (N₁, N₂, N₃)\n N = polynomialorders(grid)\n # In 2D, we assume same polynomial order in the horizontal\n dim = dimensionality(grid)\n # Currently only support same polynomial in both horizontal directions\n @assert N[1] == N[2]\n\n device = typeof(Q.data) <: Array ? CPU() : CUDADevice()\n\n nelem = length(topology.elems)\n # Number of Gauss-Lobatto quadrature points in each direction\n Nq = N .+ 1\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nq3 = dim == 2 ? 1 : Nq[dim]\n\n nrealelem = length(topology.realelems)\n event = dependencies\n\n if direction isa EveryDirection || direction isa HorizontalDirection\n @assert dim == 2 || Nq1 == Nq2\n filtermatrix = filter.filter_matrices[1]\n event = kernel_apply_filter!(device, (Nq1, Nq2, Nq3))(\n Val(dim),\n Val(N),\n Val(vars(Q)),\n Val(isnothing(state_auxiliary) ? nothing : vars(state_auxiliary)),\n HorizontalDirection(),\n Q.data,\n isnothing(state_auxiliary) ? nothing : state_auxiliary.data,\n target,\n filtermatrix,\n ndrange = (nrealelem * Nq1, Nq2, Nq3),\n dependencies = event,\n )\n end\n if direction isa EveryDirection || direction isa VerticalDirection\n filtermatrix = filter.filter_matrices[end]\n event = kernel_apply_filter!(device, (Nq1, Nq2, Nq3))(\n Val(dim),\n Val(N),\n Val(vars(Q)),\n Val(isnothing(state_auxiliary) ? nothing : vars(state_auxiliary)),\n VerticalDirection(),\n Q.data,\n isnothing(state_auxiliary) ? nothing : state_auxiliary.data,\n target,\n filtermatrix,\n ndrange = (nrealelem * Nq1, Nq2, Nq3),\n dependencies = event,\n )\n end\n return event\nend\n\n\nfunction apply_async!(\n Q,\n target::AbstractFilterTarget,\n grid::DiscontinuousSpectralElementGrid,\n ::TMARFilter;\n dependencies,\n)\n topology = grid.topology\n\n device = typeof(Q.data) <: Array ? CPU() : CUDADevice()\n\n dim = dimensionality(grid)\n N = polynomialorders(grid)\n # Currently only support same polynomial in both horizontal directions\n @assert dim == 2 || N[1] == N[2]\n Nqs = N .+ 1\n Nq = Nqs[1]\n Nqj = dim == 2 ? 1 : Nqs[2]\n\n nrealelem = length(topology.realelems)\n nreduce = 2^ceil(Int, log2(Nq * Nqj))\n\n event = dependencies\n event = kernel_apply_TMAR_filter!(device, (Nq, Nqj), (nrealelem * Nq, Nqj))(\n Val(nreduce),\n Val(dim),\n Val(N),\n Q.data,\n target,\n grid.vgeo,\n dependencies = event,\n )\n return event\nend\n\nfunction apply_async!(\n Q,\n target::AbstractFilterTarget,\n grid::DiscontinuousSpectralElementGrid,\n filter::MassPreservingCutoffFilter;\n dependencies,\n state_auxiliary = nothing,\n direction = EveryDirection(),\n)\n topology = grid.topology\n\n device = typeof(Q.data) <: Array ? CPU() : CUDADevice()\n\n dim = dimensionality(grid)\n N = polynomialorders(grid)\n # Currently only support same polynomial in both horizontal directions\n @assert dim == 2 || N[1] == N[2]\n Nq = N .+ 1\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nq3 = dim == 2 ? 1 : Nq[dim]\n\n nrealelem = length(topology.realelems)\n # parallel sum info\n nreduce = 2^ceil(Int, log2(Nq1 * Nq2 * Nq3))\n event = dependencies\n\n if direction isa EveryDirection || direction isa HorizontalDirection\n @assert dim == 2 || Nq1 == Nq2\n filtermatrix = filter.filter_matrices[1]\n event = kernel_apply_mp_filter!(device, (Nq1, Nq2, Nq3))(\n Val(nreduce),\n Val(dim),\n Val(N),\n Val(vars(Q)),\n Val(isnothing(state_auxiliary) ? nothing : vars(state_auxiliary)),\n HorizontalDirection(),\n Q.data,\n isnothing(state_auxiliary) ? nothing : state_auxiliary.data,\n target,\n filtermatrix,\n grid.vgeo,\n ndrange = (nrealelem * Nq1, Nq2, Nq3),\n dependencies = event,\n )\n end\n if direction isa EveryDirection || direction isa VerticalDirection\n filtermatrix = filter.filter_matrices[end]\n event = kernel_apply_mp_filter!(device, (Nq1, Nq2, Nq3))(\n Val(nreduce),\n Val(dim),\n Val(N),\n Val(vars(Q)),\n Val(isnothing(state_auxiliary) ? nothing : vars(state_auxiliary)),\n VerticalDirection(),\n Q.data,\n isnothing(state_auxiliary) ? nothing : state_auxiliary.data,\n target,\n filtermatrix,\n grid.vgeo,\n ndrange = (nrealelem * Nq1, Nq2, Nq3),\n dependencies = event,\n )\n end\n return event\nend\n\nfunction apply_async!(\n Q,\n indices::Union{Colon, AbstractRange, Tuple{Vararg{Integer}}},\n grid::DiscontinuousSpectralElementGrid,\n filter::AbstractFilter;\n kwargs...,\n)\n if indices isa Colon\n indices = 1:size(Q, 2)\n end\n apply_async!(Q, FilterIndices(indices...), grid, filter; kwargs...)\nend\n\nfunction apply_async!(\n Q,\n vs::Tuple,\n grid::DiscontinuousSpectralElementGrid,\n filter::AbstractFilter;\n kwargs...,\n)\n apply_async!(\n Q,\n FilterIndices(varsindices(vars(Q), vs)...),\n grid,\n filter;\n kwargs...,\n )\nend\n\nconst _M = Grids._M\n\n@doc \"\"\"\n kernel_apply_filter!(::Val{dim}, ::Val{N}, direction,\n Q, state_auxiliary, target, filtermatrix\n ) where {dim, N}\n\nComputational kernel: Applies the `filtermatrix` to `Q` given a\ncustom target `target` while preserving the cell average.\n\nThe `direction` argument is used to control if the filter is applied in the\nhorizontal and/or vertical reference directions.\n\"\"\" kernel_apply_filter!\n@kernel function kernel_apply_filter!(\n ::Val{dim},\n ::Val{N},\n ::Val{vars_Q},\n ::Val{vars_state_auxiliary},\n direction,\n Q,\n state_auxiliary,\n target::AbstractFilterTarget,\n filtermatrix,\n) where {dim, N, vars_Q, vars_state_auxiliary}\n @uniform begin\n FT = eltype(Q)\n\n Nqs = N .+ 1\n Nq1 = Nqs[1]\n Nq2 = Nqs[2]\n Nq3 = dim == 2 ? 1 : Nqs[dim]\n\n if direction isa EveryDirection\n filterinξ1 = filterinξ2 = true\n filterinξ3 = dim == 2 ? false : true\n elseif direction isa HorizontalDirection\n filterinξ1 = true\n filterinξ2 = dim == 2 ? false : true\n filterinξ3 = false\n elseif direction isa VerticalDirection\n filterinξ1 = false\n filterinξ2 = dim == 2 ? true : false\n filterinξ3 = dim == 2 ? false : true\n end\n\n nstates = varsize(vars_Q)\n nfilterstates = number_state_filtered(target, FT)\n nfilteraux =\n isnothing(state_auxiliary) ? 0 : varsize(vars_state_auxiliary)\n\n # ugly workaround around problems with @private\n # hopefully will be soon fixed in KA\n l_Q2 = MVector{nstates, FT}(undef)\n l_Qfiltered2 = MVector{nfilterstates, FT}(undef)\n end\n\n s_Q = @localmem FT (Nq1, Nq2, Nq3, nfilterstates)\n l_Q = @private FT (nstates,)\n l_Qfiltered = @private FT (nfilterstates,)\n l_aux = @private FT (nfilteraux,)\n\n e = @index(Group, Linear)\n i, j, k = @index(Local, NTuple)\n\n @inbounds begin\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n\n @unroll for s in 1:nstates\n l_Q[s] = Q[ijk, s, e]\n end\n\n @unroll for s in 1:nfilteraux\n l_aux[s] = state_auxiliary[ijk, s, e]\n end\n\n fill!(l_Qfiltered2, -zero(FT))\n\n compute_filter_argument!(\n target,\n Vars{vars_state_filtered(target, FT)}(l_Qfiltered2),\n Vars{vars_Q}(l_Q[:]),\n Vars{vars_state_auxiliary}(l_aux[:]),\n )\n\n @unroll for fs in 1:nfilterstates\n l_Qfiltered[fs] = zero(FT)\n end\n\n @unroll for fs in 1:nfilterstates\n s_Q[i, j, k, fs] = l_Qfiltered2[fs]\n end\n\n if filterinξ1\n @synchronize\n @unroll for n in 1:Nq1\n @unroll for fs in 1:nfilterstates\n l_Qfiltered[fs] += filtermatrix[i, n] * s_Q[n, j, k, fs]\n end\n end\n\n if filterinξ2 || filterinξ3\n @synchronize\n @unroll for fs in 1:nfilterstates\n s_Q[i, j, k, fs] = l_Qfiltered[fs]\n l_Qfiltered[fs] = zero(FT)\n end\n end\n end\n\n if filterinξ2\n @synchronize\n @unroll for n in 1:Nq2\n @unroll for fs in 1:nfilterstates\n l_Qfiltered[fs] += filtermatrix[j, n] * s_Q[i, n, k, fs]\n end\n end\n\n if filterinξ3\n @synchronize\n @unroll for fs in 1:nfilterstates\n s_Q[i, j, k, fs] = l_Qfiltered[fs]\n l_Qfiltered[fs] = zero(FT)\n end\n end\n end\n\n if filterinξ3\n @synchronize\n @unroll for n in 1:Nq3\n @unroll for fs in 1:nfilterstates\n l_Qfiltered[fs] += filtermatrix[k, n] * s_Q[i, j, n, fs]\n end\n end\n end\n\n @unroll for s in 1:nstates\n l_Q2[s] = l_Q[s]\n end\n\n compute_filter_result!(\n target,\n Vars{vars_Q}(l_Q2),\n Vars{vars_state_filtered(target, FT)}(l_Qfiltered[:]),\n Vars{vars_state_auxiliary}(l_aux[:]),\n )\n\n # Store result\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n @unroll for s in 1:nstates\n Q[ijk, s, e] = l_Q2[s]\n end\n\n @synchronize\n end\nend\n\n@kernel function kernel_apply_TMAR_filter!(\n ::Val{nreduce},\n ::Val{dim},\n ::Val{N},\n Q,\n target::FilterIndices{I},\n vgeo,\n) where {nreduce, dim, N, I}\n @uniform begin\n FT = eltype(Q)\n\n Nqs = N .+ 1\n Nq1 = Nqs[1]\n Nq2 = dim == 2 ? 1 : Nqs[2]\n Nq3 = Nqs[end]\n\n nfilterstates = number_state_filtered(target, FT)\n nelemperblock = 1\n end\n\n l_Q = @private FT (nfilterstates, Nq1)\n l_MJ = @private FT (Nq1,)\n\n s_MJQ = @localmem FT (Nq1 * Nq2, nfilterstates)\n s_MJQclipped = @localmem FT (Nq1 * Nq2, nfilterstates)\n\n e = @index(Group, Linear)\n i, j = @index(Local, NTuple)\n\n @inbounds begin\n # loop up the pencil and load Q and MJ\n @unroll for k in 1:Nq3\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n\n @unroll for sf in 1:nfilterstates\n s = I[sf]\n l_Q[sf, k] = Q[ijk, s, e]\n end\n\n l_MJ[k] = vgeo[ijk, _M, e]\n end\n\n @unroll for sf in 1:nfilterstates\n MJQ, MJQclipped = zero(FT), zero(FT)\n\n @unroll for k in 1:Nq3\n MJ = l_MJ[k]\n Qs = l_Q[sf, k]\n Qsclipped = Qs ≥ 0 ? Qs : zero(Qs)\n\n MJQ += MJ * Qs\n MJQclipped += MJ * Qsclipped\n end\n\n ij = i + Nq1 * (j - 1)\n\n s_MJQ[ij, sf] = MJQ\n s_MJQclipped[ij, sf] = MJQclipped\n end\n @synchronize\n\n @unroll for n in 11:-1:1\n if nreduce ≥ 2^n\n ij = i + Nq1 * (j - 1)\n ijshift = ij + 2^(n - 1)\n if ij ≤ 2^(n - 1) && ijshift ≤ Nq1 * Nq2\n @unroll for sf in 1:nfilterstates\n s_MJQ[ij, sf] += s_MJQ[ijshift, sf]\n s_MJQclipped[ij, sf] += s_MJQclipped[ijshift, sf]\n end\n end\n @synchronize\n end\n end\n\n @unroll for sf in 1:nfilterstates\n qs_average = s_MJQ[1, sf]\n qs_clipped_average = s_MJQclipped[1, sf]\n\n r = qs_average > 0 ? qs_average / qs_clipped_average : zero(FT)\n\n s = I[sf]\n @unroll for k in 1:Nq3\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n\n Qs = l_Q[sf, k]\n Q[ijk, s, e] = Qs ≥ 0 ? r * Qs : zero(Qs)\n end\n end\n end\nend\n\n\n\"\"\"\n kernel_apply_mp_filter!(::Val{dim}, ::Val{N}, direction,\n Q, state_auxiliary, target, filtermatrix\n ) where {dim, N}\nComputational kernel: Applies the `filtermatrix` to `Q` given a\ncustom target `target`.\nThe `direction` argument is used to control if the filter is applied in the\n\"\"\"\n@kernel function kernel_apply_mp_filter!(\n ::Val{nreduce},\n ::Val{dim},\n ::Val{N},\n ::Val{vars_Q},\n ::Val{vars_state_auxiliary},\n direction,\n Q,\n state_auxiliary,\n target::AbstractFilterTarget,\n filtermatrix,\n vgeo,\n) where {nreduce, dim, N, vars_Q, vars_state_auxiliary}\n @uniform begin\n FT = eltype(Q)\n\n Nqs = N .+ 1\n Nq1 = Nqs[1]\n Nq2 = Nqs[2]\n Nq3 = dim == 2 ? 1 : Nqs[dim]\n\n if direction isa EveryDirection\n filterinξ1 = filterinξ2 = true\n filterinξ3 = dim == 2 ? false : true\n elseif direction isa HorizontalDirection\n filterinξ1 = true\n filterinξ2 = dim == 2 ? false : true\n filterinξ3 = false\n elseif direction isa VerticalDirection\n filterinξ1 = false\n filterinξ2 = dim == 2 ? true : false\n filterinξ3 = dim == 2 ? false : true\n end\n\n nstates = varsize(vars_Q)\n nfilterstates = number_state_filtered(target, FT)\n nfilteraux =\n isnothing(state_auxiliary) ? 0 : varsize(vars_state_auxiliary)\n\n # ugly workaround around problems with @private\n # hopefully will be soon fixed in KA\n u_Q = MVector{nstates, FT}(undef)\n u_Qfiltered = MVector{nfilterstates, FT}(undef)\n end\n\n l_Q = @localmem FT (Nq1, Nq2, Nq3, nfilterstates) # element local \n l_MQᴮ = @localmem FT (Nq1 * Nq2 * Nq3, nstates) # before applying filter\n l_MQᴬ = @localmem FT (Nq1 * Nq2 * Nq3, nstates) # after applying filter\n l_M = @localmem FT (Nq1 * Nq2 * Nq3) # local mass matrix\n\n p_Q = @private FT (nstates,)\n p_Qfiltered = @private FT (nfilterstates,) # scratch space for storing mat mul\n p_aux = @private FT (nfilteraux,)\n\n e = @index(Group, Linear)\n i, j, k = @index(Local, NTuple)\n ijk = @index(Local, Linear)\n\n @inbounds begin\n\n @unroll for s in 1:nstates\n p_Q[s] = Q[ijk, s, e]\n end\n\n @unroll for s in 1:nfilteraux\n p_aux[s] = state_auxiliary[ijk, s, e]\n end\n\n # Load mass matrix and pre-filtered mass weighted quantities to shared memory\n l_M[ijk] = vgeo[ijk, _M, e]\n @unroll for s in 1:nstates\n l_MQᴮ[ijk, s] = l_M[ijk] * p_Q[s]\n end\n\n fill!(u_Qfiltered, -zero(FT))\n\n compute_filter_argument!(\n target,\n Vars{vars_state_filtered(target, FT)}(u_Qfiltered),\n Vars{vars_Q}(p_Q[:]),\n Vars{vars_state_auxiliary}(p_aux[:]),\n )\n\n @unroll for fs in 1:nfilterstates\n p_Qfiltered[fs] = zero(FT)\n end\n\n @unroll for fs in 1:nfilterstates\n l_Q[i, j, k, fs] = u_Qfiltered[fs]\n end\n\n if filterinξ1\n @synchronize\n @unroll for n in 1:Nq1\n @unroll for fs in 1:nfilterstates\n p_Qfiltered[fs] += filtermatrix[i, n] * l_Q[n, j, k, fs]\n end\n end\n\n if filterinξ2 || filterinξ3\n @synchronize\n @unroll for fs in 1:nfilterstates\n l_Q[i, j, k, fs] = p_Qfiltered[fs]\n p_Qfiltered[fs] = zero(FT)\n end\n end\n end\n\n if filterinξ2\n @synchronize\n @unroll for n in 1:Nq2\n @unroll for fs in 1:nfilterstates\n p_Qfiltered[fs] += filtermatrix[j, n] * l_Q[i, n, k, fs]\n end\n end\n\n if filterinξ3\n @synchronize\n @unroll for fs in 1:nfilterstates\n l_Q[i, j, k, fs] = p_Qfiltered[fs]\n p_Qfiltered[fs] = zero(FT)\n end\n end\n end\n\n if filterinξ3\n @synchronize\n @unroll for n in 1:Nq3\n @unroll for fs in 1:nfilterstates\n p_Qfiltered[fs] += filtermatrix[k, n] * l_Q[i, j, n, fs]\n end\n end\n end\n\n # work around for not being able to `Vars` `@private` arrays\n @unroll for s in 1:nstates\n u_Q[s] = p_Q[s]\n end\n\n compute_filter_result!(\n target,\n Vars{vars_Q}(u_Q),\n Vars{vars_state_filtered(target, FT)}(p_Qfiltered[:]),\n Vars{vars_state_auxiliary}(p_aux[:]),\n )\n # Store result and post-filtered mass weighted quantities\n @unroll for s in 1:nstates\n p_Q[s] = u_Q[s]\n l_MQᴬ[ijk, s] = l_M[ijk] * p_Q[s]\n end\n\n @synchronize\n @unroll for n in 11:-1:1\n if nreduce ≥ (1 << n)\n ijkshift = ijk + (1 << (n - 1))\n if ijk ≤ (1 << (n - 1)) && ijkshift ≤ Nq1 * Nq2 * Nq3\n l_M[ijk] += l_M[ijkshift]\n @unroll for s in 1:nstates\n l_MQᴮ[ijk, s] += l_MQᴮ[ijkshift, s]\n l_MQᴬ[ijk, s] += l_MQᴬ[ijkshift, s]\n end\n end\n @synchronize\n end\n end\n\n @synchronize\n M⁻¹ = 1 / l_M[1]\n # Reset the element average and store result\n @unroll for s in 1:nstates\n Q[ijk, s, e] = p_Q[s] + M⁻¹ * (l_MQᴮ[1, s] - l_MQᴬ[1, s])\n end\n\n @synchronize\n end\n\nend\n\nend # end of module\n" }, { "path": "src/Numerics/Mesh/GeometricFactors.jl", "content": "module GeometricFactors\n\nexport VolumeGeometry, SurfaceGeometry\n\n\"\"\"\n VolumeGeometry{Nq, AA <: AbstractArray, A <: AbstractArray}\n\nA struct that collects `VolumeGeometry` fields:\n- array: Array contatining the data stored in a VolumeGeometry struct (the following fields are views into this array)\n- ∂ξk/∂xi: Derivative of the Cartesian reference element coordinate `ξ_k` with respect to the Cartesian physical coordinate `x_i`\n- ωJ: Mass matrix. This is the physical mass matrix, and thus contains the Jacobian determinant, J .* (ωᵢ ⊗ ωⱼ ⊗ ωₖ), where ωᵢ are the quadrature weights and J is the Jacobian determinant, det(∂x/∂ξ)\n- ωJI: Inverse mass matrix: 1 ./ ωJ\n- ωJH: Horizontal mass matrix (used in diagnostics), J .* norm(∂ξ3/∂x) * (ωᵢ ⊗ ωⱼ); for integrating over a plane (in 2-D ξ2 is used, not ξ3)\n- xi: Nodal degrees of freedom locations in Cartesian physical space\n- JcV: Metric terms for vertical line integrals norm(∂x/∂ξ3) (in 2-D ξ2 is used, not ξ3)\n- ∂xk/∂ξi: Inverse of matrix `∂ξk/∂xi` that represents the derivative of Cartesian physical coordinate `x_i` with respect to Cartesian reference element coordinate `ξ_k`\n\"\"\"\nstruct VolumeGeometry{Nq, AA <: AbstractArray, A <: AbstractArray}\n array::AA\n ξ1x1::A\n ξ2x1::A\n ξ3x1::A\n ξ1x2::A\n ξ2x2::A\n ξ3x2::A\n ξ1x3::A\n ξ2x3::A\n ξ3x3::A\n ωJ::A\n ωJI::A\n ωJH::A\n x1::A\n x2::A\n x3::A\n JcV::A\n x1ξ1::A\n x2ξ1::A\n x3ξ1::A\n x1ξ2::A\n x2ξ2::A\n x3ξ2::A\n x1ξ3::A\n x2ξ3::A\n x3ξ3::A\n function VolumeGeometry{Nq}(\n array::AA,\n args::A...,\n ) where {Nq, AA, A <: AbstractArray}\n new{Nq, AA, A}(array, args...)\n end\nend\n\n\"\"\"\n VolumeGeometry(FT, Nq::NTuple{N, Int}, nelems::Int)\n\nConstruct an empty `VolumeGeometry` object, in `FT` precision.\n- `Nq` is a tuple containing the number of quadrature points in each direction.\n- `nelem` is the number of elements.\n\"\"\"\nfunction VolumeGeometry(FT, Nq::NTuple{N, Int}, nelem::Int) where {N}\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n array = zeros(FT, prod(Nq), fieldcount(VolumeGeometry) - 1, nelem)\n VolumeGeometry{Nq}(\n array,\n ntuple(j -> @view(array[:, j, :]), fieldcount(VolumeGeometry) - 1)...,\n )\nend\n\n\"\"\"\n SurfaceGeometry{Nq, AA, A <: AbstractArray}\n\nA struct that collects `VolumeGeometry` fields:\n- array: Array contatining the data stored in a SurfaceGeometry struct (the following fields are views into this array)\n- ni: Outward pointing unit normal in physical space\n- sωJ: Surface mass matrix. This is the physical mass matrix, and thus contains the surface Jacobian determinant, sJ .* (ωⱼ ⊗ ωₖ), where ωᵢ are the quadrature weights and sJ is the surface Jacobian determinant\n- vωJI: Volume mass matrix at the surface nodes (needed in the lift operation, i.e., the projection of a face field back to the volume). Since DGSEM is used only collocated, volume mass matrices are required.\n\"\"\"\nstruct SurfaceGeometry{Nq, AA, A <: AbstractArray}\n array::AA\n n1::A\n n2::A\n n3::A\n sωJ::A\n vωJI::A\n function SurfaceGeometry{Nq}(\n array::AA,\n args::A...,\n ) where {Nq, AA, A <: AbstractArray}\n new{Nq, AA, A}(array, args...)\n end\nend\n\n\"\"\"\n SurfaceGeometry(FT, Nq::NTuple{N, Int}, nface::Int, nelem::Int)\n\nConstruct an empty `SurfaceGeometry` object, in `FT` precision.\n- `Nq` is a tuple containing the number of quadrature points in each direction.\n- `nface` is the number of faces.\n- `nelem` is the number of elements.\n\"\"\"\nfunction SurfaceGeometry(\n FT,\n Nq::NTuple{N, Int},\n nface::Int,\n nelem::Int,\n) where {N}\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n array =\n zeros(FT, fieldcount(SurfaceGeometry) - 1, maximum(Nfp), nface, nelem)\n SurfaceGeometry{Nfp}(\n array,\n ntuple(\n j -> @view(array[j, :, :, :]),\n fieldcount(SurfaceGeometry) - 1,\n )...,\n )\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/Geometry.jl", "content": "module Geometry\n\nusing StaticArrays, LinearAlgebra, DocStringExtensions\nusing KernelAbstractions.Extras: @unroll\nusing ..Grids:\n _ξ1x1,\n _ξ2x1,\n _ξ3x1,\n _ξ1x2,\n _ξ2x2,\n _ξ3x2,\n _ξ1x3,\n _ξ2x3,\n _ξ3x3,\n _M,\n _MI,\n _x1,\n _x2,\n _x3,\n _JcV\n\nexport LocalGeometry, lengthscale, resolutionmetric, lengthscale_horizontal\n\n\"\"\"\n LocalGeometry\n\nThe local geometry at a nodal point.\n\n# Constructors\n\n LocalGeometry{Np, N}(vgeo::AbstractArray{T}, n::Integer, e::Integer)\n\nExtracts a `LocalGeometry` object from the `vgeo` array at node `n` in element\n`e` with `Np` being the number of points in the element and `N` being the\npolynomial order\n\n# Fields\n\n- `polyorder`\n\n polynomial order of the element\n\n- `coord`\n\n local degree of freedom Cartesian coordinate \n\n- `invJ`\n\n Jacobian from Cartesian to element coordinates: `invJ[i,j]` is ``∂ξ_i / ∂x_j``\n\n$(DocStringExtensions.FIELDS)\n\"\"\"\nstruct LocalGeometry{Np, N, AT, IT}\n \"Global volume geometry array\"\n vgeo::AT\n \"element local linear node index\"\n n::IT\n \"process local element index\"\n e::IT\n\n LocalGeometry{Np, N}(vgeo::AT, n::IT, e::IT) where {Np, N, AT, IT} =\n new{Np, N, AT, IT}(vgeo, n, e)\nend\n\n@inline function Base.getproperty(\n geo::LocalGeometry{Np, N},\n sym::Symbol,\n) where {Np, N}\n @inbounds if sym === :polyorder\n return N\n elseif sym === :coord\n vgeo, n, e = getfield(geo, :vgeo), getfield(geo, :n), getfield(geo, :e)\n FT = eltype(vgeo)\n return @SVector FT[vgeo[n, _x1, e], vgeo[n, _x2, e], vgeo[n, _x3, e]]\n elseif sym === :invJ\n vgeo, n, e = getfield(geo, :vgeo), getfield(geo, :n), getfield(geo, :e)\n FT = eltype(vgeo)\n return @SMatrix FT[\n vgeo[n, _ξ1x1, e] vgeo[n, _ξ1x2, e] vgeo[n, _ξ1x3, e]\n vgeo[n, _ξ2x1, e] vgeo[n, _ξ2x2, e] vgeo[n, _ξ2x3, e]\n vgeo[n, _ξ3x1, e] vgeo[n, _ξ3x2, e] vgeo[n, _ξ3x3, e]\n ]\n elseif sym === :center_coord\n vgeo, n, e = getfield(geo, :vgeo), getfield(geo, :n), getfield(geo, :e)\n FT = eltype(vgeo)\n coords = SVector(vgeo[n, _x1, e], vgeo[n, _x2, e], vgeo[n, _x3, e])\n V = FT(0)\n xc = FT(0)\n yc = FT(0)\n zc = FT(0)\n @unroll for i in 1:Np\n M = vgeo[i, _M, e]\n V += M\n xc += M * vgeo[i, _x1, e]\n yc += M * vgeo[i, _x2, e]\n zc += M * vgeo[i, _x3, e]\n end\n return SVector(xc / V, yc / V, zc / V)\n else\n return getfield(geo, sym)\n end\nend\n\n\"\"\"\n resolutionmetric(g::LocalGeometry)\n\nThe metric tensor of the discretisation resolution. Given a unit vector `u` in\nCartesian coordinates and `M = resolutionmetric(g)`, `sqrt(u'*M*u)` is the\ndegree-of-freedom density in the direction of `u`.\n\"\"\"\nfunction resolutionmetric(g::LocalGeometry)\n S = g.polyorder * g.invJ / 2\n S' * S # TODO: return an eigendecomposition / symmetric object?\nend\n\n\"\"\"\n lengthscale(g::LocalGeometry)\n\nThe effective geometric mean grid resolution at the point.\n\"\"\"\nlengthscale(g::LocalGeometry) =\n 2 / (cbrt(det(g.invJ)) * maximum(max.(1, g.polyorder)))\n\n\"\"\"\n lengthscale_horizontal(g::LocalGeometry)\n\nThe effective horizontal grid resolution at the point.\n\"\"\"\nfunction lengthscale_horizontal(g::LocalGeometry)\n # inverse Jacobian matrix:\n #\n # invJ[i,j] = ∂ξ_i / ∂x_j\n invJ = g.invJ\n\n # The local horizontal grid stretchings are:\n #\n # horizontal direction 1: √( ∑_j ∂x_1 / ∂ξ_j) * 2\n # horizontal direction 2: √( ∑_j ∂x_2 / ∂ξ_j) * 2\n #\n # To get these we need to have the Jacobian matrix:\n #\n # J[i,j] = ∂x_i / ∂ξ_j\n #\n # To get this we can solve the system with the unit basis vectors\n # (division by `N` gives the local average node distance)\n Δ1 = norm(invJ \\ SVector(1, 0, 0)) * 2 / g.polyorder[1]\n Δ2 = norm(invJ \\ SVector(0, 1, 0)) * 2 / g.polyorder[2]\n\n # Set the horizontal length scale to the average of the two calculated values\n return (Δ1 + Δ2) / 2\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/Grids.jl", "content": "module Grids\nusing ..Topologies, ..GeometricFactors\nimport ..Metrics, ..Elements\nimport ..BrickMesh\nusing ClimateMachine.MPIStateArrays\n\nusing MPI\nusing LinearAlgebra\nusing KernelAbstractions\nusing DocStringExtensions\n\nexport DiscontinuousSpectralElementGrid, AbstractGrid\nexport dofs_per_element, arraytype, dimensionality, polynomialorders\nexport referencepoints, min_node_distance, get_z, computegeometry\nexport EveryDirection, HorizontalDirection, VerticalDirection, Direction\n\nabstract type Direction end\nstruct EveryDirection <: Direction end\nstruct HorizontalDirection <: Direction end\nstruct VerticalDirection <: Direction end\nBase.in(::T, ::S) where {T <: Direction, S <: Direction} = T == S\n\n\"\"\"\n MinNodalDistance{FT}\n\nA struct containing the minimum nodal distance\nalong horizontal and vertical directions.\n\"\"\"\nstruct MinNodalDistance{FT}\n \"horizontal\"\n h::FT\n \"vertical\"\n v::FT\nend\n\nabstract type AbstractGrid{\n FloatType,\n dim,\n polynomialorder,\n numberofDOFs,\n DeviceArray,\n} end\n\ndofs_per_element(::AbstractGrid{T, D, N, Np}) where {T, D, N, Np} = Np\n\npolynomialorders(::AbstractGrid{T, dim, N}) where {T, dim, N} = N\n\ndimensionality(::AbstractGrid{T, dim}) where {T, dim} = dim\n\nBase.eltype(::AbstractGrid{T}) where {T} = T\n\narraytype(::AbstractGrid{T, D, N, Np, DA}) where {T, D, N, Np, DA} = DA\n\n\"\"\"\n referencepoints(::AbstractGrid)\n\nReturns the points on the reference element.\n\"\"\"\nreferencepoints(::AbstractGrid) = error(\"needs to be implemented\")\n\n\"\"\"\n min_node_distance(::AbstractGrid, direction::Direction=EveryDirection() )\n\nReturns an approximation of the minimum node distance in physical space.\n\"\"\"\nfunction min_node_distance(\n ::AbstractGrid,\n direction::Direction = EveryDirection(),\n)\n error(\"needs to be implemented\")\nend\n\n# {{{\n# `vgeo` stores geometry and metrics terms at the volume quadrature /\n# interpolation points\nconst _nvgeo = 16\nconst _ξ1x1,\n_ξ2x1,\n_ξ3x1,\n_ξ1x2,\n_ξ2x2,\n_ξ3x2,\n_ξ1x3,\n_ξ2x3,\n_ξ3x3,\n_M,\n_MI,\n_MH,\n_x1,\n_x2,\n_x3,\n_JcV = 1:_nvgeo\nconst vgeoid = (\n # ∂ξk/∂xi: derivative of the Cartesian reference element coordinate `ξ_k`\n # with respect to the Cartesian physical coordinate `x_i`\n ξ1x1id = _ξ1x1,\n ξ2x1id = _ξ2x1,\n ξ3x1id = _ξ3x1,\n ξ1x2id = _ξ1x2,\n ξ2x2id = _ξ2x2,\n ξ3x2id = _ξ3x2,\n ξ1x3id = _ξ1x3,\n ξ2x3id = _ξ2x3,\n ξ3x3id = _ξ3x3,\n # M refers to the mass matrix. This is the physical mass matrix, and thus\n # contains the Jacobian determinant:\n # J .* (ωᵢ ⊗ ωⱼ ⊗ ωₖ)\n # where ωᵢ are the quadrature weights and J is the Jacobian determinant\n # det(∂x/∂ξ)\n Mid = _M,\n # Inverse mass matrix: 1 ./ M\n MIid = _MI,\n # Horizontal mass matrix (used in diagnostics)\n # J .* norm(∂ξ3/∂x) * (ωᵢ ⊗ ωⱼ); for integrating over a plane\n # (in 2-D ξ2 not ξ3 is used)\n MHid = _MH,\n # Nodal degrees of freedom locations in Cartesian physical space\n x1id = _x1,\n x2id = _x2,\n x3id = _x3,\n # Metric terms for vertical line integrals\n # norm(∂x/∂ξ3)\n # (in 2-D ξ2 not ξ3 is used)\n JcVid = _JcV,\n)\n\n# `sgeo` stores geometry and metrics terms at the surface quadrature /\n# interpolation points\nconst _nsgeo = 5\nconst _n1, _n2, _n3, _sM, _vMI = 1:_nsgeo\nconst sgeoid = (\n # outward pointing unit normal in physical space\n n1id = _n1,\n n2id = _n2,\n n3id = _n3,\n # sM refers to the surface mass matrix. This is the physical mass matrix,\n # and thus contains the surface Jacobian determinant:\n # sJ .* (ωⱼ ⊗ ωₖ)\n # where ωᵢ are the quadrature weights and sJ is the surface Jacobian\n # determinant\n sMid = _sM,\n # Volume mass matrix at the surface nodes (needed in the lift operation,\n # i.e., the projection of a face field back to the volume). Since DGSEM is\n # used only collocated, volume mass matrices are required.\n vMIid = _vMI,\n)\n# }}}\n\n\"\"\"\n DiscontinuousSpectralElementGrid(topology; FloatType, DeviceArray,\n polynomialorder,\n meshwarp = (x...)->identity(x))\n\nGenerate a discontinuous spectral element (tensor product,\nLegendre-Gauss-Lobatto) grid/mesh from a `topology`, where the order of the\nelements is given by `polynomialorder`. `DeviceArray` gives the array type used\nto store the data (`CuArray` or `Array`), and the coordinate points will be of\n`FloatType`.\n\nThe polynomial order can be different in each direction (specified as a\n`NTuple`). If only a single integer is specified, then each dimension will use\nthe same order. If the topology dimension is 3 and the `polynomialorder` has\ndimension 2, then the first value will be used for horizontal and the second for\nthe vertical.\n\nThe optional `meshwarp` function allows the coordinate points to be warped after\nthe mesh is created; the mesh degrees of freedom are orginally assigned using a\ntrilinear blend of the element corner locations.\n\"\"\"\nstruct DiscontinuousSpectralElementGrid{\n T,\n dim,\n N,\n Np,\n DA,\n DAT1,\n DAT2,\n DAT3,\n DAT4,\n DAI1,\n DAI2,\n DAI3,\n TOP,\n TVTK,\n MINΔ,\n} <: AbstractGrid{T, dim, N, Np, DA}\n \"mesh topology\"\n topology::TOP\n\n \"volume metric terms\"\n vgeo::DAT3\n\n \"surface metric terms\"\n sgeo::DAT4\n\n \"element to boundary condition map\"\n elemtobndy::DAI2\n\n \"volume DOF to element minus side map\"\n vmap⁻::DAI3\n\n \"volume DOF to element plus side map\"\n vmap⁺::DAI3\n\n \"list of DOFs that need to be received (in neighbors order)\"\n vmaprecv::DAI1\n\n \"list of DOFs that need to be sent (in neighbors order)\"\n vmapsend::DAI1\n\n \"An array of ranges in `vmaprecv` to receive from each neighbor\"\n nabrtovmaprecv::Any\n\n \"An array of ranges in `vmapsend` to send to each neighbor\"\n nabrtovmapsend::Any\n\n \"Array of real elements that do not have a ghost element as a neighbor\"\n interiorelems::Any\n\n \"Array of real elements that have at least one ghost element as a neighbor\"\n exteriorelems::Any\n\n \"Array indicating if a degree of freedom (real or ghost) is active\"\n activedofs::Any\n\n \"1-D LGL weights on the device (one for each dimension)\"\n ω::DAT1\n\n \"1-D derivative operator on the device (one for each dimension)\"\n D::DAT2\n\n \"1-D indefinite integral operator on the device (one for each dimension)\"\n Imat::DAT2\n\n \"\"\"\n tuple of (x1, x2, x3) to use for vtk output (Needed for the `N = 0` case) in\n other cases these match `vgeo` values\n \"\"\"\n x_vtk::TVTK\n\n \"\"\"\n Minimum nodal distances for horizontal and vertical directions\n \"\"\"\n minΔ::MINΔ\n\n \"\"\"\n Map to vertex degrees of freedom: `vertmap[v]` contains the degree of freedom located at vertex `v`.\n \"\"\"\n vertmap::Union{DAI1, Nothing}\n\n \"\"\"\n Map to edge degrees of freedom: `edgemap[i, edgno, orient]` contains the element node index of \n the `i`th interior node on edge `edgno`, under orientation `orient`.\n \"\"\"\n edgemap::Union{DAI3, Nothing}\n\n \"\"\"\n Map to face degrees of freedom: `facemap[ij, fcno, orient]` contains the element node index of the `ij`th \n interior node on face `fcno` under orientation `orient`.\n\n Note that only the two orientations that are generated for stacked meshes are currently supported, i.e.,\n mesh orientation `3` as defined by `BrickMesh` gets mapped to orientation `2` for this data structure. \n \"\"\"\n facemap::Union{DAI3, Nothing}\n\n # Constructor for a tuple of polynomial orders\n function DiscontinuousSpectralElementGrid(\n topology::AbstractTopology{dim};\n polynomialorder,\n FloatType,\n DeviceArray,\n meshwarp::Function = (x...) -> identity(x),\n ) where {dim}\n\n if polynomialorder isa Integer\n polynomialorder = ntuple(j -> polynomialorder, dim)\n elseif polynomialorder isa NTuple{2} && dim == 3\n polynomialorder =\n (polynomialorder[1], polynomialorder[1], polynomialorder[2])\n end\n\n @assert dim == length(polynomialorder)\n\n N = polynomialorder\n\n (vmap⁻, vmap⁺) = mappings(\n N,\n topology.elemtoelem,\n topology.elemtoface,\n topology.elemtoordr,\n )\n\n (vmaprecv, nabrtovmaprecv) = commmapping(\n N,\n topology.ghostelems,\n topology.ghostfaces,\n topology.nabrtorecv,\n )\n (vmapsend, nabrtovmapsend) = commmapping(\n N,\n topology.sendelems,\n topology.sendfaces,\n topology.nabrtosend,\n )\n\n Np = prod(N .+ 1)\n\n vertmap, edgemap, facemap = init_vertex_edge_face_mappings(N)\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n N[j] == 0 ? Elements.glpoints(FloatType, N[j]) :\n Elements.lglpoints(FloatType, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n\n Imat = ntuple(\n j -> indefinite_integral_interpolation_matrix(ξ[j], ω[j]),\n dim,\n )\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n (vgeo, sgeo, x_vtk) =\n computegeometry(topology.elemtocoord, D, ξ, ω, meshwarp)\n\n vgeo = vgeo.array\n sgeo = sgeo.array\n @assert Np == size(vgeo, 1)\n\n activedofs = zeros(Bool, Np * length(topology.elems))\n activedofs[1:(Np * length(topology.realelems))] .= true\n activedofs[vmaprecv] .= true\n\n # Create arrays on the device\n vgeo = DeviceArray(vgeo)\n sgeo = DeviceArray(sgeo)\n elemtobndy = DeviceArray(topology.elemtobndy)\n vmap⁻ = DeviceArray(vmap⁻)\n vmap⁺ = DeviceArray(vmap⁺)\n vmapsend = DeviceArray(vmapsend)\n vmaprecv = DeviceArray(vmaprecv)\n activedofs = DeviceArray(activedofs)\n ω = DeviceArray.(ω)\n D = DeviceArray.(D)\n Imat = DeviceArray.(Imat)\n\n # FIXME: There has got to be a better way!\n DAT1 = typeof(ω)\n DAT2 = typeof(D)\n DAT3 = typeof(vgeo)\n DAT4 = typeof(sgeo)\n DAI1 = typeof(vmapsend)\n DAI2 = typeof(elemtobndy)\n DAI3 = typeof(vmap⁻)\n TOP = typeof(topology)\n TVTK = typeof(x_vtk)\n if vertmap isa Array\n vertmap = DAI1(vertmap)\n end\n if edgemap isa Array\n edgemap = DAI3(edgemap)\n end\n if facemap isa Array\n facemap = DAI3(facemap)\n end\n FT = FloatType\n minΔ = MinNodalDistance(\n min_node_distance(vgeo, topology, N, FT, HorizontalDirection()),\n min_node_distance(vgeo, topology, N, FT, VerticalDirection()),\n )\n MINΔ = typeof(minΔ)\n\n new{\n FloatType,\n dim,\n N,\n Np,\n DeviceArray,\n DAT1,\n DAT2,\n DAT3,\n DAT4,\n DAI1,\n DAI2,\n DAI3,\n TOP,\n TVTK,\n MINΔ,\n }(\n topology,\n vgeo,\n sgeo,\n elemtobndy,\n vmap⁻,\n vmap⁺,\n vmaprecv,\n vmapsend,\n nabrtovmaprecv,\n nabrtovmapsend,\n DeviceArray(topology.interiorelems),\n DeviceArray(topology.exteriorelems),\n activedofs,\n ω,\n D,\n Imat,\n x_vtk,\n minΔ,\n vertmap,\n edgemap,\n facemap,\n )\n end\nend\n\n\"\"\"\n referencepoints(::DiscontinuousSpectralElementGrid)\n\nReturns the 1D interpolation points used for the reference element.\n\"\"\"\nfunction referencepoints(\n ::DiscontinuousSpectralElementGrid{FT, dim, N},\n) where {FT, dim, N}\n ξω = ntuple(\n j ->\n N[j] == 0 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, _ = ntuple(j -> map(x -> x[j], ξω), 2)\n return ξ\nend\n\n\"\"\"\n min_node_distance(\n ::DiscontinuousSpectralElementGrid,\n direction::Direction=EveryDirection())\n )\n\nReturns an approximation of the minimum node distance in physical space along\nthe reference coordinate directions. The direction controls which reference\ndirections are considered.\n\"\"\"\nmin_node_distance(\n grid::DiscontinuousSpectralElementGrid,\n direction::Direction = EveryDirection(),\n) = min_node_distance(grid.minΔ, direction)\n\nmin_node_distance(minΔ::MinNodalDistance, ::VerticalDirection) = minΔ.v\nmin_node_distance(minΔ::MinNodalDistance, ::HorizontalDirection) = minΔ.h\nmin_node_distance(minΔ::MinNodalDistance, ::EveryDirection) =\n min(minΔ.h, minΔ.v)\n\nfunction min_node_distance(\n vgeo,\n topology::AbstractTopology{dim},\n N,\n ::Type{T},\n direction::Direction = EveryDirection(),\n) where {T, dim}\n topology = topology\n nrealelem = length(topology.realelems)\n if nrealelem > 0\n Nq = N .+ 1\n Np = prod(Nq)\n device = vgeo isa Array ? CPU() : CUDADevice()\n min_neighbor_distance = similar(vgeo, Np, nrealelem)\n event = Event(device)\n event = kernel_min_neighbor_distance!(device, min(Np, 1024))(\n Val(N),\n Val(dim),\n direction,\n min_neighbor_distance,\n vgeo,\n topology.realelems;\n ndrange = (Np * nrealelem),\n dependencies = (event,),\n )\n wait(device, event)\n locmin = minimum(min_neighbor_distance)\n else\n locmin = typemax(T)\n end\n\n MPI.Allreduce(locmin, min, topology.mpicomm)\nend\n\n\"\"\"\n get_z(grid; z_scale = 1, rm_dupes = false)\n\nGet the Gauss-Lobatto points along the Z-coordinate.\n\n - `grid`: DG grid\n - `z_scale`: multiplies `z-coordinate`\n - `rm_dupes`: removes duplicate Gauss-Lobatto points\n\"\"\"\nfunction get_z(\n grid::DiscontinuousSpectralElementGrid{T, dim, N};\n z_scale = 1,\n rm_dupes = false,\n) where {T, dim, N}\n # Assumes same polynomial orders in all horizontal directions\n @assert dim < 3 || N[1] == N[2]\n Nhoriz = N[1]\n Nvert = N[end]\n Nph = (Nhoriz + 1)^2\n Np = Nph * (Nvert + 1)\n if last(polynomialorders(grid)) == 0\n rm_dupes = false # no duplicates in FVM\n end\n if rm_dupes\n ijk_range = (1:Nph:(Np - Nph))\n vgeo = Array(grid.vgeo)\n z = reshape(vgeo[ijk_range, _x3, :], :)\n z = [z..., vgeo[Np, _x3, end]]\n return z * z_scale\n else\n ijk_range = (1:Nph:Np)\n z = Array(reshape(grid.vgeo[ijk_range, _x3, :], :))\n return z * z_scale\n end\n return reshape(grid.vgeo[(1:Nph:Np), _x3, :], :) * z_scale\nend\n\nfunction Base.getproperty(G::DiscontinuousSpectralElementGrid, s::Symbol)\n if s ∈ keys(vgeoid)\n vgeoid[s]\n elseif s ∈ keys(sgeoid)\n sgeoid[s]\n else\n getfield(G, s)\n end\nend\n\nfunction Base.propertynames(G::DiscontinuousSpectralElementGrid)\n (\n fieldnames(DiscontinuousSpectralElementGrid)...,\n keys(vgeoid)...,\n keys(sgeoid)...,\n )\nend\n\n# {{{ mappings\n\"\"\"\n mappings(N, elemtoelem, elemtoface, elemtoordr)\n\nThis function takes in a tuple of polynomial orders `N` and parts of a topology\n(as returned from `connectmesh`) and returns index mappings for the element\nsurface flux computation. The returned `Tuple` contains:\n\n - `vmap⁻` an array of linear indices into the volume degrees of freedom where\n `vmap⁻[:,f,e]` are the degrees of freedom indices for face `f` of element\n `e`.\n\n - `vmap⁺` an array of linear indices into the volume degrees of freedom where\n `vmap⁺[:,f,e]` are the degrees of freedom indices for the face neighboring\n face `f` of element `e`.\n\"\"\"\nfunction mappings(N, elemtoelem, elemtoface, elemtoordr)\n nfaces, nelem = size(elemtoelem)\n\n d = div(nfaces, 2)\n Nq = N .+ 1\n # number of points in the element\n Np = prod(Nq)\n\n # Compute the maximum number of points on a face\n Nfp = div.(Np, Nq)\n\n # linear index for each direction, e.g., (i, j, k) -> n\n p = reshape(1:Np, ntuple(j -> Nq[j], d))\n\n # fmask[f] -> returns an array of all degrees of freedom on face f\n fmask = if d == 1\n (\n p[1:1], # Face 1\n p[Nq[1]:Nq[1]], # Face 2\n )\n elseif d == 2\n (\n p[1, :][:], # Face 1\n p[Nq[1], :][:], # Face 2\n p[:, 1][:], # Face 3\n p[:, Nq[2]][:], # Face 4\n )\n elseif d == 3\n (\n p[1, :, :][:], # Face 1\n p[Nq[1], :, :][:], # Face 2\n p[:, 1, :][:], # Face 3\n p[:, Nq[2], :][:], # Face 4\n p[:, :, 1][:], # Face 5\n p[:, :, Nq[3]][:], # Face 6\n )\n else\n error(\"unknown dimensionality\")\n end\n\n # Create a map from Cartesian face dof number to linear face dof numbering\n # inds[face][i, j] -> n\n inds = ntuple(\n f -> dropdims(\n LinearIndices(ntuple(j -> j == cld(f, 2) ? 1 : Nq[j], d));\n dims = cld(f, 2),\n ),\n nfaces,\n )\n\n # Use the largest possible storage\n vmap⁻ = fill!(similar(elemtoelem, maximum(Nfp), nfaces, nelem), 0)\n vmap⁺ = fill!(similar(elemtoelem, maximum(Nfp), nfaces, nelem), 0)\n\n for e1 in 1:nelem, f1 in 1:nfaces\n e2 = elemtoelem[f1, e1]\n f2 = elemtoface[f1, e1]\n o2 = elemtoordr[f1, e1]\n d1, d2 = cld(f1, 2), cld(f2, 2)\n\n # Check to make sure the dof grid is conforming\n @assert Nfp[d1] == Nfp[d2]\n\n # Always pull out minus side without any flips / rotations\n vmap⁻[1:Nfp[d1], f1, e1] .= Np * (e1 - 1) .+ fmask[f1][1:Nfp[d1]][:]\n\n # Orientation codes defined in BrickMesh.jl (arbitrary numbers in 3D)\n if o2 == 1 # Neighbor oriented same as minus\n vmap⁺[1:Nfp[d1], f1, e1] .= Np * (e2 - 1) .+ fmask[f2][1:Nfp[d1]][:]\n elseif d == 3 && o2 == 3 # Neighbor fliped in first index\n vmap⁺[1:Nfp[d1], f1, e1] =\n Np * (e2 - 1) .+ fmask[f2][inds[f2][end:-1:1, :]][:]\n else\n error(\"Orientation '$o2' with dim '$d' not supported yet\")\n end\n end\n\n (vmap⁻, vmap⁺)\nend\n# }}}\n\nfunction init_vertex_edge_face_mappings(N)\n dim = length(N)\n Np = N .+ 1\n if dim == 3 && Np[end] > 2\n nodes = reshape(1:prod(Np), Np)\n vertmap =\n Int64.([\n nodes[1, 1, 1],\n nodes[Np[1], 1, 1],\n nodes[1, Np[2], 1],\n nodes[Np[1], Np[2], 1],\n nodes[1, 1, Np[3]],\n nodes[Np[1], 1, Np[3]],\n nodes[1, Np[2], Np[3]],\n nodes[Np[1], Np[2], Np[3]],\n ])\n Ne = Np .- 2\n Ne_max = maximum(Ne)\n if Ne_max ≥ 1\n edgemap = -ones(Int64, Ne_max, 12, 2)\n\n if Np[1] > 2\n edgemap[1:Ne[1], 1, 1] .= nodes[2:(end - 1), 1, 1]\n edgemap[1:Ne[1], 2, 1] .= nodes[2:(end - 1), Np[2], 1]\n edgemap[1:Ne[1], 3, 1] .= nodes[2:(end - 1), 1, Np[3]]\n edgemap[1:Ne[1], 4, 1] .= nodes[2:(end - 1), Np[2], Np[3]]\n\n edgemap[1:Ne[1], 1, 2] .= nodes[(end - 1):-1:2, 1, 1]\n edgemap[1:Ne[1], 2, 2] .= nodes[(end - 1):-1:2, Np[2], 1]\n edgemap[1:Ne[1], 3, 2] .= nodes[(end - 1):-1:2, 1, Np[3]]\n edgemap[1:Ne[1], 4, 2] .= nodes[(end - 1):-1:2, Np[2], Np[3]]\n end\n\n if Np[2] > 2\n edgemap[1:Ne[2], 5, 1] .= nodes[1, 2:(end - 1), 1]\n edgemap[1:Ne[2], 6, 1] .= nodes[Np[1], 2:(end - 1), 1]\n edgemap[1:Ne[2], 7, 1] .= nodes[1, 2:(end - 1), Np[3]]\n edgemap[1:Ne[2], 8, 1] .= nodes[Np[1], 2:(end - 1), Np[3]]\n\n edgemap[1:Ne[2], 5, 2] .= nodes[1, (end - 1):-1:2, 1]\n edgemap[1:Ne[2], 6, 2] .= nodes[Np[1], (end - 1):-1:2, 1]\n edgemap[1:Ne[2], 7, 2] .= nodes[1, (end - 1):-1:2, Np[3]]\n edgemap[1:Ne[2], 8, 2] .= nodes[Np[1], (end - 1):-1:2, Np[3]]\n end\n\n if Np[3] > 2\n edgemap[1:Ne[3], 9, 1] .= nodes[1, 1, 2:(end - 1)]\n edgemap[1:Ne[3], 10, 1] .= nodes[Np[1], 1, 2:(end - 1)]\n edgemap[1:Ne[3], 11, 1] .= nodes[1, Np[2], 2:(end - 1)]\n edgemap[1:Ne[3], 12, 1] .= nodes[Np[1], Np[2], 2:(end - 1)]\n\n edgemap[1:Ne[3], 9, 2] .= nodes[1, 1, (end - 1):-1:2]\n edgemap[1:Ne[3], 10, 2] .= nodes[Np[1], 1, (end - 1):-1:2]\n edgemap[1:Ne[3], 11, 2] .= nodes[1, Np[2], (end - 1):-1:2]\n edgemap[1:Ne[3], 12, 2] .= nodes[Np[1], Np[2], (end - 1):-1:2]\n end\n else\n edgemap = nothing\n end\n\n Nf = Np .- 2\n Nf_max = maximum([Nf[1] * Nf[2], Nf[2] * Nf[3], Nf[1] * Nf[3]])\n if Nf_max ≥ 1\n facemap = -ones(Int64, Nf_max, 6, 2)\n\n if Nf[2] > 0 && Nf[3] > 0\n nfc = Nf[2] * Nf[3]\n facemap[1:nfc, 1, 1] .= nodes[1, 2:(end - 1), 2:(end - 1)][:]\n facemap[1:nfc, 2, 1] .=\n nodes[Np[1], 2:(end - 1), 2:(end - 1)][:]\n\n facemap[1:nfc, 1, 2] .= nodes[1, (end - 1):-1:2, 2:(end - 1)][:]\n facemap[1:nfc, 2, 2] .=\n nodes[Np[1], (end - 1):-1:2, 2:(end - 1)][:]\n end\n\n if Nf[1] > 0 && Nf[3] > 0\n nfc = Nf[1] * Nf[3]\n facemap[1:nfc, 3, 1] .= nodes[2:(end - 1), 1, 2:(end - 1)][:]\n facemap[1:nfc, 4, 1] .=\n nodes[2:(end - 1), Np[2], 2:(end - 1)][:]\n\n facemap[1:nfc, 3, 2] .= nodes[(end - 1):-1:2, 1, 2:(end - 1)][:]\n facemap[1:nfc, 4, 2] .=\n nodes[(end - 1):-1:2, Np[2], 2:(end - 1)][:]\n end\n\n if Nf[1] > 0 && Nf[2] > 0\n nfc = Nf[1] * Nf[2]\n facemap[1:nfc, 5, 1] .= nodes[2:(end - 1), 2:(end - 1), 1][:]\n facemap[1:nfc, 6, 1] .=\n nodes[2:(end - 1), 2:(end - 1), Np[3]][:]\n\n facemap[1:nfc, 5, 2] .= nodes[(end - 1):-1:2, 2:(end - 1), 1][:]\n facemap[1:nfc, 6, 2] .=\n nodes[(end - 1):-1:2, 2:(end - 1), Np[3]][:]\n end\n else\n facemap = nothing\n end\n else\n vertmap, edgemap, facemap = nothing, nothing, nothing\n end\n\n return vertmap, edgemap, facemap\nend\n\n\n\"\"\"\n commmapping(N, commelems, commfaces, nabrtocomm)\n\nThis function takes in a tuple of polynomial orders `N` and parts of a mesh (as\nreturned from `connectmesh` such as `sendelems`, `sendfaces`, and `nabrtosend`)\nand returns index mappings for the element surface flux parallel communcation.\nThe returned `Tuple` contains:\n\n - `vmapC` an array of linear indices into the volume degrees of freedom to be\n communicated.\n\n - `nabrtovmapC` a range in `vmapC` to communicate with each neighbor.\n\"\"\"\nfunction commmapping(N, commelems, commfaces, nabrtocomm)\n nface, nelem = size(commfaces)\n\n @assert nelem == length(commelems)\n\n d = div(nface, 2)\n Nq = N .+ 1\n Np = prod(Nq)\n\n vmapC = similar(commelems, nelem * Np)\n nabrtovmapC = similar(nabrtocomm)\n\n i = 1\n e = 1\n for neighbor in 1:length(nabrtocomm)\n rbegin = i\n for ne in nabrtocomm[neighbor]\n ce = commelems[ne]\n\n # Whole element sending\n # for n = 1:Np\n # vmapC[i] = (ce-1)*Np + n\n # i += 1\n # end\n\n CI = CartesianIndices(ntuple(j -> 1:Nq[j], d))\n for (ci, li) in zip(CI, LinearIndices(CI))\n addpoint = false\n for j in 1:d\n addpoint |=\n (commfaces[2 * (j - 1) + 1, e] && ci[j] == 1) ||\n (commfaces[2 * (j - 1) + 2, e] && ci[j] == Nq[j])\n end\n\n if addpoint\n vmapC[i] = (ce - 1) * Np + li\n i += 1\n end\n end\n\n e += 1\n end\n rend = i - 1\n\n nabrtovmapC[neighbor] = rbegin:rend\n end\n\n resize!(vmapC, i - 1)\n\n (vmapC, nabrtovmapC)\nend\n\n# Compute geometry FVM version\nfunction computegeometry_fvm(elemtocoord, D, ξ, ω, meshwarp)\n FT = eltype(D[1])\n dim = length(D)\n nface = 2dim\n nelem = size(elemtocoord, 3)\n\n Nq = ntuple(j -> size(D[j], 1), dim)\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n Nq_N1 = max.(Nq, 2)\n Np_N1 = prod(Nq_N1)\n Nfp_N1 = div.(Np_N1, Nq_N1)\n\n # First we compute the geometry with all the N = 0 dimension set to N = 1\n # so that we can later compute the geometry for the case N = 0, as the\n # average of two N = 1 cases\n ξ1, ω1 = Elements.lglpoints(FT, 1)\n D1 = Elements.spectralderivative(ξ1)\n D_N1 = ntuple(j -> Nq[j] == 1 ? D1 : D[j], dim)\n ξ_N1 = ntuple(j -> Nq[j] == 1 ? ξ1 : ξ[j], dim)\n ω_N1 = ntuple(j -> Nq[j] == 1 ? ω1 : ω[j], dim)\n (vgeo_N1, sgeo_N1, x_vtk) =\n computegeometry(elemtocoord, D_N1, ξ_N1, ω_N1, meshwarp)\n\n # Allocate the storage for N = 0 volume metrics\n vgeo = VolumeGeometry(FT, Nq, nelem)\n\n # Counter to make sure we got all the vgeo terms\n num_vgeo_handled = 0\n\n Metrics.creategrid!(vgeo, elemtocoord, ξ)\n\n x1 = vgeo.x1\n x2 = vgeo.x2\n x3 = vgeo.x3\n @inbounds for j in 1:length(vgeo.x1)\n (x1[j], x2[j], x3[j]) = meshwarp(vgeo.x1[j], vgeo.x2[j], vgeo.x3[j])\n end\n\n # Update data in vgeo\n vgeo.x1 .= x1\n vgeo.x2 .= x2\n vgeo.x3 .= x3\n num_vgeo_handled += 3\n\n @views begin\n # ωJ should be a sum\n ωJ_N1 = reshape(vgeo_N1.ωJ, (Nq_N1..., nelem))\n vgeo.ωJ[:] .= sum(ωJ_N1, dims = findall(Nq .== 1))[:]\n num_vgeo_handled += 1\n\n # need to recompute ωJI\n vgeo.ωJI .= 1 ./ vgeo.ωJ\n num_vgeo_handled += 1\n\n # coordinates should just be averages\n avg_den = 2^sum(Nq .== 1)\n JcV_N1 = reshape(vgeo_N1.JcV, (Nq_N1..., nelem))\n vgeo.JcV[:] .= sum(JcV_N1, dims = findall(Nq .== 1))[:] ./ avg_den\n num_vgeo_handled += 1\n\n # For the metrics it is J * ξixk we approximate so multiply and divide the\n # mass matrix (which has the Jacobian determinant and the proper averaging\n # due to the quadrature weights)\n ωJ_N1 = reshape(vgeo_N1.ωJ, (Nq_N1..., nelem))\n ωJI = vgeo.ωJI\n\n ξ1x1_N1 = reshape(vgeo_N1.ξ1x1, (Nq_N1..., nelem))\n ξ2x1_N1 = reshape(vgeo_N1.ξ2x1, (Nq_N1..., nelem))\n ξ3x1_N1 = reshape(vgeo_N1.ξ3x1, (Nq_N1..., nelem))\n ξ1x2_N1 = reshape(vgeo_N1.ξ1x2, (Nq_N1..., nelem))\n ξ2x2_N1 = reshape(vgeo_N1.ξ2x2, (Nq_N1..., nelem))\n ξ3x2_N1 = reshape(vgeo_N1.ξ3x2, (Nq_N1..., nelem))\n ξ1x3_N1 = reshape(vgeo_N1.ξ1x3, (Nq_N1..., nelem))\n ξ2x3_N1 = reshape(vgeo_N1.ξ2x3, (Nq_N1..., nelem))\n ξ3x3_N1 = reshape(vgeo_N1.ξ3x3, (Nq_N1..., nelem))\n\n vgeo.ξ1x1[:] .=\n sum(ωJ_N1 .* ξ1x1_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ2x1[:] .=\n sum(ωJ_N1 .* ξ2x1_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ3x1[:] .=\n sum(ωJ_N1 .* ξ3x1_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ1x2[:] .=\n sum(ωJ_N1 .* ξ1x2_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ2x2[:] .=\n sum(ωJ_N1 .* ξ2x2_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ3x2[:] .=\n sum(ωJ_N1 .* ξ3x2_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ1x3[:] .=\n sum(ωJ_N1 .* ξ1x3_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ2x3[:] .=\n sum(ωJ_N1 .* ξ2x3_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n vgeo.ξ3x3[:] .=\n sum(ωJ_N1 .* ξ3x3_N1, dims = findall(Nq .== 1))[:] .* ωJI[:]\n num_vgeo_handled += 9\n\n # compute ωJH and JvC\n horizontal_metrics!(vgeo, Nq, ω)\n num_vgeo_handled += 1\n\n # Make sure we handled all the vgeo terms\n @assert _nvgeo == num_vgeo_handled\n end\n\n # Sort out the sgeo terms\n @views begin\n sgeo = SurfaceGeometry(FT, Nq, nface, nelem)\n\n # for the volume inverse mass matrix\n p = reshape(1:Np, Nq)\n if dim == 1\n fmask = (p[1:1], p[Nq[1]:Nq[1]])\n elseif dim == 2\n fmask = (p[1, :][:], p[Nq[1], :][:], p[:, 1][:], p[:, Nq[2]][:])\n elseif dim == 3\n fmask = (\n p[1, :, :][:],\n p[Nq[1], :, :][:],\n p[:, 1, :][:],\n p[:, Nq[2], :][:],\n p[:, :, 1][:],\n p[:, :, Nq[3]][:],\n )\n end\n for d in 1:dim\n for f in (2d - 1):(2d)\n # number of points matches means that we keep all the data\n # (N = 0 is not on the face)\n if Nfp[d] == Nfp_N1[d]\n sgeo.n1[:, f, :] .= sgeo_N1.n1[:, f, :]\n sgeo.n2[:, f, :] .= sgeo_N1.n2[:, f, :]\n sgeo.n3[:, f, :] .= sgeo_N1.n3[:, f, :]\n sgeo.sωJ[:, f, :] .= sgeo_N1.sωJ[:, f, :]\n\n # Volume inverse mass will be wrong so reset it\n sgeo.vωJI[:, f, :] .= vgeo.ωJI[fmask[f], :]\n else\n # Counter to make sure we got all the sgeo terms\n num_sgeo_handled = 0\n\n # sum to get sM\n Nq_f = (Nq[1:(d - 1)]..., Nq[(d + 1):dim]...)\n Nq_f_N1 = (Nq_N1[1:(d - 1)]..., Nq_N1[(d + 1):dim]...)\n sM_N1 = reshape(\n sgeo_N1.sωJ[1:Nfp_N1[d], f, :],\n Nq_f_N1...,\n nelem,\n )\n sgeo.sωJ[1:Nfp[d], f, :][:] .=\n sum(sM_N1, dims = findall(Nq_f .== 1))[:]\n num_sgeo_handled += 1\n\n # Normals (like metrics in the volume) need to be computed\n # scaled by surface Jacobian which we can do with the\n # surface mass matrices\n sM = sgeo.sωJ[1:Nfp[d], f, :]\n\n fld_N1_n1 = reshape(\n sgeo_N1.n1[1:Nfp_N1[d], f, :],\n Nq_f_N1...,\n nelem,\n )\n fld_N1_n2 = reshape(\n sgeo_N1.n2[1:Nfp_N1[d], f, :],\n Nq_f_N1...,\n nelem,\n )\n fld_N1_n3 = reshape(\n sgeo_N1.n3[1:Nfp_N1[d], f, :],\n Nq_f_N1...,\n nelem,\n )\n\n sgeo.n1[1:Nfp[d], f, :][:] .=\n sum(sM_N1 .* fld_N1_n1, dims = findall(Nq_f .== 1))[:] ./\n sM[:]\n\n sgeo.n2[1:Nfp[d], f, :][:] .=\n sum(sM_N1 .* fld_N1_n2, dims = findall(Nq_f .== 1))[:] ./\n sM[:]\n\n sgeo.n3[1:Nfp[d], f, :][:] .=\n sum(sM_N1 .* fld_N1_n3, dims = findall(Nq_f .== 1))[:] ./\n sM[:]\n\n num_sgeo_handled += 3\n\n # set the volume inverse mass matrix\n sgeo.vωJI[1:Nfp[d], f, :] .= vgeo.ωJI[fmask[f], :]\n num_sgeo_handled += 1\n\n # Make sure we handled all the vgeo terms\n @assert _nsgeo == num_sgeo_handled\n end\n end\n end\n end\n\n (vgeo, sgeo, x_vtk)\nend\n\n\"\"\"\n computegeometry(elemtocoord, D, ξ, ω, meshwarp)\n\nCompute the geometric factors data needed to define metric terms at\neach quadrature point. First, compute the so called \"topology coordinates\"\nfrom reference coordinates ξ. Then map these topology coordinate\nto physical coordinates. Then compute the Jacobian of the mapping from\nreference coordinates to physical coordinates, i.e., ∂x/∂ξ, by calling\n`compute_reference_to_physical_coord_jacobian!`.\nFinally, compute the metric terms by calling the function `computemetric!`.\n\"\"\"\nfunction computegeometry(elemtocoord, D, ξ, ω, meshwarp)\n FT = eltype(D[1])\n dim = length(D)\n nface = 2dim\n nelem = size(elemtocoord, 3)\n\n Nq = ntuple(j -> size(D[j], 1), dim)\n\n # Compute metric terms for FVM\n if any(Nq .== 1)\n return computegeometry_fvm(elemtocoord, D, ξ, ω, meshwarp)\n end\n\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n # Initialize volume and surface geometric term data structures\n vgeo = VolumeGeometry(FT, Nq, nelem)\n sgeo = SurfaceGeometry(FT, Nq, nface, nelem)\n\n # a) Compute \"topology coordinates\" from reference coordinates ξ\n Metrics.creategrid!(vgeo, elemtocoord, ξ)\n\n # Create local variables\n x1 = vgeo.x1\n x2 = vgeo.x2\n x3 = vgeo.x3\n\n # b) Map \"topology coordinates\" -> physical coordinates\n @inbounds for j in 1:length(vgeo.x1)\n (x1[j], x2[j], x3[j]) = meshwarp(vgeo.x1[j], vgeo.x2[j], vgeo.x3[j])\n end\n\n # Update global data in vgeo\n vgeo.x1 .= x1\n vgeo.x2 .= x2\n vgeo.x3 .= x3\n\n # c) Compute Jacobian matrix, ∂x/∂ξ\n Metrics.compute_reference_to_physical_coord_jacobian!(vgeo, nelem, D)\n\n # d) Compute the metric terms\n Metrics.computemetric!(vgeo, sgeo, D)\n\n # Note:\n # To get analytic derivatives, we need to be able differentiate through (a,b) and combine (a,b,c)\n\n # Compute the metric terms\n p = reshape(1:Np, Nq)\n if dim == 1\n fmask = (p[1:1], p[Nq[1]:Nq[1]])\n elseif dim == 2\n fmask = (p[1, :][:], p[Nq[1], :][:], p[:, 1][:], p[:, Nq[2]][:])\n elseif dim == 3\n fmask = (\n p[1, :, :][:],\n p[Nq[1], :, :][:],\n p[:, 1, :][:],\n p[:, Nq[2], :][:],\n p[:, :, 1][:],\n p[:, :, Nq[3]][:],\n )\n end\n\n # since `ξ1` is the fastest dimension and `ξdim` the slowest the tensor\n # product order is reversed\n M = kron(1, reverse(ω)...)\n vgeo.ωJ .*= M\n vgeo.ωJI .= 1 ./ vgeo.ωJ\n for d in 1:dim\n for f in (2d - 1):(2d)\n sgeo.vωJI[1:Nfp[d], f, :] .= vgeo.ωJI[fmask[f], :]\n end\n end\n\n sM = fill!(similar(sgeo.sωJ, maximum(Nfp), nface), NaN)\n for d in 1:dim\n for f in (2d - 1):(2d)\n ωf = ntuple(j -> ω[mod1(d + j, dim)], dim - 1)\n # Because of the `mod1` this face is already flipped\n if !(dim == 3 && d == 2)\n ωf = reverse(ωf)\n end\n sM[1:Nfp[d], f] .= dim > 1 ? kron(1, ωf...) : one(FT)\n end\n end\n sgeo.sωJ .*= sM\n\n # compute MH and JvC\n horizontal_metrics!(vgeo, Nq, ω)\n\n # This is mainly done to support FVM plotting when N=0 (since we need cell\n # edge values)\n x_vtk = (vgeo.x1, vgeo.x2, vgeo.x3)\n\n return (vgeo, sgeo, x_vtk)\nend\n\n\"\"\"\n horizontal_metrics!(vgeo::VolumeGeometry, Nq, ω)\n\nCompute the horizontal mass matrix `ωJH` field of `vgeo`\n```\nJ .* norm(∂ξ3/∂x) * (ωᵢ ⊗ ωⱼ); for integrating over a plane\n```\n(in 2-D ξ2 not ξ3 is used).\n\"\"\"\nfunction horizontal_metrics!(vgeo::VolumeGeometry, Nq, ω)\n dim = length(Nq)\n\n MH = dim == 1 ? 1 : kron(ones(1, Nq[dim]), reverse(ω[1:(dim - 1)])...)[:]\n M = vec(kron(1, reverse(ω)...))\n\n J = vgeo.ωJ ./ M\n\n # Compute |r'(ξ3)| for vertical line integrals\n if dim == 1\n vgeo.ωJH .= 1\n elseif dim == 2\n vgeo.ωJH .= MH .* hypot.(J .* vgeo.ξ2x1, J .* vgeo.ξ2x2)\n elseif dim == 3\n vgeo.ωJH .= MH .* hypot.(J .* vgeo.ξ3x1, J .* vgeo.ξ3x2, J .* vgeo.ξ3x3)\n else\n error(\"dim $dim not implemented\")\n end\n return vgeo\nend\n\n\"\"\"\n indefinite_integral_interpolation_matrix(r, ω)\n\nGiven a set of integration points `r` and integration weights `ω` this computes\na matrix that will compute the indefinite integral of the (interpolant) of a\nfunction and evaluate the indefinite integral at the points `r`.\n\nNamely, let\n```math\n q(ξ) = ∫_{ξ_{0}}^{ξ} f(ξ') dξ'\n```\nthen we have that\n```\nI∫ * f.(r) = q.(r)\n```\nwhere `I∫` is the integration and interpolation matrix defined by this function.\n\n!!! note\n\n The integration is done using the provided quadrature weight, so if these\n cannot integrate `f(ξ)` exactly, `f` is first interpolated and then\n integrated using quadrature. Namely, we have that:\n ```math\n q(ξ) = ∫_{ξ_{0}}^{ξ} I(f(ξ')) dξ'\n ```\n where `I` is the interpolation operator.\n\n\"\"\"\nfunction indefinite_integral_interpolation_matrix(r, ω)\n Nq = length(r)\n\n I∫ = similar(r, Nq, Nq)\n # first value is zero\n I∫[1, :] .= Nq == 1 ? ω[1] : 0\n\n # barycentric weights for interpolation\n wbary = Elements.baryweights(r)\n\n # Compute the interpolant of the indefinite integral\n for n in 2:Nq\n # grid from first dof to current point\n rdst = (1 .- r) / 2 * r[1] + (1 .+ r) / 2 * r[n]\n # interpolation matrix\n In = Elements.interpolationmatrix(r, rdst, wbary)\n # scaling from LGL to current of the interval\n Δ = (r[n] - r[1]) / 2\n # row of the matrix we have computed\n I∫[n, :] .= (Δ * ω' * In)[:]\n end\n I∫\nend\n# }}}\n\nusing KernelAbstractions.Extras: @unroll\n\nusing StaticArrays\n\nconst _x1 = Grids._x1\nconst _x2 = Grids._x2\nconst _x3 = Grids._x3\nconst _JcV = Grids._JcV\n\n@doc \"\"\"\n kernel_min_neighbor_distance!(::Val{N}, ::Val{dim}, direction,\n min_neighbor_distance, vgeo, topology.realelems)\n\nComputational kernel: Computes the minimum physical distance between node\nneighbors within an element.\n\nThe `direction` in the reference element controls which nodes are considered\nneighbors.\n\"\"\" kernel_min_neighbor_distance!\n@kernel function kernel_min_neighbor_distance!(\n ::Val{N},\n ::Val{dim},\n direction,\n min_neighbor_distance,\n vgeo,\n elems,\n) where {N, dim}\n\n @uniform begin\n FT = eltype(min_neighbor_distance)\n Nq = N .+ 1\n Np = prod(Nq)\n\n if direction isa EveryDirection\n mininξ = (true, true, true)\n elseif direction isa HorizontalDirection\n mininξ = (true, dim == 2 ? false : true, false)\n elseif direction isa VerticalDirection\n mininξ = (false, dim == 2 ? true : false, dim == 2 ? false : true)\n end\n\n @inbounds begin\n # 2D Nq = (nh, nv)\n # 3D Nq = (nh, nh, nv)\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nqk = dim == 2 ? 1 : Nq[end]\n mininξ1 = mininξ[1]\n mininξ2 = mininξ[2]\n mininξ3 = mininξ[3]\n end\n end\n\n\n I = @index(Global, Linear)\n # local element id\n e = (I - 1) ÷ Np + 1\n # local quadrature id\n ijk = (I - 1) % Np + 1\n\n # local i, j, k quadrature id\n i = (ijk - 1) % Nq1 + 1\n j = (ijk - 1) ÷ Nq1 % Nq2 + 1\n k = (ijk - 1) ÷ (Nq1 * Nq2) % Nqk + 1\n\n md = typemax(FT)\n\n x = SVector(vgeo[ijk, _x1, e], vgeo[ijk, _x2, e], vgeo[ijk, _x3, e])\n\n # first horizontal distance\n if mininξ1\n @unroll for î in (i - 1, i + 1)\n if 1 ≤ î ≤ Nq1\n îjk = î + Nq1 * (j - 1) + Nq1 * Nq2 * (k - 1)\n x̂ = SVector(\n vgeo[îjk, _x1, e],\n vgeo[îjk, _x2, e],\n vgeo[îjk, _x3, e],\n )\n md = min(md, norm(x - x̂))\n end\n end\n end\n\n # second horizontal distance or vertical distance (dim=2)\n if mininξ2\n # FV Vercial direction, use 2vgeo[ijk, _JcV, e]\n if dim == 2 && Nq2 == 1\n md = min(md, 2vgeo[ijk, _JcV, e])\n else\n @unroll for ĵ in (j - 1, j + 1)\n if 1 ≤ ĵ ≤ Nq2\n iĵk = i + Nq1 * (ĵ - 1) + Nq1 * Nq2 * (k - 1)\n x̂ = SVector(\n vgeo[iĵk, _x1, e],\n vgeo[iĵk, _x2, e],\n vgeo[iĵk, _x3, e],\n )\n md = min(md, norm(x - x̂))\n end\n end\n end\n end\n\n # vertical distance (dim=3)\n if mininξ3\n # FV Vercial direction, use 2vgeo[ijk, _JcV, e]\n if dim == 3 && Nqk == 1\n md = min(md, 2vgeo[ijk, _JcV, e])\n else\n @unroll for k̂ in (k - 1, k + 1)\n if 1 ≤ k̂ ≤ Nqk\n ijk̂ = i + Nq1 * (j - 1) + Nq1 * Nq2 * (k̂ - 1)\n x̂ = SVector(\n vgeo[ijk̂, _x1, e],\n vgeo[ijk̂, _x2, e],\n vgeo[ijk̂, _x3, e],\n )\n md = min(md, norm(x - x̂))\n end\n end\n end\n end\n\n min_neighbor_distance[ijk, e] = md\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/Interpolation.jl", "content": "module Interpolation\n\nusing CUDA\nusing DocStringExtensions\nusing LinearAlgebra\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing KernelAbstractions\n\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Geometry\nimport ClimateMachine.Mesh.Elements: baryweights\nimport ClimateMachine.MPIStateArrays: array_device\n\nexport dimensions,\n accumulate_interpolated_data,\n accumulate_interpolated_data!,\n InterpolationBrick,\n InterpolationCubedSphere,\n interpolate_local!,\n project_cubed_sphere!,\n InterpolationTopology\n\nabstract type InterpolationTopology end\n\ndimensions(nothing) = OrderedDict()\n\n\"\"\"\n InterpolationBrick{\n FT <: AbstractFloat,CuArrays\n UI8AD <: AbstractArray{UInt8, 2},\n UI16VD <: AbstractVector{UInt16},\n I32V <: AbstractVector{Int32},\n } <: InterpolationTopology\n\nThis interpolation data structure and the corresponding functions works for a\nbrick, where stretching/compression happens only along the x1, x2 & x3 axis.\nHere x1 = X1(ξ1), x2 = X2(ξ2) and x3 = X3(ξ3).\n\n# Fields\n\n$(DocStringExtensions.FIELDS)\n\n# Usage\n\n InterpolationBrick(\n grid::DiscontinuousSpectralElementGrid{FT},\n xbnd::Array{FT,2},\n xres,\n ) where FT <: AbstractFloat\n\nThis interpolation structure and the corresponding functions works for a brick,\nwhere stretching/compression happens only along the x1, x2 & x3 axis. Here x1\n= X1(ξ1), x2 = X2(ξ2) and x3 = X3(ξ3).\n\n# Arguments for the inner constructor\n - `grid`: DiscontinousSpectralElementGrid\n - `xbnd`: Domain boundaries in x1, x2 and x3 directions\n - `x1g`: Interpolation grid in x1 direction\n - `x2g`: Interpolation grid in x2 direction\n - `x3g`: Interpolation grid in x3 direction\n\"\"\"\nstruct InterpolationBrick{\n FT <: AbstractFloat,\n I <: Int,\n FTV <: AbstractVector{FT},\n FTVD <: AbstractVector{FT},\n IVD <: AbstractVector{I},\n FTA2 <: Array{FT, 2},\n UI8AD <: AbstractArray{UInt8, 2},\n UI16VD <: AbstractVector{UInt16},\n I32V <: AbstractVector{Int32},\n} <: InterpolationTopology\n \"Number of elements\"\n Nel::I\n \"Total number of interpolation points\"\n Np::I\n \"Total number of interpolation points on local process\"\n Npl::I\n \"Domain bounds in x1, x2 and x3 directions\"\n xbnd::FTA2\n \"Interpolation grid in x1 direction\"\n x1g::FTV\n \"Interpolation grid in x2 direction\"\n x2g::FTV\n \"Interpolation grid in x3 direction\"\n x3g::FTV\n \"Unique ξ1 coordinates of interpolation points within each spectral element\"\n ξ1::FTVD\n \"Unique ξ2 coordinates of interpolation points within each spectral element\"\n ξ2::FTVD\n \"Unique ξ3 coordinates of interpolation points within each spectral element\"\n ξ3::FTVD\n \"Flags when ξ1/ξ2/ξ3 interpolation point matches with a GLL point\"\n flg::UI8AD\n \"Normalization factor\"\n fac::FTVD\n \"x1 interpolation grid index of interpolation points within each element on the local process\"\n x1i::UI16VD\n \"x2 interpolation grid index of interpolation points within each element on the local process\"\n x2i::UI16VD\n \"x3 interpolation grid index of interpolation points within each element on the local process\"\n x3i::UI16VD\n \"Offsets for each element\"\n offset::IVD # offsets for each element for v\n \"GLL points in ξ1 direction\"\n m_ξ1::FTVD\n \"GLL points in ξ2 direction\"\n m_ξ2::FTVD\n \"GLL points in ξ3 direction\"\n m_ξ3::FTVD\n \"Barycentric weights\"\n wb1::FTVD\n \"Barycentric weights\"\n wb2::FTVD\n \"Barycentric weights\"\n wb3::FTVD\n # MPI setup for gathering interpolated variable on proc # 0\n \"Number of interpolation points on each of the processes\"\n Np_all::I32V\n\n \"x1 interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n x1i_all::UI16VD\n \"x2 interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n x2i_all::UI16VD\n \"x3 interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n x3i_all::UI16VD\n\n function InterpolationBrick(\n grid::DiscontinuousSpectralElementGrid{FT},\n xbnd::Array{FT, 2},\n x1g::AbstractArray{FT, 1},\n x2g::AbstractArray{FT, 1},\n x3g::AbstractArray{FT, 1},\n ) where {FT <: AbstractFloat}\n mpicomm = grid.topology.mpicomm\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n\n DA = arraytype(grid) # device array\n device = arraytype(grid) <: Array ? CPU() : CUDADevice()\n\n qm = polynomialorders(grid) .+ 1\n ndim = 3\n toler = 4 * eps(FT) # tolerance\n\n n1g = length(x1g)\n n2g = length(x2g)\n n3g = length(x3g)\n\n Np = n1g * n2g * n3g\n marker = BitArray{3}(undef, n1g, n2g, n3g)\n fill!(marker, true)\n\n Nel = length(grid.topology.realelems) # # of elements on local process\n offset = Vector{Int}(undef, Nel + 1) # offsets for the interpolated variable\n n123 = zeros(Int, ndim) # # of unique ξ1, ξ2, ξ3 points in each cell\n xsten = zeros(Int, 2, ndim) # x1, x2, x3 start and end for each brick element\n xbndl = zeros(FT, 2, ndim) # x1,x2,x3 limits (min,max) for each brick element\n\n ξ1 = map(i -> zeros(FT, i), zeros(Int, Nel))\n ξ2 = map(i -> zeros(FT, i), zeros(Int, Nel))\n ξ3 = map(i -> zeros(FT, i), zeros(Int, Nel))\n\n x1i = map(i -> zeros(UInt16, i), zeros(UInt16, Nel))\n x2i = map(i -> zeros(UInt16, i), zeros(UInt16, Nel))\n x3i = map(i -> zeros(UInt16, i), zeros(UInt16, Nel))\n\n x = map(i -> zeros(FT, ndim, i), zeros(Int, Nel)) # interpolation grid points embedded in each cell\n\n offset[1] = 0\n\n for el in 1:Nel\n for (xg, dim) in zip((x1g, x2g, x3g), 1:ndim)\n xbndl[1, dim], xbndl[2, dim] =\n extrema(grid.topology.elemtocoord[dim, :, el])\n\n st = findfirst(xg .≥ xbndl[1, dim] .- toler)\n if st ≠ nothing\n if xg[st] > (xbndl[2, dim] + toler)\n st = nothing\n end\n end\n\n if st ≠ nothing\n xsten[1, dim] = st\n xsten[2, dim] =\n findlast(temp -> temp .≤ xbndl[2, dim] .+ toler, xg)\n n123[dim] = xsten[2, dim] - xsten[1, dim] + 1\n else\n n123[dim] = 0\n end\n end\n\n if prod(n123) > 0\n for k in xsten[1, 3]:xsten[2, 3],\n j in xsten[1, 2]:xsten[2, 2],\n i in xsten[1, 1]:xsten[2, 1]\n\n if marker[i, j, k]\n push!(\n ξ1[el],\n 2 * (x1g[i] - xbndl[1, 1]) /\n (xbndl[2, 1] - xbndl[1, 1]) - 1,\n )\n push!(\n ξ2[el],\n 2 * (x2g[j] - xbndl[1, 2]) /\n (xbndl[2, 2] - xbndl[1, 2]) - 1,\n )\n push!(\n ξ3[el],\n 2 * (x3g[k] - xbndl[1, 3]) /\n (xbndl[2, 3] - xbndl[1, 3]) - 1,\n )\n\n push!(x1i[el], UInt16(i))\n push!(x2i[el], UInt16(j))\n push!(x3i[el], UInt16(k))\n marker[i, j, k] = false\n end\n end\n offset[el + 1] = offset[el] + length(ξ1[el])\n else\n offset[el + 1] = offset[el]\n end\n\n end # el loop\n\n m_ξ1, m_ξ2, m_ξ3 = referencepoints(grid)\n wb1, wb2, wb3 = baryweights(m_ξ1), baryweights(m_ξ2), baryweights(m_ξ3)\n\n Npl = offset[end]\n\n ξ1_d = Array{FT}(undef, Npl)\n ξ2_d = Array{FT}(undef, Npl)\n ξ3_d = Array{FT}(undef, Npl)\n x1i_d = Array{UInt16}(undef, Npl)\n x2i_d = Array{UInt16}(undef, Npl)\n x3i_d = Array{UInt16}(undef, Npl)\n fac_d = zeros(FT, Npl)\n flg_d = zeros(UInt8, 3, Npl)\n\n for i in 1:Nel\n ctr = 1\n for j in (offset[i] + 1):offset[i + 1]\n ξ1_d[j] = ξ1[i][ctr]\n ξ2_d[j] = ξ2[i][ctr]\n ξ3_d[j] = ξ3[i][ctr]\n x1i_d[j] = x1i[i][ctr]\n x2i_d[j] = x2i[i][ctr]\n x3i_d[j] = x3i[i][ctr]\n # set up interpolation\n fac1 = FT(0)\n fac2 = FT(0)\n fac3 = FT(0)\n for ib in 1:qm[1]\n if abs(m_ξ1[ib] - ξ1_d[j]) < toler\n @inbounds flg_d[1, j] = UInt8(ib)\n else\n @inbounds fac1 += wb1[ib] / (ξ1_d[j] - m_ξ1[ib])\n end\n end\n\n for ib in 1:qm[2]\n if abs(m_ξ2[ib] - ξ2_d[j]) < toler\n @inbounds flg_d[2, j] = UInt8(ib)\n else\n @inbounds fac2 += wb2[ib] / (ξ2_d[j] - m_ξ2[ib])\n end\n end\n\n for ib in 1:qm[3]\n if abs(m_ξ3[ib] - ξ3_d[j]) < toler\n @inbounds flg_d[3, j] = UInt8(ib)\n else\n @inbounds fac3 += wb3[ib] / (ξ3_d[j] - m_ξ3[ib])\n end\n end\n\n flg_d[1, j] ≠ UInt8(0) && (fac1 = FT(1))\n flg_d[2, j] ≠ UInt8(0) && (fac2 = FT(1))\n flg_d[3, j] ≠ UInt8(0) && (fac3 = FT(1))\n\n fac_d[j] = FT(1) / (fac1 * fac2 * fac3)\n\n ctr += 1\n end\n end\n # MPI setup for gathering data on proc 0\n root = 0\n Np_all = zeros(Int32, npr)\n Np_all[pid + 1] = Npl\n\n MPI.Allreduce!(Np_all, +, mpicomm)\n\n if pid ≠ root\n x1i_all = zeros(UInt16, 0)\n x2i_all = zeros(UInt16, 0)\n x3i_all = zeros(UInt16, 0)\n MPI.Gatherv!(x1i_d, nothing, root, mpicomm)\n MPI.Gatherv!(x2i_d, nothing, root, mpicomm)\n MPI.Gatherv!(x3i_d, nothing, root, mpicomm)\n else\n x1i_all = Array{UInt16}(undef, sum(Np_all))\n x2i_all = Array{UInt16}(undef, sum(Np_all))\n x3i_all = Array{UInt16}(undef, sum(Np_all))\n MPI.Gatherv!(x1i_d, VBuffer(x1i_all, Np_all), root, mpicomm)\n MPI.Gatherv!(x2i_d, VBuffer(x2i_all, Np_all), root, mpicomm)\n MPI.Gatherv!(x3i_d, VBuffer(x3i_all, Np_all), root, mpicomm)\n end\n\n\n if device isa CUDADevice\n ξ1_d = DA(ξ1_d)\n ξ2_d = DA(ξ2_d)\n ξ3_d = DA(ξ3_d)\n x1i_d = DA(x1i_d)\n x2i_d = DA(x2i_d)\n x3i_d = DA(x3i_d)\n flg_d = DA(flg_d)\n fac_d = DA(fac_d)\n offset = DA(offset)\n m_ξ1 = DA(m_ξ1)\n m_ξ2 = DA(m_ξ2)\n m_ξ3 = DA(m_ξ3)\n wb1 = DA(wb1)\n wb2 = DA(wb2)\n wb3 = DA(wb3)\n x1i_all = DA(x1i_all)\n x2i_all = DA(x2i_all)\n x3i_all = DA(x3i_all)\n end\n return new{\n FT,\n Int,\n typeof(x1g),\n typeof(ξ1_d),\n typeof(offset),\n typeof(xbnd),\n typeof(flg_d),\n typeof(x1i_d),\n typeof(Np_all),\n }(\n Nel,\n Np,\n Npl,\n xbnd,\n x1g,\n x2g,\n x3g,\n ξ1_d,\n ξ2_d,\n ξ3_d,\n flg_d,\n fac_d,\n x1i_d,\n x2i_d,\n x3i_d,\n offset,\n m_ξ1,\n m_ξ2,\n m_ξ3,\n wb1,\n wb2,\n wb3,\n Np_all,\n x1i_all,\n x2i_all,\n x3i_all,\n )\n\n end\n\nend # struct InterpolationBrick\n\n\"\"\"\n interpolate_local!(\n intrp_brck::InterpolationBrick{FT},\n sv::AbstractArray{FT},\n v::AbstractArray{FT},\n ) where {FT <: AbstractFloat}\n\nThis interpolation function works for a brick, where stretching/compression\nhappens only along the x1, x2 & x3 axis. Here x1 = X1(ξ1), x2 = X2(ξ2) and x3\n= X3(ξ3)\n\n# Arguments\n - `intrp_brck`: Initialized InterpolationBrick structure\n - `sv`: State Array consisting of various variables on the discontinuous\n Galerkin grid\n - `v`: Interpolated variables\n\"\"\"\nfunction interpolate_local!(\n intrp_brck::InterpolationBrick{FT},\n sv::AbstractArray{FT},\n v::AbstractArray{FT},\n) where {FT <: AbstractFloat}\n\n offset = intrp_brck.offset\n m_ξ1 = intrp_brck.m_ξ1\n m_ξ2 = intrp_brck.m_ξ2\n m_ξ3 = intrp_brck.m_ξ3\n wb1 = intrp_brck.wb1\n wb2 = intrp_brck.wb2\n wb3 = intrp_brck.wb3\n ξ1 = intrp_brck.ξ1\n ξ2 = intrp_brck.ξ2\n ξ3 = intrp_brck.ξ3\n flg = intrp_brck.flg\n fac = intrp_brck.fac\n\n qm = (length(m_ξ1), length(m_ξ2), length(m_ξ3))\n Nel = length(offset) - 1\n nvars = size(sv, 2)\n\n device = array_device(sv)\n comp_stream = Event(device)\n\n workgroup = (qm[2], qm[3])\n ndrange = (qm[2] * Nel, qm[3] * nvars)\n\n comp_stream = interpolate_local_kernel!(device, workgroup)(\n offset,\n m_ξ1,\n m_ξ2,\n m_ξ3,\n wb1,\n wb2,\n wb3,\n ξ1,\n ξ2,\n ξ3,\n flg,\n fac,\n sv,\n v,\n Val(qm),\n ndrange = ndrange,\n dependencies = (comp_stream,),\n )\n wait(comp_stream)\n return nothing\nend\n#---------------------------------------------------------------\n@kernel function interpolate_local_kernel!(\n offset::AbstractArray{T, 1},\n m_ξ1::AbstractArray{FT, 1},\n m_ξ2::AbstractArray{FT, 1},\n m_ξ3::AbstractArray{FT, 1},\n wb1::AbstractArray{FT, 1},\n wb2::AbstractArray{FT, 1},\n wb3::AbstractArray{FT, 1},\n ξ1::AbstractArray{FT, 1},\n ξ2::AbstractArray{FT, 1},\n ξ3::AbstractArray{FT, 1},\n flg::AbstractArray{UInt8, 2},\n fac::AbstractArray{FT, 1},\n sv::AbstractArray{FT},\n v::AbstractArray{FT},\n ::Val{qm},\n) where {qm, T <: Int, FT <: AbstractFloat}\n\n el, st_idx = @index(Group, NTuple)\n tj, tk = @index(Local, NTuple)\n\n\n vout_jk = @localmem FT (qm[2], qm[3])\n m_ξ1_sh = @localmem FT (qm[1],)\n m_ξ2_sh = @localmem FT (qm[2],)\n m_ξ3_sh = @localmem FT (qm[3],)\n wb1_sh = @localmem FT (qm[1],)\n wb2_sh = @localmem FT (qm[2],)\n wb3_sh = @localmem FT (qm[3],)\n\n np = @localmem T (1,)\n off = @localmem T (1,)\n\n # load shared memory\n if tk == 1\n m_ξ2_sh[tj] = m_ξ2[tj]\n wb2_sh[tj] = wb2[tj]\n end\n if tj == 1\n m_ξ3_sh[tk] = m_ξ3[tk]\n wb3_sh[tk] = wb3[tk]\n end\n if tj == 1 && tk == 1\n for i in 1:qm[1]\n m_ξ1_sh[i] = m_ξ1[i]\n wb1_sh[i] = wb1[i]\n end\n np[1] = offset[el + 1] - offset[el]\n off[1] = offset[el]\n end\n @synchronize\n\n\n for i in 1:np[1] # interpolate point-by-point\n\n ξ1l = ξ1[off[1] + i]\n ξ2l = ξ2[off[1] + i]\n\n f1 = flg[1, off[1] + i]\n\n if f1 == 0 # apply phir\n @inbounds vout_jk[tj, tk] =\n sv[\n 1 + (tj - 1) * qm[1] + (tk - 1) * qm[1] * qm[2],\n st_idx,\n el,\n ] * wb1_sh[1] / (ξ1l - m_ξ1_sh[1])\n for ii in 2:qm[1]\n @inbounds vout_jk[tj, tk] +=\n sv[\n ii + (tj - 1) * qm[1] + (tk - 1) * qm[1] * qm[2],\n st_idx,\n el,\n ] * wb1_sh[ii] / (ξ1l - m_ξ1_sh[ii])\n end\n else\n @inbounds vout_jk[tj, tk] =\n sv[f1 + (tj - 1) * qm[1] + (tk - 1) * qm[1] * qm[2], st_idx, el]\n end\n\n if flg[2, off[1] + i] == 0 # apply phis\n @inbounds vout_jk[tj, tk] *= (wb2_sh[tj] / (ξ2l - m_ξ2_sh[tj]))\n end\n @synchronize\n\n f2 = flg[2, off[1] + i]\n ξ3l = ξ3[off[1] + i]\n\n if tj == 1 # reduction\n if f2 == 0\n for ij in 2:qm[2]\n @inbounds vout_jk[1, tk] += vout_jk[ij, tk]\n end\n else\n if f2 ≠ 1\n @inbounds vout_jk[1, tk] = vout_jk[f2, tk]\n end\n end\n\n if flg[3, off[1] + i] == 0 # apply phit\n @inbounds vout_jk[1, tk] *= (wb3_sh[tk] / (ξ3l - m_ξ3_sh[tk]))\n end\n end\n @synchronize\n\n f3 = flg[3, off[1] + i]\n\n if tj == 1 && tk == 1 # reduction\n if f3 == 0\n for ik in 2:qm[3]\n @inbounds vout_jk[1, 1] += vout_jk[1, ik]\n end\n else\n if f3 ≠ 1\n @inbounds vout_jk[1, 1] = vout_jk[1, f3]\n end\n end\n @inbounds v[off[1] + i, st_idx] = vout_jk[1, 1] * fac[off[1] + i]\n end\n @synchronize\n end\nend\n#---------------------------------------------------------------\nfunction dimensions(interpol::InterpolationBrick)\n if Array ∈ typeof(interpol.x1g).parameters\n h_x1g = interpol.x1g\n h_x2g = interpol.x2g\n h_x3g = interpol.x3g\n else\n h_x1g = Array(interpol.x1g)\n h_x2g = Array(interpol.x2g)\n h_x3g = Array(interpol.x3g)\n end\n return OrderedDict(\n \"x\" => (h_x1g, OrderedDict()),\n \"y\" => (h_x2g, OrderedDict()),\n \"z\" => (h_x3g, OrderedDict()),\n )\nend\n\n\"\"\"\n InterpolationCubedSphere{\n FT <: AbstractFloat,\n T <: Int,\n FTV <: AbstractVector{FT},\n FTVD <: AbstractVector{FT},\n TVD <: AbstractVector{T},\n UI8AD <: AbstractArray{UInt8, 2},\n UI16VD <: AbstractVector{UInt16},\n I32V <: AbstractVector{Int32},\n } <: InterpolationTopology\n\nThis interpolation structure and the corresponding functions works for a cubed sphere topology. The data is interpolated along a lat/long/rad grid.\n\n-90⁰ ≤ lat ≤ 90⁰\n\n-180⁰ ≤ long ≤ 180⁰\n\nRᵢ ≤ r ≤ Rₒ\n\n# Fields\n\n$(DocStringExtensions.FIELDS)\n\n# Usage\n\n InterpolationCubedSphere(grid::DiscontinuousSpectralElementGrid, vert_range::AbstractArray{FT}, nhor::Int, lat_res::FT, long_res::FT, rad_res::FT) where {FT <: AbstractFloat}\n\nThis interpolation structure and the corresponding functions works for a cubed sphere topology. The data is interpolated along a lat/long/rad grid.\n\n-90⁰ ≤ lat ≤ 90⁰\n\n-180⁰ ≤ long ≤ 180⁰\n\nRᵢ ≤ r ≤ Rₒ\n\n# Arguments for the inner constructor\n - `grid`: DiscontinousSpectralElementGrid\n - `vert_range`: Vertex range along the radial coordinate\n - `lat_res`: Resolution of the interpolation grid along the latitude coordinate in radians\n - `long_res`: Resolution of the interpolation grid along the longitude coordinate in radians\n - `rad_res`: Resolution of the interpolation grid along the radial coordinate\n\"\"\"\nstruct InterpolationCubedSphere{\n FT <: AbstractFloat,\n T <: Int,\n FTV <: AbstractVector{FT},\n FTVD <: AbstractVector{FT},\n TVD <: AbstractVector{T},\n UI8AD <: AbstractArray{UInt8, 2},\n UI16VD <: AbstractVector{UInt16},\n I32V <: AbstractVector{Int32},\n} <: InterpolationTopology\n \"Number of elements\"\n Nel::T\n \"Number of interpolation points\"\n Np::T\n \"Number of interpolation points on local process\"\n Npl::T # # of interpolation points on the local process\n \"Number of interpolation points in radial direction\"\n n_rad::T\n \"Number of interpolation points in lat direction\"\n n_lat::T\n \"Number of interpolation points in long direction\"\n n_long::T\n \"Interpolation grid in radial direction\"\n rad_grd::FTV\n \"Interpolation grid in lat direction\"\n lat_grd::FTV\n \"Interpolation grid in long direction\"\n long_grd::FTV # rad, lat & long locations of interpolation grid\n \"Device array containing ξ1 coordinates of interpolation points within each element\"\n ξ1::FTVD\n \"Device array containing ξ2 coordinates of interpolation points within each element\"\n ξ2::FTVD\n \"Device array containing ξ3 coordinates of interpolation points within each element\"\n ξ3::FTVD\n \"flags when ξ1/ξ2/ξ3 interpolation point matches with a GLL point\"\n flg::UI8AD\n \"Normalization factor\"\n fac::FTVD\n \"Radial coordinates of interpolation points withing each element\"\n radi::UI16VD\n \"Latitude coordinates of interpolation points withing each element\"\n lati::UI16VD\n \"Longitude coordinates of interpolation points withing each element\"\n longi::UI16VD\n \"Offsets for each element\"\n offset::TVD\n \"GLL points in ξ1 direction\"\n m_ξ1::FTVD\n \"GLL points in ξ2 direction\"\n m_ξ2::FTVD\n \"GLL points in ξ3 direction\"\n m_ξ3::FTVD\n \"Barycentric weights in ξ1 direction\"\n wb1::FTVD\n \"Barycentric weights in ξ2 direction\"\n wb2::FTVD\n \"Barycentric weights in ξ3 direction\"\n wb3::FTVD\n # MPI setup for gathering interpolated variable on proc 0\n \"Number of interpolation points on each of the processes\"\n Np_all::I32V\n \"Radial interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n radi_all::UI16VD\n \"Latitude interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n lati_all::UI16VD\n \"Longitude interpolation grid index of interpolation points within each element on all processes stored only on proc 0\"\n longi_all::UI16VD\n\n function InterpolationCubedSphere(\n grid::DiscontinuousSpectralElementGrid,\n vert_range::AbstractArray{FT},\n nhor::Int,\n lat_grd::AbstractArray{FT, 1},\n long_grd::AbstractArray{FT, 1},\n rad_grd::AbstractArray{FT};\n nr_toler = nothing,\n ) where {FT <: AbstractFloat}\n mpicomm = MPI.COMM_WORLD\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n\n DA = arraytype(grid) # device array\n device = arraytype(grid) <: Array ? CPU() : CUDADevice()\n\n qm = polynomialorders(grid) .+ 1\n toler1 = FT(eps(FT) * vert_range[1] * 2.0) # tolerance for unwarp function\n toler2 = FT(eps(FT) * 4.0) # tolerance\n # tolerance for Newton-Raphson\n if isnothing(nr_toler)\n nr_toler = FT(eps(FT) * vert_range[1] * 10.0)\n end\n\n Nel = length(grid.topology.realelems) # # of local elements on the local process\n\n nvert_range = length(vert_range)\n nvert = nvert_range - 1 # # of elements in vertical direction\n Nel_glob = nvert * nhor * nhor * 6\n\n nblck = nhor * nhor * nvert\n Δh = 2 / nhor # horizontal grid spacing in unwarped grid\n\n n_lat, n_long, n_rad =\n Int(length(lat_grd)), Int(length(long_grd)), Int(length(rad_grd))\n\n Np = n_lat * n_long * n_rad\n\n uw_grd = zeros(FT, 3)\n diffv = zeros(FT, 3)\n ξ = zeros(FT, 3)\n\n glob_ord = grid.topology.origsendorder # to account for reordering of elements after the partitioning process\n\n glob_elem_no = zeros(Int, nvert * length(glob_ord))\n\n for i in 1:length(glob_ord), j in 1:nvert\n glob_elem_no[j + (i - 1) * nvert] = (glob_ord[i] - 1) * nvert + j\n end\n glob_to_loc = Dict(glob_elem_no[i] => Int(i) for i in 1:Nel) # using dictionary for speedup\n\n ξ1, ξ2, ξ3 = map(i -> zeros(FT, i), zeros(Int, Nel)),\n map(i -> zeros(FT, i), zeros(Int, Nel)),\n map(i -> zeros(FT, i), zeros(Int, Nel))\n\n radi, lati, longi = map(i -> zeros(UInt16, i), zeros(UInt16, Nel)),\n map(i -> zeros(UInt16, i), zeros(UInt16, Nel)),\n map(i -> zeros(UInt16, i), zeros(UInt16, Nel))\n\n\n offset_d = zeros(Int, Nel + 1)\n\n for i in 1:n_rad\n rad = rad_grd[i]\n if rad ≤ vert_range[1] # accounting for minor rounding errors from unwarp function at boundaries\n vert_range[1] - rad < toler1 ? l_nrm = 1 :\n error(\n \"fatal error, rad lower than inner radius: \",\n vert_range[1] - rad,\n \" $rad_grd /// $lat_grd //// $long_grd\",\n )\n elseif rad ≥ vert_range[end] # accounting for minor rounding errors from unwarp function at boundaries\n rad - vert_range[end] < toler1 ? l_nrm = nvert :\n error(\"fatal error, rad greater than outer radius\")\n else # normal scenario\n for l in 2:nvert_range\n if vert_range[l] - rad > FT(0)\n l_nrm = l - 1\n break\n end\n end\n end\n\n for j in 1:n_lat\n @inbounds x3_grd = rad * sind(lat_grd[j])\n for k in 1:n_long\n @inbounds x1_grd =\n rad * cosd(lat_grd[j]) * cosd(long_grd[k]) # inclination -> latitude; azimuthal -> longitude.\n @inbounds x2_grd =\n rad * cosd(lat_grd[j]) * sind(long_grd[k]) # inclination -> latitude; azimuthal -> longitude.\n\n uw_grd[1], uw_grd[2], uw_grd[3] =\n Topologies.cubed_sphere_unwarp(\n EquiangularCubedSphere(),\n x1_grd,\n x2_grd,\n x3_grd,\n ) # unwarping from sphere to cubed shell\n\n x1_uw2_grd = uw_grd[1] / rad # unwrapping cubed shell on to a 2D grid (in 3D space, -1 to 1 cube)\n x2_uw2_grd = uw_grd[2] / rad\n x3_uw2_grd = uw_grd[3] / rad\n\n if abs(x1_uw2_grd + 1) < toler2 # face 1 (x1 == -1 plane)\n l2 = min(div(x2_uw2_grd + 1, Δh) + 1, nhor)\n l3 = min(div(x3_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (nhor - l2) * nvert +\n (l3 - 1) * nvert * nhor,\n )\n elseif abs(x2_uw2_grd + 1) < toler2 # face 2 (x2 == -1 plane)\n l1 = min(div(x1_uw2_grd + 1, Δh) + 1, nhor)\n l3 = min(div(x3_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (l1 - 1) * nvert +\n (l3 - 1) * nvert * nhor +\n nblck * 1,\n )\n elseif abs(x1_uw2_grd - 1) < toler2 # face 3 (x1 == +1 plane)\n l2 = min(div(x2_uw2_grd + 1, Δh) + 1, nhor)\n l3 = min(div(x3_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (l2 - 1) * nvert +\n (l3 - 1) * nvert * nhor +\n nblck * 2,\n )\n elseif abs(x3_uw2_grd - 1) < toler2 # face 4 (x3 == +1 plane)\n l1 = min(div(x1_uw2_grd + 1, Δh) + 1, nhor)\n l2 = min(div(x2_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (l1 - 1) * nvert +\n (l2 - 1) * nvert * nhor +\n nblck * 3,\n )\n elseif abs(x2_uw2_grd - 1) < toler2 # face 5 (x2 == +1 plane)\n l1 = min(div(x1_uw2_grd + 1, Δh) + 1, nhor)\n l3 = min(div(x3_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (l1 - 1) * nvert +\n (nhor - l3) * nvert * nhor +\n nblck * 4,\n )\n elseif abs(x3_uw2_grd + 1) < toler2 # face 6 (x3 == -1 plane)\n l1 = min(div(x1_uw2_grd + 1, Δh) + 1, nhor)\n l2 = min(div(x2_uw2_grd + 1, Δh) + 1, nhor)\n el_glob = Int(\n l_nrm +\n (l1 - 1) * nvert +\n (nhor - l2) * nvert * nhor +\n nblck * 5,\n )\n else\n error(\"error: unwrapped grid does not lie on any of the 6 faces\")\n end\n\n el_loc = get(glob_to_loc, el_glob, nothing)\n if el_loc ≠ nothing # computing inner coordinates for local elements\n invert_trilear_mapping_hex!(\n view(grid.topology.elemtocoord, 1, :, el_loc),\n view(grid.topology.elemtocoord, 2, :, el_loc),\n view(grid.topology.elemtocoord, 3, :, el_loc),\n uw_grd,\n diffv,\n nr_toler,\n ξ,\n )\n push!(ξ1[el_loc], ξ[1])\n push!(ξ2[el_loc], ξ[2])\n push!(ξ3[el_loc], ξ[3])\n push!(radi[el_loc], UInt16(i))\n push!(lati[el_loc], UInt16(j))\n push!(longi[el_loc], UInt16(k))\n offset_d[el_loc + 1] += 1\n end\n end\n end\n end\n\n for i in 2:(Nel + 1)\n @inbounds offset_d[i] += offset_d[i - 1]\n end\n\n Npl = offset_d[Nel + 1]\n\n v = Vector{FT}(undef, offset_d[Nel + 1]) # Allocating storage for interpolation variable\n\n ξ1_d = Vector{FT}(undef, Npl)\n ξ2_d = Vector{FT}(undef, Npl)\n ξ3_d = Vector{FT}(undef, Npl)\n\n flg_d = zeros(UInt8, 3, Npl)\n fac_d = ones(FT, Npl)\n\n rad_d = Vector{UInt16}(undef, Npl)\n lat_d = Vector{UInt16}(undef, Npl)\n long_d = Vector{UInt16}(undef, Npl)\n\n m_ξ1, m_ξ2, m_ξ3 = referencepoints(grid)\n wb1, wb2, wb3 = baryweights(m_ξ1), baryweights(m_ξ2), baryweights(m_ξ3)\n for i in 1:Nel\n ctr = 1\n for j in (offset_d[i] + 1):offset_d[i + 1]\n @inbounds ξ1_d[j] = ξ1[i][ctr]\n @inbounds ξ2_d[j] = ξ2[i][ctr]\n @inbounds ξ3_d[j] = ξ3[i][ctr]\n @inbounds rad_d[j] = radi[i][ctr]\n @inbounds lat_d[j] = lati[i][ctr]\n @inbounds long_d[j] = longi[i][ctr]\n # set up interpolation\n fac1 = FT(0)\n fac2 = FT(0)\n fac3 = FT(0)\n for ib in 1:qm[1]\n if abs(m_ξ1[ib] - ξ1_d[j]) < toler2\n @inbounds flg_d[1, j] = UInt8(ib)\n else\n @inbounds fac1 += wb1[ib] / (ξ1_d[j] - m_ξ1[ib])\n end\n end\n\n for ib in 1:qm[2]\n if abs(m_ξ2[ib] - ξ2_d[j]) < toler2\n @inbounds flg_d[2, j] = UInt8(ib)\n else\n @inbounds fac2 += wb2[ib] / (ξ2_d[j] - m_ξ2[ib])\n end\n end\n\n for ib in 1:qm[3]\n if abs(m_ξ3[ib] - ξ3_d[j]) < toler2\n @inbounds flg_d[3, j] = UInt8(ib)\n else\n @inbounds fac3 += wb3[ib] / (ξ3_d[j] - m_ξ3[ib])\n end\n end\n\n flg_d[1, j] ≠ 0 && (fac1 = FT(1))\n flg_d[2, j] ≠ 0 && (fac2 = FT(1))\n flg_d[3, j] ≠ 0 && (fac3 = FT(1))\n\n fac_d[j] = FT(1) / (fac1 * fac2 * fac3)\n\n ctr += 1\n end\n end\n # MPI setup for gathering data on proc 0\n root = 0\n Np_all = zeros(Int32, npr)\n Np_all[pid + 1] = Int32(Npl)\n\n MPI.Allreduce!(Np_all, +, mpicomm)\n\n if pid ≠ root\n radi_all = zeros(UInt16, 0)\n lati_all = zeros(UInt16, 0)\n longi_all = zeros(UInt16, 0)\n MPI.Gatherv!(rad_d, nothing, root, mpicomm)\n MPI.Gatherv!(lat_d, nothing, root, mpicomm)\n MPI.Gatherv!(long_d, nothing, root, mpicomm)\n else\n radi_all = Array{UInt16}(undef, sum(Np_all))\n lati_all = Array{UInt16}(undef, sum(Np_all))\n longi_all = Array{UInt16}(undef, sum(Np_all))\n MPI.Gatherv!(rad_d, VBuffer(radi_all, Np_all), root, mpicomm)\n MPI.Gatherv!(lat_d, VBuffer(lati_all, Np_all), root, mpicomm)\n MPI.Gatherv!(long_d, VBuffer(longi_all, Np_all), root, mpicomm)\n end\n\n\n if device isa CUDADevice\n ξ1_d = DA(ξ1_d)\n ξ2_d = DA(ξ2_d)\n ξ3_d = DA(ξ3_d)\n\n flg_d = DA(flg_d)\n fac_d = DA(fac_d)\n\n rad_d = DA(rad_d)\n lat_d = DA(lat_d)\n long_d = DA(long_d)\n\n m_ξ1 = DA(m_ξ1)\n m_ξ2 = DA(m_ξ2)\n m_ξ3 = DA(m_ξ3)\n wb1 = DA(wb1)\n wb2 = DA(wb2)\n wb3 = DA(wb3)\n\n offset_d = DA(offset_d)\n\n rad_grd = DA(rad_grd)\n lat_grd = DA(lat_grd)\n long_grd = DA(long_grd)\n\n radi_all = DA(radi_all)\n lati_all = DA(lati_all)\n longi_all = DA(longi_all)\n end\n\n return new{\n FT,\n Int,\n typeof(rad_grd),\n typeof(ξ1_d),\n typeof(offset_d),\n typeof(flg_d),\n typeof(rad_d),\n typeof(Np_all),\n }(\n Nel,\n Np,\n Npl,\n n_rad,\n n_lat,\n n_long,\n rad_grd,\n lat_grd,\n long_grd,\n ξ1_d,\n ξ2_d,\n ξ3_d,\n flg_d,\n fac_d,\n rad_d,\n lat_d,\n long_d,\n offset_d,\n m_ξ1,\n m_ξ2,\n m_ξ3,\n wb1,\n wb2,\n wb3,\n Np_all,\n radi_all,\n lati_all,\n longi_all,\n )\n\n end # Inner constructor InterpolationCubedSphere\n\nend # struct InterpolationCubedSphere\n\n\"\"\"\n invert_trilear_mapping_hex!(X1::AbstractArray{FT,1},\n X2::AbstractArray{FT,1},\n X3::AbstractArray{FT,1},\n x::AbstractArray{FT,1},\n d::AbstractArray{FT,1},\n tol::FT,\n ξ::AbstractArray{FT,1}) where FT <: AbstractFloat\n\nThis function computes ξ = (ξ1,ξ2,ξ3) given x = (x1,x2,x3) and the (8) vertex coordinates of a Hexahedron. Newton-Raphson method is used.\n\n# Arguments\n - `X1`: X1 coordinates of the (8) vertices of the hexahedron\n - `X2`: X2 coordinates of the (8) vertices of the hexahedron\n - `X3`: X3 coordinates of the (8) vertices of the hexahedron\n - `x`: (x1,x2,x3) coordinates of the point\n - `d`: (x1,x2,x3) coordinates, temporary storage\n - `tol`: Tolerance for convergence\n - `ξ`: (ξ1,ξ2,ξ3) coordinates of the point\n\"\"\"\nfunction invert_trilear_mapping_hex!(\n X1::AbstractArray{FT, 1},\n X2::AbstractArray{FT, 1},\n X3::AbstractArray{FT, 1},\n x::AbstractArray{FT, 1},\n d::AbstractArray{FT, 1},\n tol::FT,\n ξ::AbstractArray{FT, 1},\n) where {FT <: AbstractFloat}\n max_it = 10 # maximum # of iterations\n ξ .= FT(0) #zeros(FT,3,1) # initial guess => cell centroid\n trilinear_map_minus_x!(ξ, X1, X2, X3, x, d)\n err = sqrt(d[1] * d[1] + d[2] * d[2] + d[3] * d[3])\n ctr = 0\n # Newton-Raphson iterations\n while err > tol\n trilinear_map_IJac_x_vec!(ξ, X1, X2, X3, d)\n ξ .-= d\n trilinear_map_minus_x!(ξ, X1, X2, X3, x, d)\n err = sqrt(d[1] * d[1] + d[2] * d[2] + d[3] * d[3]) #norm(d)\n ctr += 1\n if ctr > max_it\n error(\n \"invert_trilinear_mapping_hex: Newton-Raphson not converging to desired tolerance after max_it = \",\n max_it,\n \" iterations; err = \",\n err,\n \"; toler = \",\n tol,\n )\n end\n end\n\n clamp!(ξ, FT(-1), FT(1))\n return nothing\nend\n\nfunction trilinear_map_minus_x!(\n ξ::AbstractArray{FT, 1},\n x1v::AbstractArray{FT, 1},\n x2v::AbstractArray{FT, 1},\n x3v::AbstractArray{FT, 1},\n x::AbstractArray{FT, 1},\n d::AbstractArray{FT, 1},\n) where {FT <: AbstractFloat}\n p1 = 1 + ξ[1]\n p2 = 1 + ξ[2]\n p3 = 1 + ξ[3]\n m1 = 1 - ξ[1]\n m2 = 1 - ξ[2]\n m3 = 1 - ξ[3]\n\n\n d[1] =\n (\n m1 * (\n m2 * (m3 * x1v[1] + p3 * x1v[5]) +\n p2 * (m3 * x1v[3] + p3 * x1v[7])\n ) +\n p1 * (\n m2 * (m3 * x1v[2] + p3 * x1v[6]) +\n p2 * (m3 * x1v[4] + p3 * x1v[8])\n )\n ) / 8.0 - x[1]\n\n d[2] =\n (\n m1 * (\n m2 * (m3 * x2v[1] + p3 * x2v[5]) +\n p2 * (m3 * x2v[3] + p3 * x2v[7])\n ) +\n p1 * (\n m2 * (m3 * x2v[2] + p3 * x2v[6]) +\n p2 * (m3 * x2v[4] + p3 * x2v[8])\n )\n ) / 8.0 - x[2]\n\n d[3] =\n (\n m1 * (\n m2 * (m3 * x3v[1] + p3 * x3v[5]) +\n p2 * (m3 * x3v[3] + p3 * x3v[7])\n ) +\n p1 * (\n m2 * (m3 * x3v[2] + p3 * x3v[6]) +\n p2 * (m3 * x3v[4] + p3 * x3v[8])\n )\n ) / 8.0 - x[3]\n\n return nothing\nend\n\nfunction trilinear_map_IJac_x_vec!(\n ξ::AbstractArray{FT, 1},\n x1v::AbstractArray{FT, 1},\n x2v::AbstractArray{FT, 1},\n x3v::AbstractArray{FT, 1},\n v::AbstractArray{FT, 1},\n) where {FT <: AbstractFloat}\n p1 = 1 + ξ[1]\n p2 = 1 + ξ[2]\n p3 = 1 + ξ[3]\n m1 = 1 - ξ[1]\n m2 = 1 - ξ[2]\n m3 = 1 - ξ[3]\n\n Jac11 =\n (\n m2 * (m3 * (x1v[2] - x1v[1]) + p3 * (x1v[6] - x1v[5])) +\n p2 * (m3 * (x1v[4] - x1v[3]) + p3 * (x1v[8] - x1v[7]))\n ) / 8.0\n\n Jac12 =\n (\n m1 * (m3 * (x1v[3] - x1v[1]) + p3 * (x1v[7] - x1v[5])) +\n p1 * (m3 * (x1v[4] - x1v[2]) + p3 * (x1v[8] - x1v[6]))\n ) / 8.0\n\n Jac13 =\n (\n m1 * (m2 * (x1v[5] - x1v[1]) + p2 * (x1v[7] - x1v[3])) +\n p1 * (m2 * (x1v[6] - x1v[2]) + p2 * (x1v[8] - x1v[4]))\n ) / 8.0\n\n Jac21 =\n (\n m2 * (m3 * (x2v[2] - x2v[1]) + p3 * (x2v[6] - x2v[5])) +\n p2 * (m3 * (x2v[4] - x2v[3]) + p3 * (x2v[8] - x2v[7]))\n ) / 8.0\n\n Jac22 =\n (\n m1 * (m3 * (x2v[3] - x2v[1]) + p3 * (x2v[7] - x2v[5])) +\n p1 * (m3 * (x2v[4] - x2v[2]) + p3 * (x2v[8] - x2v[6]))\n ) / 8.0\n\n Jac23 =\n (\n m1 * (m2 * (x2v[5] - x2v[1]) + p2 * (x2v[7] - x2v[3])) +\n p1 * (m2 * (x2v[6] - x2v[2]) + p2 * (x2v[8] - x2v[4]))\n ) / 8.0\n\n Jac31 =\n (\n m2 * (m3 * (x3v[2] - x3v[1]) + p3 * (x3v[6] - x3v[5])) +\n p2 * (m3 * (x3v[4] - x3v[3]) + p3 * (x3v[8] - x3v[7]))\n ) / 8.0\n\n Jac32 =\n (\n m1 * (m3 * (x3v[3] - x3v[1]) + p3 * (x3v[7] - x3v[5])) +\n p1 * (m3 * (x3v[4] - x3v[2]) + p3 * (x3v[8] - x3v[6]))\n ) / 8.0\n\n Jac33 =\n (\n m1 * (m2 * (x3v[5] - x3v[1]) + p2 * (x3v[7] - x3v[3])) +\n p1 * (m2 * (x3v[6] - x3v[2]) + p2 * (x3v[8] - x3v[4]))\n ) / 8.0\n\n # computing cofactor matrix\n C11 = Jac22 * Jac33 - Jac23 * Jac32\n C12 = -Jac21 * Jac33 + Jac23 * Jac31\n C13 = Jac21 * Jac32 - Jac22 * Jac31\n C21 = -Jac12 * Jac33 + Jac13 * Jac32\n C22 = Jac11 * Jac33 - Jac13 * Jac31\n C23 = -Jac11 * Jac32 + Jac12 * Jac31\n C31 = Jac12 * Jac23 - Jac13 * Jac22\n C32 = -Jac11 * Jac23 + Jac13 * Jac21\n C33 = Jac11 * Jac22 - Jac12 * Jac21\n\n # computing determinant\n det = Jac11 * C11 + Jac12 * C12 + Jac13 * C13\n\n Jac11 = (C11 * v[1] + C21 * v[2] + C31 * v[3]) / det\n Jac21 = (C12 * v[1] + C22 * v[2] + C32 * v[3]) / det\n Jac31 = (C13 * v[1] + C23 * v[2] + C33 * v[3]) / det\n\n v[1] = Jac11\n v[2] = Jac21\n v[3] = Jac31\n\n return nothing\nend\n\n\"\"\"\n interpolate_local!(intrp_cs::InterpolationCubedSphere{FT},\n sv::AbstractArray{FT},\n v::AbstractArray{FT}) where {FT <: AbstractFloat}\n\nThis interpolation function works for cubed spherical shell geometry.\n\n# Arguments\n - `intrp_cs`: Initialized cubed sphere structure\n - `sv`: Array consisting of various variables on the discontinuous Galerkin grid\n - `v`: Array consisting of variables on the interpolated grid\n\"\"\"\nfunction interpolate_local!(\n intrp_cs::InterpolationCubedSphere{FT},\n sv::AbstractArray{FT},\n v::AbstractArray{FT},\n) where {FT <: AbstractFloat}\n\n offset = intrp_cs.offset\n m_ξ1 = intrp_cs.m_ξ1\n m_ξ2 = intrp_cs.m_ξ2\n m_ξ3 = intrp_cs.m_ξ3\n wb1 = intrp_cs.wb1\n wb2 = intrp_cs.wb2\n wb3 = intrp_cs.wb3\n ξ1 = intrp_cs.ξ1\n ξ2 = intrp_cs.ξ2\n ξ3 = intrp_cs.ξ3\n flg = intrp_cs.flg\n fac = intrp_cs.fac\n\n qm = (length(m_ξ1), length(m_ξ2), length(m_ξ3))\n nvars = size(sv, 2)\n Nel = length(offset) - 1\n np_tot = size(v, 1)\n\n device = array_device(sv)\n comp_stream = Event(device)\n\n workgroup = (qm[2], qm[3])\n ndrange = (qm[2] * Nel, qm[3] * nvars)\n\n comp_stream = interpolate_local_kernel!(device, workgroup)(\n offset,\n m_ξ1,\n m_ξ2,\n m_ξ3,\n wb1,\n wb2,\n wb3,\n ξ1,\n ξ2,\n ξ3,\n flg,\n fac,\n sv,\n v,\n Val(qm),\n ndrange = ndrange,\n dependencies = (comp_stream,),\n )\n wait(comp_stream)\n\n return nothing\nend\n\n\"\"\"\n project_cubed_sphere!(intrp_cs::InterpolationCubedSphere{FT},\n v::AbstractArray{FT},\n uvwi::Tuple{Int,Int,Int}) where {FT <: AbstractFloat}\n\nThis function projects the velocity field along unit vectors in radial, lat and long directions for cubed spherical shell geometry.\n\n# Fields\n - `intrp_cs`: Initialized cubed sphere structure\n - `v`: Array consisting of x1, x2 and x3 components of the vector field\n - `uvwi`: Tuple providing the column numbers for x1, x2 and x3 components of vector field in the array.\n These columns will be replaced with projected vector fields along unit vectors in long, lat and rad directions.\n\"\"\"\nfunction project_cubed_sphere!(\n intrp_cs::InterpolationCubedSphere{FT},\n v::AbstractArray{FT},\n uvwi::Tuple{Int, Int, Int},\n) where {FT <: AbstractFloat}\n # projecting velocity onto unit vectors in long, lat and rad directions\n # assumes u, v and w are located in columns specified in vector uvwi\n @assert length(uvwi) == 3 \"length(uvwi) is not 3\"\n lati = intrp_cs.lati\n longi = intrp_cs.longi\n lat_grd = intrp_cs.lat_grd\n long_grd = intrp_cs.long_grd\n _ρu = uvwi[1]\n _ρv = uvwi[2]\n _ρw = uvwi[3]\n np_tot = size(v, 1)\n\n device = array_device(v)\n comp_stream = Event(device)\n\n thr_x = min(256, np_tot)\n\n workgroup = (thr_x,)\n ndrange = (np_tot,)\n\n comp_stream = project_cubed_sphere_kernel!(device, workgroup)(\n lat_grd,\n long_grd,\n lati,\n longi,\n v,\n _ρu,\n _ρv,\n _ρw,\n ndrange = ndrange,\n dependencies = (comp_stream,),\n )\n wait(comp_stream)\n return nothing\nend\n\n@kernel function project_cubed_sphere_kernel!(\n lat_grd::AbstractArray{FT, 1},\n long_grd::AbstractArray{FT, 1},\n lati::AbstractVector{UInt16},\n longi::AbstractVector{UInt16},\n v::AbstractArray{FT},\n _ρu::Int,\n _ρv::Int,\n _ρw::Int,\n) where {FT <: AbstractFloat}\n\n idx = @index(Global, Linear) # global thread ids\n # projecting velocity onto unit vectors in long, lat and rad directions\n # assumed u, v and w are located in columns 2, 3 and 4\n deg2rad = FT(π) / FT(180)\n vrad =\n v[idx, _ρu] *\n cos(lat_grd[lati[idx]] * deg2rad) *\n cos(long_grd[longi[idx]] * deg2rad) +\n v[idx, _ρv] *\n cos(lat_grd[lati[idx]] * deg2rad) *\n sin(long_grd[longi[idx]] * deg2rad) +\n v[idx, _ρw] * sin(lat_grd[lati[idx]] * deg2rad)\n\n vlat =\n -v[idx, _ρu] *\n sin(lat_grd[lati[idx]] * deg2rad) *\n cos(long_grd[longi[idx]] * deg2rad) -\n v[idx, _ρv] *\n sin(lat_grd[lati[idx]] * deg2rad) *\n sin(long_grd[longi[idx]] * deg2rad) +\n v[idx, _ρw] * cos(lat_grd[lati[idx]] * deg2rad)\n\n vlon =\n -v[idx, _ρu] * sin(long_grd[longi[idx]] * deg2rad) +\n v[idx, _ρv] * cos(long_grd[longi[idx]] * deg2rad)\n\n v[idx, _ρu] = vlon\n v[idx, _ρv] = vlat\n v[idx, _ρw] = vrad\n # TODO: cosd / sind having issues on GPU. Unable to isolate the issue at this point. Needs to be revisited.\nend\n\nfunction dimensions(interpol::InterpolationCubedSphere)\n if Array ∈ typeof(interpol.rad_grd).parameters\n h_long_grd = interpol.long_grd\n h_lat_grd = interpol.lat_grd\n h_rad_grd = interpol.rad_grd\n else\n h_long_grd = Array(interpol.long_grd)\n h_lat_grd = Array(interpol.lat_grd)\n h_rad_grd = Array(interpol.rad_grd)\n end\n FT = eltype(h_rad_grd)\n return OrderedDict(\n \"long\" => (\n h_long_grd,\n OrderedDict(\"units\" => \"degrees_east\", \"long_name\" => \"longitude\"),\n ),\n \"lat\" => (\n h_lat_grd,\n OrderedDict(\"units\" => \"degrees_north\", \"long_name\" => \"latitude\"),\n ),\n \"level\" =>\n (h_rad_grd, OrderedDict(\"units\" => \"m\", \"long_name\" => \"level\")),\n )\nend\n\n\"\"\"\n accumulate_interpolated_data!(intrp::InterpolationTopology,\n iv::AbstractArray{FT,2},\n fiv::AbstractArray{FT,4}) where {FT <: AbstractFloat}\n\nThis interpolation function gathers interpolated data onto process # 0.\n\n# Fields\n - `intrp`: Initialized interpolation topology structure\n - `iv`: Interpolated variables on local process\n - `fiv`: Full interpolated variables accumulated on process # 0\n\"\"\"\nfunction accumulate_interpolated_data!(\n intrp::InterpolationTopology,\n iv::AbstractArray{FT, 2},\n fiv::AbstractArray{FT, 4},\n) where {FT <: AbstractFloat}\n\n device = array_device(iv)\n mpicomm = MPI.COMM_WORLD\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n root = 0\n nvars = size(iv, 2)\n\n if intrp isa InterpolationCubedSphere\n nx1 = length(intrp.long_grd)\n nx2 = length(intrp.lat_grd)\n nx3 = length(intrp.rad_grd)\n np_tot = length(intrp.radi_all)\n i1 = intrp.longi_all\n i2 = intrp.lati_all\n i3 = intrp.radi_all\n elseif intrp isa InterpolationBrick\n nx1 = length(intrp.x1g)\n nx2 = length(intrp.x2g)\n nx3 = length(intrp.x3g)\n np_tot = length(intrp.x1i_all)\n i1 = intrp.x1i_all\n i2 = intrp.x2i_all\n i3 = intrp.x3i_all\n else\n error(\"Unsupported topology; only InterpolationCubedSphere and InterpolationBrick supported\")\n end\n\n if pid == 0 && size(fiv) ≠ (nx1, nx2, nx3, nvars)\n error(\"size of fiv = $(size(fiv)); which does not match with ($nx1,$nx2,$nx3,$nvars) \")\n end\n\n if npr > 1\n Np_all = intrp.Np_all\n pid == 0 ? v_all = Array{FT}(undef, np_tot, nvars) :\n v_all = Array{FT}(undef, 0, nvars)\n if device isa CPU\n\n for vari in 1:nvars\n MPI.Gatherv!(\n view(iv, :, vari),\n pid == root ? VBuffer(view(v_all, :, vari), Np_all) :\n nothing,\n root,\n mpicomm,\n )\n end\n\n elseif device isa CUDADevice\n\n v = Array(iv)\n for vari in 1:nvars\n MPI.Gatherv!(\n view(v, :, vari),\n pid == root ? VBuffer(view(v_all, :, vari), Np_all) :\n nothing,\n root,\n mpicomm,\n )\n end\n v_all = CuArray(v_all)\n\n else\n error(\"accumulate_interpolate_data: unsupported device, only CPU() and CUDADevice() supported\")\n end\n else\n v_all = iv\n end\n\n if pid == 0\n comp_stream = Event(device)\n thr_x = min(256, np_tot)\n workgroup = (thr_x,)\n ndrange = (np_tot,)\n\n comp_stream = accumulate_helper_kernel!(device, workgroup)(\n i1,\n i2,\n i3,\n v_all,\n fiv,\n ndrange = ndrange,\n dependencies = (comp_stream,),\n )\n wait(comp_stream)\n end\n MPI.Barrier(mpicomm)\n return nothing\nend\n\n@kernel function accumulate_helper_kernel!(\n i1::AbstractArray{UInt16, 1},\n i2::AbstractArray{UInt16, 1},\n i3::AbstractArray{UInt16, 1},\n v_all::AbstractArray{FT, 2},\n fiv::AbstractArray{FT, 4},\n) where {FT <: AbstractFloat}\n idx = @index(Global, Linear)\n nvars = size(v_all, 2)\n\n for vari in 1:nvars\n @inbounds fiv[i1[idx], i2[idx], i3[idx], vari] = v_all[idx, vari]\n end\nend\n\nfunction accumulate_interpolated_data(\n mpicomm::MPI.Comm,\n intrp::InterpolationTopology,\n iv::AbstractArray{FT, 2},\n) where {FT <: AbstractFloat}\n\n mpirank = MPI.Comm_rank(mpicomm)\n numranks = MPI.Comm_size(mpicomm)\n nvars = size(iv, 2)\n\n if intrp isa InterpolationCubedSphere\n nx1 = length(intrp.long_grd)\n nx2 = length(intrp.lat_grd)\n nx3 = length(intrp.rad_grd)\n np_tot = length(intrp.radi_all)\n i1 = intrp.longi_all\n i2 = intrp.lati_all\n i3 = intrp.radi_all\n elseif intrp isa InterpolationBrick\n nx1 = length(intrp.x1g)\n nx2 = length(intrp.x2g)\n nx3 = length(intrp.x3g)\n np_tot = length(intrp.x1i_all)\n i1 = intrp.x1i_all\n i2 = intrp.x2i_all\n i3 = intrp.x3i_all\n else\n error(\"Unsupported topology; only InterpolationCubedSphere and InterpolationBrick supported\")\n end\n\n if array_device(iv) isa CPU\n h_iv = iv\n h_i1 = i1\n h_i2 = i2\n h_i3 = i3\n else\n h_iv = Array(iv)\n h_i1 = Array(i1)\n h_i2 = Array(i2)\n h_i3 = Array(i3)\n end\n\n if numranks == 1\n v_all = h_iv\n else\n v_all = Array{FT}(undef, mpirank == 0 ? np_tot : 0, nvars)\n for vari in 1:nvars\n MPI.Gatherv!(\n view(h_iv, :, vari),\n mpirank == 0 ? VBuffer(view(v_all, :, vari), intrp.Np_all) :\n nothing,\n 0,\n mpicomm,\n )\n end\n end\n\n if mpirank == 0\n fiv = Array{FT}(undef, nx1, nx2, nx3, nvars)\n for i in 1:np_tot\n for vari in 1:nvars\n @inbounds fiv[h_i1[i], h_i2[i], h_i3[i], vari] = v_all[i, vari]\n end\n end\n else\n fiv = nothing\n end\n\n return fiv\nend\n\nend # module Interpolation\n" }, { "path": "src/Numerics/Mesh/Mesh.jl", "content": "module Mesh\ninclude(\"BrickMesh.jl\")\ninclude(\"Topologies.jl\")\ninclude(\"GeometricFactors.jl\")\ninclude(\"Metrics.jl\")\ninclude(\"Elements.jl\")\ninclude(\"Grids.jl\")\ninclude(\"DSS.jl\")\ninclude(\"Filters.jl\")\ninclude(\"Geometry.jl\")\ninclude(\"Interpolation.jl\")\nend # module\n" }, { "path": "src/Numerics/Mesh/Metrics.jl", "content": "module Metrics\n\nusing ..GeometricFactors\n\nexport creategrid, compute_reference_to_physical_coord_jacobian, computemetric\n\n\"\"\"\n creategrid!(vgeo, elemtocoord, ξ)\n\nCreate a 1-D grid using `elemtocoord` (see `brickmesh`) using the 1-D\n`(-1, 1)` reference coordinates `ξ` (in 1D, `ξ = ξ1`). The element grids\nare filled using linear interpolation of the element coordinates.\n\nIf `Nq = length(ξ)` and `nelem = size(elemtocoord, 3)` then the preallocated\narray `vgeo.x1` should be `Nq * nelem == length(x1)`.\n\"\"\"\nfunction creategrid!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n e2c,\n ξ::NTuple{1, Vector{FT}},\n) where {Nq, FT}\n (d, nvert, nelem) = size(e2c)\n @assert d == 1\n (ξ1,) = ξ\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n\n # Linear blend\n @inbounds for e in 1:nelem\n for i in 1:Nq[1]\n vgeo.x1[i, e] =\n ((1 - ξ1[i]) * e2c[1, 1, e] + (1 + ξ1[i]) * e2c[1, 2, e]) / 2\n end\n end\n nothing\nend\n\n\"\"\"\n creategrid!(vgeo, elemtocoord, ξ)\n\nCreate a 2-D tensor product grid using `elemtocoord` (see `brickmesh`)\nusing the tuple `ξ = (ξ1, ξ2)`, composed by the 1D reference coordinates `ξ1` and `ξ2` in `(-1, 1)^2`.\nThe element grids are filled using bilinear interpolation of the element coordinates.\n\nIf `Nq = (length(ξ1), length(ξ2))` and `nelem = size(elemtocoord, 3)` then the\npreallocated arrays `vgeo.x1` and `vgeo.x2` should be\n`prod(Nq) * nelem == size(vgeo.x1) == size(vgeo.x2)`.\n\"\"\"\nfunction creategrid!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n e2c,\n ξ::NTuple{2, Vector{FT}},\n) where {Nq, FT}\n (d, nvert, nelem) = size(e2c)\n @assert d == 2\n (ξ1, ξ2) = ξ\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n\n # Bilinear blend of corners\n @inbounds for (f, n) in zip((x1, x2), 1:d)\n for e in 1:nelem, j in 1:Nq[2], i in 1:Nq[1]\n f[i, j, e] =\n (\n (1 - ξ1[i]) * (1 - ξ2[j]) * e2c[n, 1, e] +\n (1 + ξ1[i]) * (1 - ξ2[j]) * e2c[n, 2, e] +\n (1 - ξ1[i]) * (1 + ξ2[j]) * e2c[n, 3, e] +\n (1 + ξ1[i]) * (1 + ξ2[j]) * e2c[n, 4, e]\n ) / 4\n end\n end\n nothing\nend\n\n\"\"\"\n creategrid!(vgeo, elemtocoord, ξ)\n\nCreate a 3-D tensor product grid using `elemtocoord` (see `brickmesh`)\nusing the tuple `ξ = (ξ1, ξ2, ξ3)`, composed by the 1D reference coordinates `ξ1`, `ξ2`, `ξ3` in `(-1, 1)^3`.\nThe element grids are filled using trilinear interpolation of the element coordinates.\n\nIf `Nq = (length(ξ1), length(ξ2), length(ξ3))` and\n`nelem = size(elemtocoord, 3)` then the preallocated arrays `vgeo.x1`, `vgeo.x2`,\nand `vgeo.x3` should be `prod(Nq) * nelem == size(vgeo.x1) == size(vgeo.x2) == size(vgeo.x3)`.\n\"\"\"\nfunction creategrid!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n e2c,\n ξ::NTuple{3, Vector{FT}},\n) where {Nq, FT}\n (d, nvert, nelem) = size(e2c)\n @assert d == 3\n (ξ1, ξ2, ξ3) = ξ\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n x3 = reshape(vgeo.x3, (Nq..., nelem))\n\n # Trilinear blend of corners\n @inbounds for (f, n) in zip((x1, x2, x3), 1:d)\n for e in 1:nelem, k in 1:Nq[3], j in 1:Nq[2], i in 1:Nq[1]\n f[i, j, k, e] =\n (\n (1 - ξ1[i]) * (1 - ξ2[j]) * (1 - ξ3[k]) * e2c[n, 1, e] +\n (1 + ξ1[i]) * (1 - ξ2[j]) * (1 - ξ3[k]) * e2c[n, 2, e] +\n (1 - ξ1[i]) * (1 + ξ2[j]) * (1 - ξ3[k]) * e2c[n, 3, e] +\n (1 + ξ1[i]) * (1 + ξ2[j]) * (1 - ξ3[k]) * e2c[n, 4, e] +\n (1 - ξ1[i]) * (1 - ξ2[j]) * (1 + ξ3[k]) * e2c[n, 5, e] +\n (1 + ξ1[i]) * (1 - ξ2[j]) * (1 + ξ3[k]) * e2c[n, 6, e] +\n (1 - ξ1[i]) * (1 + ξ2[j]) * (1 + ξ3[k]) * e2c[n, 7, e] +\n (1 + ξ1[i]) * (1 + ξ2[j]) * (1 + ξ3[k]) * e2c[n, 8, e]\n ) / 8\n end\n end\n nothing\nend\n\n\"\"\"\n compute_reference_to_physical_coord_jacobian!(vgeo, nelem, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- D::NTuple{2,Int}, a tuple of derivative matrices, i.e., D = (D1,), where:\n - D1::DAT2, 1-D derivative operator on the device in the first dimension\n\nCompute the Jacobian matrix, ∂x / ∂ξ, of the transformation from reference coordinates,\n`ξ1`, to physical coordinates, `vgeo.x1`, for each quadrature point in element e.\n\"\"\"\nfunction compute_reference_to_physical_coord_jacobian!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n nelem,\n D::NTuple{1, Matrix{FT}},\n) where {Nq, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n\n T = eltype(vgeo.x1)\n (D1,) = D\n\n vgeo.x1ξ1 .= zero(T)\n\n for e in 1:nelem\n for i in 1:Nq[1]\n for n in 1:Nq[1]\n vgeo.x1ξ1[i, e] += D1[i, n] * vgeo.x1[n, e]\n end\n end\n end\n\n return vgeo\nend\n\n\"\"\"\n compute_reference_to_physical_coord_jacobian!(vgeo, nelem, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- D::NTuple{2,Int}, a tuple of derivative matrices, i.e., D = (D1, D2), where:\n - D1::DAT2, 1-D derivative operator on the device in the first dimension\n - D2::DAT2, 1-D derivative operator on the device in the second dimension\n\nCompute the Jacobian matrix, ∂x / ∂ξ, of the transformation from reference coordinates,\n`ξ1`, `ξ2`, to physical coordinates, `vgeo.x1`, `vgeo.x2`,\nfor each quadrature point in element e.\n\"\"\"\nfunction compute_reference_to_physical_coord_jacobian!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n nelem,\n D::NTuple{2, Matrix{FT}},\n) where {Nq, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n\n T = eltype(vgeo.x1)\n (D1, D2) = D\n\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n x1ξ1 = reshape(vgeo.x1ξ1, (Nq..., nelem))\n x2ξ1 = reshape(vgeo.x2ξ1, (Nq..., nelem))\n x1ξ2 = reshape(vgeo.x1ξ2, (Nq..., nelem))\n x2ξ2 = reshape(vgeo.x2ξ2, (Nq..., nelem))\n\n x1ξ1 .= x1ξ2 .= zero(T)\n x2ξ1 .= x2ξ2 .= zero(T)\n\n for e in 1:nelem\n for j in 1:Nq[2], i in 1:Nq[1]\n for n in 1:Nq[1]\n x1ξ1[i, j, e] += D1[i, n] * x1[n, j, e]\n x2ξ1[i, j, e] += D1[i, n] * x2[n, j, e]\n end\n for n in 1:Nq[2]\n x1ξ2[i, j, e] += D2[j, n] * x1[i, n, e]\n x2ξ2[i, j, e] += D2[j, n] * x2[i, n, e]\n end\n end\n end\n\n return vgeo\nend\n\n\"\"\"\n compute_reference_to_physical_coord_jacobian!(vgeo, nelem, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- D::NTuple{3,Int}, a tuple of derivative matrices, i.e., D = (D1, D2, D3), where:\n - D1::DAT2, 1-D derivative operator on the device in the first dimension\n - D2::DAT2, 1-D derivative operator on the device in the second dimension\n - D3::DAT2, 1-D derivative operator on the device in the third dimension\n\nCompute the Jacobian matrix, ∂x / ∂ξ, of the transformation from reference coordinates,\n`ξ1`, `ξ2`, `ξ3` to physical coordinates, `vgeo.x1`, `vgeo.x2`, `vgeo.x3` for\neach quadrature point in element e.\n\"\"\"\nfunction compute_reference_to_physical_coord_jacobian!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n nelem,\n D::NTuple{3, Matrix{FT}},\n) where {Nq, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n\n T = eltype(vgeo.x1)\n (D1, D2, D3) = D\n\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n x3 = reshape(vgeo.x3, (Nq..., nelem))\n x1ξ1 = reshape(vgeo.x1ξ1, (Nq..., nelem))\n x2ξ1 = reshape(vgeo.x2ξ1, (Nq..., nelem))\n x3ξ1 = reshape(vgeo.x3ξ1, (Nq..., nelem))\n x1ξ2 = reshape(vgeo.x1ξ2, (Nq..., nelem))\n x2ξ2 = reshape(vgeo.x2ξ2, (Nq..., nelem))\n x3ξ2 = reshape(vgeo.x3ξ2, (Nq..., nelem))\n x1ξ3 = reshape(vgeo.x1ξ3, (Nq..., nelem))\n x2ξ3 = reshape(vgeo.x2ξ3, (Nq..., nelem))\n x3ξ3 = reshape(vgeo.x3ξ3, (Nq..., nelem))\n\n x1ξ1 .= x1ξ2 .= x1ξ3 .= zero(T)\n x2ξ1 .= x2ξ2 .= x2ξ3 .= zero(T)\n x3ξ1 .= x3ξ2 .= x3ξ3 .= zero(T)\n\n @inbounds for e in 1:nelem\n for k in 1:Nq[3], j in 1:Nq[2], i in 1:Nq[1]\n for n in 1:Nq[1]\n x1ξ1[i, j, k, e] += D1[i, n] * x1[n, j, k, e]\n x2ξ1[i, j, k, e] += D1[i, n] * x2[n, j, k, e]\n x3ξ1[i, j, k, e] += D1[i, n] * x3[n, j, k, e]\n end\n for n in 1:Nq[2]\n x1ξ2[i, j, k, e] += D2[j, n] * x1[i, n, k, e]\n x2ξ2[i, j, k, e] += D2[j, n] * x2[i, n, k, e]\n x3ξ2[i, j, k, e] += D2[j, n] * x3[i, n, k, e]\n end\n for n in 1:Nq[3]\n x1ξ3[i, j, k, e] += D3[k, n] * x1[i, j, n, e]\n x2ξ3[i, j, k, e] += D3[k, n] * x2[i, j, n, e]\n x3ξ3[i, j, k, e] += D3[k, n] * x3[i, j, n, e]\n end\n end\n end\n\n return vgeo\nend\n\n\"\"\"\n computemetric!(vgeo, sgeo, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- sgeo::SurfaceGeometry, a struct containing the surface geometric factors\n- D::NTuple{1,Int}, tuple with 1-D derivative operator on the device\n\nCompute the 1-D metric terms from the element grid arrays `vgeo.x1`. All the arrays\nare preallocated by the user and the (square) derivative matrix `D` should be\nconsistent with the reference grid `ξ1` used in [`creategrid!`](@ref).\n\nIf `Nq = size(D, 1)` and `nelem = div(length(x1), Nq)` then the volume arrays\n`x1`, `J`, and `ξ1x1` should all have length `Nq * nelem`. Similarly, the face\narrays `sJ` and `n1` should be of length `nface * nelem` with `nface = 2`.\n\"\"\"\nfunction computemetric!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n sgeo::SurfaceGeometry{Nfp, <:AbstractArray},\n D::NTuple{1, Matrix{FT}},\n) where {Nq, Nfp, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n\n nelem = div(length(vgeo.ωJ), Nq[1])\n ωJ = reshape(vgeo.ωJ, (Nq[1], nelem))\n nface = 2\n n1 = reshape(sgeo.n1, (1, nface, nelem))\n sωJ = reshape(sgeo.sωJ, (1, nface, nelem))\n\n # Compute vertical Jacobian determinant, JcV, and Jacobian determinant, det(∂x/∂ξ), per quadrature point\n vgeo.JcV .= vgeo.x1ξ1\n vgeo.ωJ .= vgeo.x1ξ1\n\n vgeo.ξ1x1 .= 1 ./ vgeo.ωJ\n\n sgeo.n1[1, 1, :] .= -sign.(ωJ[1, :])\n sgeo.n1[1, 2, :] .= sign.(ωJ[Nq[1], :])\n sgeo.sωJ .= 1\n nothing\nend\n\n\"\"\"\n computemetric!(vgeo, sgeo, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- sgeo::SurfaceGeometry, a struct containing the surface geometric factors\n- D::NTuple{2,Int}, a tuple of derivative matrices, i.e., D = (D1, D2), where:\n - D1::DAT2, 1-D derivative operator on the device in the first dimension\n - D2::DAT2, 1-D derivative operator on the device in the second dimension\n\nCompute the 2-D metric terms from the element grid arrays `vgeo.x1` and `vgeo.x2`. All the\narrays are preallocated by the user and the (square) derivative matrice `D1` and\n`D2` should be consistent with the reference grid `ξ1` and `ξ2` used in\n[`creategrid!`](@ref).\n\nIf `Nq = (size(D1, 1), size(D2, 1))` and `nelem = div(length(vgeo.x1), prod(Nq))`\nthen the volume arrays `vgeo.x1`, `vgeo.x2`, `vgeo.ωJ`, `vgeo.ξ1x1`, `vgeo.ξ2x1`, `vgeo.ξ1x2`, and `vgeo.ξ2x2`\nshould all be of size `(Nq..., nelem)`. Similarly, the face arrays `sgeo.sωJ`, `sgeo.n1`,\nand `sgeo.n2` should be of size `(maximum(Nq), nface, nelem)` with `nface = 4`\n\"\"\"\nfunction computemetric!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n sgeo::SurfaceGeometry{Nfp, <:AbstractArray},\n D::NTuple{2, Matrix{FT}},\n) where {Nq, Nfp, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n @assert Nfp == div.(prod(Nq), Nq)\n\n nelem = div(length(vgeo.ωJ), prod(Nq))\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n ωJ = reshape(vgeo.ωJ, (Nq..., nelem))\n JcV = reshape(vgeo.JcV, (Nq..., nelem))\n ξ1x1 = reshape(vgeo.ξ1x1, (Nq..., nelem))\n ξ2x1 = reshape(vgeo.ξ2x1, (Nq..., nelem))\n ξ1x2 = reshape(vgeo.ξ1x2, (Nq..., nelem))\n ξ2x2 = reshape(vgeo.ξ2x2, (Nq..., nelem))\n x1ξ1 = reshape(vgeo.x1ξ1, (Nq..., nelem))\n x1ξ2 = reshape(vgeo.x1ξ2, (Nq..., nelem))\n x2ξ1 = reshape(vgeo.x2ξ1, (Nq..., nelem))\n x2ξ2 = reshape(vgeo.x2ξ2, (Nq..., nelem))\n nface = 4\n n1 = reshape(sgeo.n1, (maximum(Nfp), nface, nelem))\n n2 = reshape(sgeo.n2, (maximum(Nfp), nface, nelem))\n sωJ = reshape(sgeo.sωJ, (maximum(Nfp), nface, nelem))\n\n for e in 1:nelem\n for j in 1:Nq[2], i in 1:Nq[1]\n\n # Compute vertical Jacobian determinant, JcV, per quadrature point\n JcV[i, j, e] = hypot(x1ξ2[i, j, e], x2ξ2[i, j, e])\n # Compute Jacobian determinant, det(∂x/∂ξ), per quadrature point\n ωJ[i, j, e] =\n x1ξ1[i, j, e] * x2ξ2[i, j, e] - x2ξ1[i, j, e] * x1ξ2[i, j, e]\n\n ξ1x1[i, j, e] = x2ξ2[i, j, e] / ωJ[i, j, e]\n ξ2x1[i, j, e] = -x2ξ1[i, j, e] / ωJ[i, j, e]\n ξ1x2[i, j, e] = -x1ξ2[i, j, e] / ωJ[i, j, e]\n ξ2x2[i, j, e] = x1ξ1[i, j, e] / ωJ[i, j, e]\n end\n\n # Compute surface struct field entries\n for i in 1:maximum(Nfp)\n if i <= Nfp[1]\n sgeo.n1[i, 1, e] = -ωJ[1, i, e] * ξ1x1[1, i, e]\n sgeo.n2[i, 1, e] = -ωJ[1, i, e] * ξ1x2[1, i, e]\n sgeo.n1[i, 2, e] = ωJ[Nq[1], i, e] * ξ1x1[Nq[1], i, e]\n sgeo.n2[i, 2, e] = ωJ[Nq[1], i, e] * ξ1x2[Nq[1], i, e]\n else\n sgeo.n1[i, 1:2, e] .= NaN\n sgeo.n2[i, 1:2, e] .= NaN\n end\n if i <= Nfp[2]\n sgeo.n1[i, 3, e] = -ωJ[i, 1, e] * ξ2x1[i, 1, e]\n sgeo.n2[i, 3, e] = -ωJ[i, 1, e] * ξ2x2[i, 1, e]\n sgeo.n1[i, 4, e] = ωJ[i, Nq[2], e] * ξ2x1[i, Nq[2], e]\n sgeo.n2[i, 4, e] = ωJ[i, Nq[2], e] * ξ2x2[i, Nq[2], e]\n else\n sgeo.n1[i, 3:4, e] .= NaN\n sgeo.n2[i, 3:4, e] .= NaN\n end\n\n for n in 1:nface\n sgeo.sωJ[i, n, e] = hypot(n1[i, n, e], n2[i, n, e])\n sgeo.n1[i, n, e] /= sωJ[i, n, e]\n sgeo.n2[i, n, e] /= sωJ[i, n, e]\n end\n end\n\n end\n\n nothing\nend\n\n\"\"\"\n computemetric!(vgeo, sgeo, D)\n\nInput arguments:\n- vgeo::VolumeGeometry, a struct containing the volumetric geometric factors\n- sgeo::SurfaceGeometry, a struct containing the surface geometric factors\n- D::NTuple{3,Int}, a tuple of derivative matrices, i.e., D = (D1, D2, D3), where:\n - D1::DAT2, 1-D derivative operator on the device in the first dimension\n - D2::DAT2, 1-D derivative operator on the device in the second dimension\n - D3::DAT2, 1-D derivative operator on the device in the third dimension\n\nCompute the 3-D metric terms from the element grid arrays `vgeo.x1`, `vgeo.x2`, and `vgeo.x3`.\nAll the arrays are preallocated by the user and the (square) derivative matrice `D1`,\n`D2`, and `D3` should be consistent with the reference grid `ξ1`, `ξ2`, and `ξ3` used in\n[`creategrid!`](@ref).\n\nIf `Nq = size(D1, 1)` and `nelem = div(length(vgeo.x1), Nq^3)` then the volume arrays\n`vgeo.x1`, `vgeo.x2`, `vgeo.x3`, `vgeo.ωJ`, `vgeo.ξ1x1`, `vgeo.ξ2x1`, `vgeo.ξ3x1`,\n`vgeo.ξ1x2`, `vgeo.ξ2x2`, `vgeo.ξ3x2`, `vgeo.ξ1x3`,`vgeo.ξ2x3`, and `vgeo.ξ3x3`\nshould all be of length `Nq^3 * nelem`. Similarly, the face\narrays `sgeo.sωJ`, `sgeo.n1`, `sgeo.n2`, and `sgeo.n3` should be of size `Nq^2 * nface * nelem` with\n`nface = 6`.\n\nThe curl invariant formulation of Kopriva (2006), equation 37, is used.\n\nReference:\n - [Kopriva2006](@cite)\n\"\"\"\nfunction computemetric!(\n vgeo::VolumeGeometry{Nq, <:AbstractArray, <:AbstractArray},\n sgeo::SurfaceGeometry{Nfp, <:AbstractArray},\n D::NTuple{3, Matrix{FT}},\n) where {Nq, Nfp, FT}\n\n @assert Nq == map(d -> size(d, 1), D)\n @assert Nfp == div.(prod(Nq), Nq)\n\n T = eltype(vgeo.x1)\n nelem = div(length(vgeo.ωJ), prod(Nq))\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n x3 = reshape(vgeo.x3, (Nq..., nelem))\n ωJ = reshape(vgeo.ωJ, (Nq..., nelem))\n JcV = reshape(vgeo.JcV, (Nq..., nelem))\n ξ1x1 = reshape(vgeo.ξ1x1, (Nq..., nelem))\n ξ2x1 = reshape(vgeo.ξ2x1, (Nq..., nelem))\n ξ3x1 = reshape(vgeo.ξ3x1, (Nq..., nelem))\n ξ1x2 = reshape(vgeo.ξ1x2, (Nq..., nelem))\n ξ2x2 = reshape(vgeo.ξ2x2, (Nq..., nelem))\n ξ3x2 = reshape(vgeo.ξ3x2, (Nq..., nelem))\n ξ1x3 = reshape(vgeo.ξ1x3, (Nq..., nelem))\n ξ2x3 = reshape(vgeo.ξ2x3, (Nq..., nelem))\n ξ3x3 = reshape(vgeo.ξ3x3, (Nq..., nelem))\n x1ξ1 = reshape(vgeo.x1ξ1, (Nq..., nelem))\n x1ξ2 = reshape(vgeo.x1ξ2, (Nq..., nelem))\n x1ξ3 = reshape(vgeo.x1ξ3, (Nq..., nelem))\n x2ξ1 = reshape(vgeo.x2ξ1, (Nq..., nelem))\n x2ξ2 = reshape(vgeo.x2ξ2, (Nq..., nelem))\n x2ξ3 = reshape(vgeo.x2ξ3, (Nq..., nelem))\n x3ξ1 = reshape(vgeo.x3ξ1, (Nq..., nelem))\n x3ξ2 = reshape(vgeo.x3ξ2, (Nq..., nelem))\n x3ξ3 = reshape(vgeo.x3ξ3, (Nq..., nelem))\n\n nface = 6\n n1 = reshape(sgeo.n1, maximum(Nfp), nface, nelem)\n n2 = reshape(sgeo.n2, maximum(Nfp), nface, nelem)\n n3 = reshape(sgeo.n3, maximum(Nfp), nface, nelem)\n sωJ = reshape(sgeo.sωJ, maximum(Nfp), nface, nelem)\n\n JI2 = similar(vgeo.ωJ, Nq...)\n (yzr, yzs, yzt) = (similar(JI2), similar(JI2), similar(JI2))\n (zxr, zxs, zxt) = (similar(JI2), similar(JI2), similar(JI2))\n (xyr, xys, xyt) = (similar(JI2), similar(JI2), similar(JI2))\n # Temporary variables to compute inverse of a 3x3 matrix\n (a11, a12, a13) = (similar(JI2), similar(JI2), similar(JI2))\n (a21, a22, a23) = (similar(JI2), similar(JI2), similar(JI2))\n (a31, a32, a33) = (similar(JI2), similar(JI2), similar(JI2))\n\n ξ1x1 .= ξ2x1 .= ξ3x1 .= zero(T)\n ξ1x2 .= ξ2x2 .= ξ3x2 .= zero(T)\n ξ1x3 .= ξ2x3 .= ξ3x3 .= zero(T)\n\n fill!(n1, NaN)\n fill!(n2, NaN)\n fill!(n3, NaN)\n fill!(sωJ, NaN)\n\n @inbounds for e in 1:nelem\n for k in 1:Nq[3], j in 1:Nq[2], i in 1:Nq[1]\n\n # Compute vertical Jacobian determinant, JcV, per quadrature point\n JcV[i, j, k, e] =\n hypot(x1ξ3[i, j, k, e], x2ξ3[i, j, k, e], x3ξ3[i, j, k, e])\n # Compute Jacobian determinant, det(∂x/∂ξ), per quadrature point\n ωJ[i, j, k, e] = (\n x1ξ1[i, j, k, e] * (\n x2ξ2[i, j, k, e] * x3ξ3[i, j, k, e] -\n x3ξ2[i, j, k, e] * x2ξ3[i, j, k, e]\n ) +\n x2ξ1[i, j, k, e] * (\n x3ξ2[i, j, k, e] * x1ξ3[i, j, k, e] -\n x1ξ2[i, j, k, e] * x3ξ3[i, j, k, e]\n ) +\n x3ξ1[i, j, k, e] * (\n x1ξ2[i, j, k, e] * x2ξ3[i, j, k, e] -\n x2ξ2[i, j, k, e] * x1ξ3[i, j, k, e]\n )\n )\n\n JI2[i, j, k] = 1 / (2 * ωJ[i, j, k, e])\n\n yzr[i, j, k] =\n x2[i, j, k, e] * x3ξ1[i, j, k, e] -\n x3[i, j, k, e] * x2ξ1[i, j, k, e]\n yzs[i, j, k] =\n x2[i, j, k, e] * x3ξ2[i, j, k, e] -\n x3[i, j, k, e] * x2ξ2[i, j, k, e]\n yzt[i, j, k] =\n x2[i, j, k, e] * x3ξ3[i, j, k, e] -\n x3[i, j, k, e] * x2ξ3[i, j, k, e]\n zxr[i, j, k] =\n x3[i, j, k, e] * x1ξ1[i, j, k, e] -\n x1[i, j, k, e] * x3ξ1[i, j, k, e]\n zxs[i, j, k] =\n x3[i, j, k, e] * x1ξ2[i, j, k, e] -\n x1[i, j, k, e] * x3ξ2[i, j, k, e]\n zxt[i, j, k] =\n x3[i, j, k, e] * x1ξ3[i, j, k, e] -\n x1[i, j, k, e] * x3ξ3[i, j, k, e]\n xyr[i, j, k] =\n x1[i, j, k, e] * x2ξ1[i, j, k, e] -\n x2[i, j, k, e] * x1ξ1[i, j, k, e]\n xys[i, j, k] =\n x1[i, j, k, e] * x2ξ2[i, j, k, e] -\n x2[i, j, k, e] * x1ξ2[i, j, k, e]\n xyt[i, j, k] =\n x1[i, j, k, e] * x2ξ3[i, j, k, e] -\n x2[i, j, k, e] * x1ξ3[i, j, k, e]\n end\n\n for k in 1:Nq[3], j in 1:Nq[2], i in 1:Nq[1]\n for n in 1:Nq[1]\n ξ2x1[i, j, k, e] -= D[1][i, n] * yzt[n, j, k]\n ξ3x1[i, j, k, e] += D[1][i, n] * yzs[n, j, k]\n ξ2x2[i, j, k, e] -= D[1][i, n] * zxt[n, j, k]\n ξ3x2[i, j, k, e] += D[1][i, n] * zxs[n, j, k]\n ξ2x3[i, j, k, e] -= D[1][i, n] * xyt[n, j, k]\n ξ3x3[i, j, k, e] += D[1][i, n] * xys[n, j, k]\n end\n for n in 1:Nq[2]\n ξ1x1[i, j, k, e] += D[2][j, n] * yzt[i, n, k]\n ξ3x1[i, j, k, e] -= D[2][j, n] * yzr[i, n, k]\n ξ1x2[i, j, k, e] += D[2][j, n] * zxt[i, n, k]\n ξ3x2[i, j, k, e] -= D[2][j, n] * zxr[i, n, k]\n ξ1x3[i, j, k, e] += D[2][j, n] * xyt[i, n, k]\n ξ3x3[i, j, k, e] -= D[2][j, n] * xyr[i, n, k]\n end\n for n in 1:Nq[3]\n ξ1x1[i, j, k, e] -= D[3][k, n] * yzs[i, j, n]\n ξ2x1[i, j, k, e] += D[3][k, n] * yzr[i, j, n]\n ξ1x2[i, j, k, e] -= D[3][k, n] * zxs[i, j, n]\n ξ2x2[i, j, k, e] += D[3][k, n] * zxr[i, j, n]\n ξ1x3[i, j, k, e] -= D[3][k, n] * xys[i, j, n]\n ξ2x3[i, j, k, e] += D[3][k, n] * xyr[i, j, n]\n end\n ξ1x1[i, j, k, e] *= JI2[i, j, k]\n ξ2x1[i, j, k, e] *= JI2[i, j, k]\n ξ3x1[i, j, k, e] *= JI2[i, j, k]\n ξ1x2[i, j, k, e] *= JI2[i, j, k]\n ξ2x2[i, j, k, e] *= JI2[i, j, k]\n ξ3x2[i, j, k, e] *= JI2[i, j, k]\n ξ1x3[i, j, k, e] *= JI2[i, j, k]\n ξ2x3[i, j, k, e] *= JI2[i, j, k]\n ξ3x3[i, j, k, e] *= JI2[i, j, k]\n\n\n # Invert ∂ξk/∂xi, since the discrete curl-invariant form that we have\n # just computed, ∂ξk/∂xi, is not equal to its inverse\n a11[i, j, k] =\n ξ2x2[i, j, k, e] * ξ3x3[i, j, k, e] -\n ξ2x3[i, j, k, e] * ξ3x2[i, j, k, e]\n a12[i, j, k] =\n ξ1x3[i, j, k, e] * ξ3x2[i, j, k, e] -\n ξ1x2[i, j, k, e] * ξ3x3[i, j, k, e]\n a13[i, j, k] =\n ξ1x2[i, j, k, e] * ξ2x3[i, j, k, e] -\n ξ1x3[i, j, k, e] * ξ2x2[i, j, k, e]\n a21[i, j, k] =\n ξ2x3[i, j, k, e] * ξ3x1[i, j, k, e] -\n ξ2x1[i, j, k, e] * ξ3x3[i, j, k, e]\n a22[i, j, k] =\n ξ1x1[i, j, k, e] * ξ3x3[i, j, k, e] -\n ξ1x3[i, j, k, e] * ξ3x1[i, j, k, e]\n a23[i, j, k] =\n ξ1x3[i, j, k, e] * ξ2x1[i, j, k, e] -\n ξ1x1[i, j, k, e] * ξ2x3[i, j, k, e]\n a31[i, j, k] =\n ξ2x1[i, j, k, e] * ξ3x2[i, j, k, e] -\n ξ2x2[i, j, k, e] * ξ3x1[i, j, k, e]\n a32[i, j, k] =\n ξ1x2[i, j, k, e] * ξ3x1[i, j, k, e] -\n ξ1x1[i, j, k, e] * ξ3x2[i, j, k, e]\n a33[i, j, k] =\n ξ1x1[i, j, k, e] * ξ2x2[i, j, k, e] -\n ξ1x2[i, j, k, e] * ξ2x1[i, j, k, e]\n\n det =\n ξ1x1[i, j, k, e] * a11[i, j, k] +\n ξ2x1[i, j, k, e] * a12[i, j, k] +\n ξ3x1[i, j, k, e] * a13[i, j, k]\n\n x1ξ1[i, j, k, e] =\n 1.0 / det * (\n a11[i, j, k] * a11[i, j, k] +\n a12[i, j, k] * a12[i, j, k] +\n a13[i, j, k] * a13[i, j, k]\n )\n x1ξ2[i, j, k, e] =\n 1.0 / det * (\n a11[i, j, k] * a21[i, j, k] +\n a12[i, j, k] * a22[i, j, k] +\n a13[i, j, k] * a23[i, j, k]\n )\n x1ξ3[i, j, k, e] =\n 1.0 / det * (\n a11[i, j, k] * a31[i, j, k] +\n a12[i, j, k] * a32[i, j, k] +\n a13[i, j, k] * a33[i, j, k]\n )\n x2ξ1[i, j, k, e] =\n 1.0 / det * (\n a21[i, j, k] * a11[i, j, k] +\n a22[i, j, k] * a12[i, j, k] +\n a23[i, j, k] * a13[i, j, k]\n )\n x2ξ2[i, j, k, e] =\n 1.0 / det * (\n a21[i, j, k] * a21[i, j, k] +\n a22[i, j, k] * a22[i, j, k] +\n a23[i, j, k] * a23[i, j, k]\n )\n x2ξ3[i, j, k, e] =\n 1.0 / det * (\n a21[i, j, k] * a31[i, j, k] +\n a22[i, j, k] * a32[i, j, k] +\n a23[i, j, k] * a33[i, j, k]\n )\n x3ξ1[i, j, k, e] =\n 1.0 / det * (\n a31[i, j, k] * a11[i, j, k] +\n a32[i, j, k] * a12[i, j, k] +\n a33[i, j, k] * a13[i, j, k]\n )\n x3ξ2[i, j, k, e] =\n 1.0 / det * (\n a31[i, j, k] * a21[i, j, k] +\n a32[i, j, k] * a22[i, j, k] +\n a33[i, j, k] * a23[i, j, k]\n )\n x3ξ3[i, j, k, e] =\n 1.0 / det * (\n a31[i, j, k] * a31[i, j, k] +\n a32[i, j, k] * a32[i, j, k] +\n a33[i, j, k] * a33[i, j, k]\n )\n end\n\n # Compute surface struct field entries\n # faces 1 & 2\n for k in 1:Nq[3], j in 1:Nq[2]\n n = j + (k - 1) * Nq[2]\n sgeo.n1[n, 1, e] = -ωJ[1, j, k, e] * ξ1x1[1, j, k, e]\n sgeo.n2[n, 1, e] = -ωJ[1, j, k, e] * ξ1x2[1, j, k, e]\n sgeo.n3[n, 1, e] = -ωJ[1, j, k, e] * ξ1x3[1, j, k, e]\n sgeo.n1[n, 2, e] = ωJ[Nq[1], j, k, e] * ξ1x1[Nq[1], j, k, e]\n sgeo.n2[n, 2, e] = ωJ[Nq[1], j, k, e] * ξ1x2[Nq[1], j, k, e]\n sgeo.n3[n, 2, e] = ωJ[Nq[1], j, k, e] * ξ1x3[Nq[1], j, k, e]\n for f in 1:2\n sgeo.sωJ[n, f, e] = hypot(n1[n, f, e], n2[n, f, e], n3[n, f, e])\n sgeo.n1[n, f, e] /= sωJ[n, f, e]\n sgeo.n2[n, f, e] /= sωJ[n, f, e]\n sgeo.n3[n, f, e] /= sωJ[n, f, e]\n end\n end\n # faces 3 & 4\n for k in 1:Nq[3], i in 1:Nq[1]\n n = i + (k - 1) * Nq[1]\n sgeo.n1[n, 3, e] = -ωJ[i, 1, k, e] * ξ2x1[i, 1, k, e]\n sgeo.n2[n, 3, e] = -ωJ[i, 1, k, e] * ξ2x2[i, 1, k, e]\n sgeo.n3[n, 3, e] = -ωJ[i, 1, k, e] * ξ2x3[i, 1, k, e]\n sgeo.n1[n, 4, e] = ωJ[i, Nq[2], k, e] * ξ2x1[i, Nq[2], k, e]\n sgeo.n2[n, 4, e] = ωJ[i, Nq[2], k, e] * ξ2x2[i, Nq[2], k, e]\n sgeo.n3[n, 4, e] = ωJ[i, Nq[2], k, e] * ξ2x3[i, Nq[2], k, e]\n for f in 3:4\n sgeo.sωJ[n, f, e] = hypot(n1[n, f, e], n2[n, f, e], n3[n, f, e])\n sgeo.n1[n, f, e] /= sωJ[n, f, e]\n sgeo.n2[n, f, e] /= sωJ[n, f, e]\n sgeo.n3[n, f, e] /= sωJ[n, f, e]\n end\n end\n # faces 5 & 6\n for j in 1:Nq[2], i in 1:Nq[1]\n n = i + (j - 1) * Nq[1]\n sgeo.n1[n, 5, e] = -ωJ[i, j, 1, e] * ξ3x1[i, j, 1, e]\n sgeo.n2[n, 5, e] = -ωJ[i, j, 1, e] * ξ3x2[i, j, 1, e]\n sgeo.n3[n, 5, e] = -ωJ[i, j, 1, e] * ξ3x3[i, j, 1, e]\n sgeo.n1[n, 6, e] = ωJ[i, j, Nq[3], e] * ξ3x1[i, j, Nq[3], e]\n sgeo.n2[n, 6, e] = ωJ[i, j, Nq[3], e] * ξ3x2[i, j, Nq[3], e]\n sgeo.n3[n, 6, e] = ωJ[i, j, Nq[3], e] * ξ3x3[i, j, Nq[3], e]\n for f in 5:6\n sgeo.sωJ[n, f, e] = hypot(n1[n, f, e], n2[n, f, e], n3[n, f, e])\n sgeo.n1[n, f, e] /= sωJ[n, f, e]\n sgeo.n2[n, f, e] /= sωJ[n, f, e]\n sgeo.n3[n, f, e] /= sωJ[n, f, e]\n end\n end\n end\n\n nothing\nend\n\nend # module\n" }, { "path": "src/Numerics/Mesh/Topologies.jl", "content": "module Topologies\nusing ClimateMachine\nusing CubedSphere, Rotations\nimport ..BrickMesh\nimport MPI\nusing CUDA\nusing DocStringExtensions\n\nexport AbstractTopology,\n BrickTopology,\n StackedBrickTopology,\n CubedShellTopology,\n StackedCubedSphereTopology,\n AnalyticalTopography,\n NoTopography,\n DCMIPMountain,\n EquidistantCubedSphere,\n EquiangularCubedSphere,\n isstacked,\n cubed_sphere_topo_warp,\n compute_lat_long,\n compute_analytical_topography,\n cubed_sphere_warp,\n cubed_sphere_unwarp,\n equiangular_cubed_sphere_warp,\n equiangular_cubed_sphere_unwarp,\n equidistant_cubed_sphere_warp,\n equidistant_cubed_sphere_unwarp,\n conformal_cubed_sphere_warp,\n conformal_cubed_sphere_unwarp\n\nexport grid1d, SingleExponentialStretching, InteriorStretching\n\nexport basic_topology_info\n\n\"\"\"\n AbstractTopology{dim, T, nb}\n\nRepresents the connectivity of individual elements, with local dimension `dim`\nwith `nb` boundary conditions types. The element coordinates are of type `T`.\n\"\"\"\nabstract type AbstractTopology{dim, T, nb} end\n\n\"\"\"\n BoxElementTopology{dim, T, nb} <: AbstractTopology{dim,T,nb}\n\nThe local topology of a larger MPI-distributed topology, represented by\n`dim`-dimensional box elements, with `nb` boundary conditions.\n\nThis contains the necessary information for the connectivity elements of the\nelements on the local process, along with \"ghost\" elements from neighbouring\nprocesses.\n\n# Fields\n\n$(DocStringExtensions.FIELDS)\n\"\"\"\nstruct BoxElementTopology{dim, T, nb} <: AbstractTopology{dim, T, nb}\n\n \"\"\"\n MPI communicator for communicating with neighbouring processes.\n \"\"\"\n mpicomm::MPI.Comm\n\n \"\"\"\n Range of element indices\n \"\"\"\n elems::UnitRange{Int64}\n\n \"\"\"\n Range of real (aka nonghost) element indices\n \"\"\"\n realelems::UnitRange{Int64}\n\n \"\"\"\n Range of ghost element indices\n \"\"\"\n ghostelems::UnitRange{Int64}\n\n \"\"\"\n Ghost element to face is received; `ghostfaces[f,ge] == true` if face `f` of\n ghost element `ge` is received.\n \"\"\"\n ghostfaces::BitArray{2}\n\n \"\"\"\n Array of send element indices\n \"\"\"\n sendelems::Array{Int64, 1}\n\n \"\"\"\n Send element to face is sent; `sendfaces[f,se] == true` if face `f` of send\n element `se` is sent.\n \"\"\"\n sendfaces::BitArray{2}\n\n \"\"\"\n Array of real elements that do not have a ghost element as a neighbor.\n \"\"\"\n interiorelems::Array{Int64, 1}\n\n \"\"\"\n Array of real elements that have at least on ghost element as a neighbor.\n\n Note that this is different from `sendelems` because `sendelems` duplicates\n elements that need to be sent to multiple neighboring processes.\n \"\"\"\n exteriorelems::Array{Int64, 1}\n\n \"\"\"\n Element to vertex coordinates; `elemtocoord[d,i,e]` is the `d`th coordinate\n of corner `i` of element `e`\n\n !!! note\n currently coordinates always are of size 3 for `(x1, x2, x3)`\n \"\"\"\n elemtocoord::Array{T, 3}\n\n \"\"\"\n Element to neighboring element; `elemtoelem[f,e]` is the number of the\n element neighboring element `e` across face `f`. If it is a boundary face,\n then it is the boundary element index.\n \"\"\"\n elemtoelem::Array{Int64, 2}\n\n \"\"\"\n Element to neighboring element face; `elemtoface[f,e]` is the face number of\n the element neighboring element `e` across face `f`. If there is no\n neighboring element then `elemtoface[f,e] == f`.\"\n \"\"\"\n elemtoface::Array{Int64, 2}\n\n \"\"\"\n element to neighboring element order; `elemtoordr[f,e]` is the ordering\n number of the element neighboring element `e` across face `f`. If there is\n no neighboring element then `elemtoordr[f,e] == 1`.\n \"\"\"\n elemtoordr::Array{Int64, 2}\n\n \"\"\"\n Element to boundary number; `elemtobndy[f,e]` is the boundary number of face\n `f` of element `e`. If there is a neighboring element then `elemtobndy[f,e]\n == 0`.\n \"\"\"\n elemtobndy::Array{Int64, 2}\n\n \"\"\"\n List of the MPI ranks for the neighboring processes\n \"\"\"\n nabrtorank::Array{Int64, 1}\n\n \"\"\"\n Range in ghost elements to receive for each neighbor\n \"\"\"\n nabrtorecv::Array{UnitRange{Int64}, 1}\n\n \"\"\"\n Range in `sendelems` to send for each neighbor\n \"\"\"\n nabrtosend::Array{UnitRange{Int64}, 1}\n\n \"\"\"\n original order in partitioning\n \"\"\"\n origsendorder::Array{Int64, 1}\n\n \"\"\"\n Tuple of boundary to element. `bndytoelem[b][i]` is the element which faces\n the `i`th boundary element of boundary `b`.\n \"\"\"\n bndytoelem::NTuple{nb, Array{Int64, 1}}\n\n \"\"\"\n Tuple of boundary to element face. `bndytoface[b][i]` is the face number of\n the element which faces the `i`th boundary element of boundary `b`.\n \"\"\"\n bndytoface::NTuple{nb, Array{Int64, 1}}\n\n \"\"\"\n Element to locally unique vertex number; `elemtouvert[v,e]` is the `v`th vertex\n of element `e`\n \"\"\"\n elemtouvert::Union{Array{Int64, 2}, Nothing}\n\n \"\"\"\n Vertex connectivity information for direct stiffness summation;\n `vtconn[vtconnoff[i]:vtconnoff[i+1]-1] == vtconn[lvt1, lelem1, lvt2, lelem2, ....]`\n for each rank-local unique vertex number,\n where `lvt1` is the element-local vertex number of element `lelem1`, etc.\n Vertices, not shared by multiple elements, are ignored.\n \"\"\"\n vtconn::Union{AbstractArray{Int64, 1}, Nothing}\n\n \"\"\"\n Vertex offset for vertex connectivity information\n \"\"\"\n vtconnoff::Union{AbstractArray{Int64, 1}, Nothing}\n\n \"\"\"\n Face connectivity information for direct stiffness summation;\n `fcconn[ufc,:] = [lfc1, lelem1, lfc2, lelem2, ordr]`\n where `ufc` is the rank-local unique face number, and\n `lfc1` is the element-local face number of element `lelem1`, etc.,\n and ordr is the relative orientation. Faces, not shared by multiple\n elements are ignored.\n \"\"\"\n fcconn::Union{AbstractArray{Int64, 2}, Nothing}\n\n \"\"\"\n Edge connectivity information for direct stiffness summation;\n `edgconn[edgconnoff[i]:edgconnoff[i+1]-1] == edgconn[ledg1, orient, lelem1, ledg2, orient, lelem2, ...]`\n for each rank-local unique edge number,\n where `ledg1` is the element-local edge number, `orient1` is orientation (forward/reverse) of dof along edge.\n Edges, not shared by multiple elements, are ignored.\n \"\"\"\n edgconn::Union{AbstractArray{Int64, 1}, Nothing}\n\n \"\"\"\n Edge offset for edge connectivity information\n \"\"\"\n edgconnoff::Union{AbstractArray{Int64, 1}, Nothing}\n\n function BoxElementTopology{dim, T, nb}(\n mpicomm,\n elems,\n realelems,\n ghostelems,\n ghostfaces,\n sendelems,\n sendfaces,\n elemtocoord,\n elemtoelem,\n elemtoface,\n elemtoordr,\n elemtobndy,\n nabrtorank,\n nabrtorecv,\n nabrtosend,\n origsendorder,\n bndytoelem,\n bndytoface;\n device_array = ClimateMachine.array_type(),\n elemtouvert = nothing,\n vtconn = nothing,\n vtconnoff = nothing,\n fcconn = nothing,\n edgconn = nothing,\n edgconnoff = nothing,\n ) where {dim, T, nb}\n\n exteriorelems = sort(unique(sendelems))\n interiorelems = sort(setdiff(realelems, exteriorelems))\n if vtconn isa Array && vtconnoff isa Array\n vtconn = device_array(vtconn)\n vtconnoff = device_array(vtconnoff)\n end\n if fcconn isa Array\n fcconn = device_array(fcconn)\n end\n if edgconn isa Array && edgconnoff isa Array\n edgconn = device_array(edgconn)\n edgconnoff = device_array(edgconnoff)\n end\n return new{dim, T, nb}(\n mpicomm,\n elems,\n realelems,\n ghostelems,\n ghostfaces,\n sendelems,\n sendfaces,\n interiorelems,\n exteriorelems,\n elemtocoord,\n elemtoelem,\n elemtoface,\n elemtoordr,\n elemtobndy,\n nabrtorank,\n nabrtorecv,\n nabrtosend,\n origsendorder,\n bndytoelem,\n bndytoface,\n elemtouvert,\n vtconn,\n vtconnoff,\n fcconn,\n edgconn,\n edgconnoff,\n )\n end\nend\n\n\"\"\"\n hasboundary(topology::AbstractTopology)\n\nquery function to check whether a topology has a boundary (i.e., not fully\nperiodic)\n\"\"\"\nhasboundary(topology::AbstractTopology{dim, T, nb}) where {dim, T, nb} = nb != 0\n\nif VERSION >= v\"1.2-\"\n isstacked(::T) where {T <: AbstractTopology} = hasfield(T, :stacksize)\nelse\n isstacked(::T) where {T <: AbstractTopology} =\n Base.fieldindex(T, :stacksize, false) > 0\nend\n\n\"\"\"\n BrickTopology{dim, T, nb} <: AbstractTopology{dim, T, nb}\n\nA simple grid-based topology. This is a convenience wrapper around\n[`BoxElementTopology`](@ref).\n\"\"\"\nstruct BrickTopology{dim, T, nb} <: AbstractTopology{dim, T, nb}\n topology::BoxElementTopology{dim, T, nb}\nend\nBase.getproperty(a::BrickTopology, p::Symbol) =\n getproperty(getfield(a, :topology), p)\n\n\"\"\"\n CubedShellTopology{T} <: AbstractTopology{2, T, 0}\n\nA cube-shell topology. This is a convenience wrapper around\n[`BoxElementTopology`](@ref).\n\"\"\"\nstruct CubedShellTopology{T} <: AbstractTopology{2, T, 0}\n topology::BoxElementTopology{2, T, 0}\nend\nBase.getproperty(a::CubedShellTopology, p::Symbol) =\n getproperty(getfield(a, :topology), p)\n\nabstract type AbstractStackedTopology{dim, T, nb} <:\n AbstractTopology{dim, T, nb} end\n\n\n\"\"\"\n StackedBrickTopology{dim, T, nb} <: AbstractStackedTopology{dim}\n\nA simple grid-based topology, where all elements on the trailing dimension are\nstacked to be contiguous. This is a convenience wrapper around\n[`BoxElementTopology`](@ref).\n\"\"\"\nstruct StackedBrickTopology{dim, T, nb} <: AbstractStackedTopology{dim, T, nb}\n topology::BoxElementTopology{dim, T, nb}\n stacksize::Int64\n periodicstack::Bool\nend\nfunction Base.getproperty(a::StackedBrickTopology, p::Symbol)\n return p in (:stacksize, :periodicstack) ? getfield(a, p) :\n getproperty(getfield(a, :topology), p)\nend\n\n\"\"\"\n StackedCubedSphereTopology{T, nb} <: AbstractStackedTopology{3, T, nb}\n\nA cube-sphere topology. All elements on the same \"vertical\" dimension are\nstacked to be contiguous. This is a convenience wrapper around\n[`BoxElementTopology`](@ref).\n\"\"\"\nstruct StackedCubedSphereTopology{T, nb} <: AbstractStackedTopology{3, T, nb}\n topology::BoxElementTopology{3, T, nb}\n stacksize::Int64\nend\nfunction Base.getproperty(a::StackedCubedSphereTopology, p::Symbol)\n if p == :periodicstack\n return false\n else\n return p == :stacksize ? getfield(a, p) :\n getproperty(getfield(a, :topology), p)\n end\nend\n\n\"\"\" A wrapper for the BrickTopology \"\"\"\nBrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where {N} =\n BrickTopology(mpicomm, map(Ne -> 0:Ne, Nelems); kw...)\n\n\"\"\"\n BrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)\n\nGenerate a brick mesh topology with coordinates given by the tuple `elemrange`\nand the periodic dimensions given by the `periodicity` tuple.\n\nThe elements of the brick are partitioned equally across the MPI ranks based\non a space-filling curve.\n\nBy default boundary faces will be marked with a one and other faces with a\nzero. Specific boundary numbers can also be passed for each face of the brick\nin `boundary`. This will mark the nonperiodic brick faces with the given\nboundary number.\n\n# Examples\n\nWe can build a 3 by 2 element two-dimensional mesh that is periodic in the\n\\$x2\\$-direction with\n```jldoctest brickmesh\n\nusing ClimateMachine.Topologies\nusing MPI\nMPI.Init()\ntopology = BrickTopology(MPI.COMM_SELF, (2:5,4:6);\n periodicity=(false,true),\n boundary=((1,2),(3,4)))\n```\nThis returns the mesh structure for\n\n x2\n\n ^\n |\n 6- +-----+-----+-----+\n | | | | |\n | | 3 | 4 | 5 |\n | | | | |\n 5- +-----+-----+-----+\n | | | | |\n | | 1 | 2 | 6 |\n | | | | |\n 4- +-----+-----+-----+\n |\n +--|-----|-----|-----|--> x1\n 2 3 4 5\n\nFor example, the (dimension by number of corners by number of elements) array\n`elemtocoord` gives the coordinates of the corners of each element.\n```jldoctest brickmesh\njulia> topology.elemtocoord\n2×4×6 Array{Int64,3}:\n[:, :, 1] =\n 2 3 2 3\n 4 4 5 5\n\n[:, :, 2] =\n 3 4 3 4\n 4 4 5 5\n\n[:, :, 3] =\n 2 3 2 3\n 5 5 6 6\n\n[:, :, 4] =\n 3 4 3 4\n 5 5 6 6\n\n[:, :, 5] =\n 4 5 4 5\n 5 5 6 6\n\n[:, :, 6] =\n 4 5 4 5\n 4 4 5 5\n```\nNote that the corners are listed in Cartesian order.\n\nThe (number of faces by number of elements) array `elemtobndy` gives the\nboundary number for each face of each element. A zero will be given for\nconnected faces.\n```jldoctest brickmesh\njulia> topology.elemtobndy\n4×6 Array{Int64,2}:\n 1 0 1 0 0 0\n 0 0 0 0 2 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n```\nNote that the faces are listed in Cartesian order.\n\n\"\"\"\nfunction BrickTopology(\n mpicomm,\n elemrange;\n boundary = ntuple(j -> (1, 1), length(elemrange)),\n periodicity = ntuple(j -> false, length(elemrange)),\n connectivity = :face,\n ghostsize = 1,\n)\n\n if boundary isa Matrix\n boundary = tuple(mapslices(x -> tuple(x...), boundary, dims = 1)...)\n end\n\n\n # We cannot handle anything else right now...\n @assert ghostsize == 1\n\n mpirank = MPI.Comm_rank(mpicomm)\n mpisize = MPI.Comm_size(mpicomm)\n topology = BrickMesh.brickmesh(\n elemrange,\n periodicity,\n part = mpirank + 1,\n numparts = mpisize,\n boundary = boundary,\n )\n topology = BrickMesh.partition(mpicomm, topology...)\n origsendorder = topology[5]\n topology =\n connectivity == :face ?\n BrickMesh.connectmesh(mpicomm, topology[1:4]...) :\n BrickMesh.connectmeshfull(mpicomm, topology[1:4]...)\n bndytoelem, bndytoface = BrickMesh.enumerateboundaryfaces!(\n topology.elemtoelem,\n topology.elemtobndy,\n periodicity,\n boundary,\n )\n\n nb = length(bndytoelem)\n dim = length(elemrange)\n T = eltype(topology.elemtocoord)\n return BrickTopology{dim, T, nb}(BoxElementTopology{dim, T, nb}(\n mpicomm,\n topology.elems,\n topology.realelems,\n topology.ghostelems,\n topology.ghostfaces,\n topology.sendelems,\n topology.sendfaces,\n topology.elemtocoord,\n topology.elemtoelem,\n topology.elemtoface,\n topology.elemtoordr,\n topology.elemtobndy,\n topology.nabrtorank,\n topology.nabrtorecv,\n topology.nabrtosend,\n origsendorder,\n bndytoelem,\n bndytoface,\n elemtouvert = topology.elemtovert,\n ))\nend\n\n\"\"\" A wrapper for the StackedBrickTopology \"\"\"\nStackedBrickTopology(mpicomm, Nelems::NTuple{N, Integer}; kw...) where {N} =\n StackedBrickTopology(mpicomm, map(Ne -> 0:Ne, Nelems); kw...)\n\n\"\"\"\n StackedBrickTopology{dim, T}(mpicomm, elemrange; boundary, periodicity)\n\nGenerate a stacked brick mesh topology with coordinates given by the tuple\n`elemrange` and the periodic dimensions given by the `periodicity` tuple.\n\nThe elements are stacked such that the elements associated with range\n`elemrange[dim]` are contiguous in the element ordering.\n\nThe elements of the brick are partitioned equally across the MPI ranks based\non a space-filling curve. Further, stacks are not split at MPI boundaries.\n\nBy default boundary faces will be marked with a one and other faces with a\nzero. Specific boundary numbers can also be passed for each face of the brick\nin `boundary`. This will mark the nonperiodic brick faces with the given\nboundary number.\n\n# Examples\n\nWe can build a 3 by 2 element two-dimensional mesh that is periodic in the\n\\$x2\\$-direction with\n```jldoctest brickmesh\n\nusing ClimateMachine.Topologies\nusing MPI\nMPI.Init()\ntopology = StackedBrickTopology(MPI.COMM_SELF, (2:5,4:6);\n periodicity=(false,true),\n boundary=((1,2),(3,4)))\n```\nThis returns the mesh structure stacked in the \\$x2\\$-direction for\n\n x2\n\n ^\n |\n 6- +-----+-----+-----+\n | | | | |\n | | 2 | 4 | 6 |\n | | | | |\n 5- +-----+-----+-----+\n | | | | |\n | | 1 | 3 | 5 |\n | | | | |\n 4- +-----+-----+-----+\n |\n +--|-----|-----|-----|--> x1\n 2 3 4 5\n\nFor example, the (dimension by number of corners by number of elements) array\n`elemtocoord` gives the coordinates of the corners of each element.\n```jldoctest brickmesh\njulia> topology.elemtocoord\n2×4×6 Array{Int64,3}:\n[:, :, 1] =\n 2 3 2 3\n 4 4 5 5\n\n[:, :, 2] =\n 2 3 2 3\n 5 5 6 6\n\n[:, :, 3] =\n 3 4 3 4\n 4 4 5 5\n\n[:, :, 4] =\n 3 4 3 4\n 5 5 6 6\n\n[:, :, 5] =\n 4 5 4 5\n 4 4 5 5\n\n[:, :, 6] =\n 4 5 4 5\n 5 5 6 6\n```\nNote that the corners are listed in Cartesian order.\n\nThe (number of faces by number of elements) array `elemtobndy` gives the\nboundary number for each face of each element. A zero will be given for\nconnected faces.\n```jldoctest brickmesh\njulia> topology.elemtobndy\n4×6 Array{Int64,2}:\n 1 0 1 0 0 0\n 0 0 0 0 2 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n```\nNote that the faces are listed in Cartesian order.\n\"\"\"\nfunction StackedBrickTopology(\n mpicomm,\n elemrange;\n boundary = ntuple(j -> (1, 1), length(elemrange)),\n periodicity = ntuple(j -> false, length(elemrange)),\n connectivity = :full,\n ghostsize = 1,\n)\n\n if boundary isa Matrix\n boundary = tuple(mapslices(x -> tuple(x...), boundary, dims = 1)...)\n end\n\n dim = length(elemrange)\n\n dim <= 1 && error(\"Stacked brick topology works for 2D and 3D\")\n\n # Build the base topology\n basetopo = BrickTopology(\n mpicomm,\n elemrange[1:(dim - 1)];\n boundary = boundary[1:(dim - 1)],\n periodicity = periodicity[1:(dim - 1)],\n connectivity = connectivity,\n ghostsize = ghostsize,\n )\n\n\n # Use the base topology to build the stacked topology\n stack = elemrange[dim]\n stacksize = length(stack) - 1\n\n nvert = 2^dim\n nface = 2dim\n\n nreal = length(basetopo.realelems) * stacksize\n nghost = length(basetopo.ghostelems) * stacksize\n\n elems = 1:(nreal + nghost)\n realelems = 1:nreal\n ghostelems = nreal .+ (1:nghost)\n\n sendelems =\n similar(basetopo.sendelems, length(basetopo.sendelems) * stacksize)\n for i in 1:length(basetopo.sendelems), j in 1:stacksize\n sendelems[stacksize * (i - 1) + j] =\n stacksize * (basetopo.sendelems[i] - 1) + j\n end\n\n ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))\n ghostfaces .= false\n\n for i in 1:length(basetopo.ghostelems), j in 1:stacksize\n e = stacksize * (i - 1) + j\n for f in 1:(2 * (dim - 1))\n ghostfaces[f, e] = basetopo.ghostfaces[f, i]\n end\n end\n\n sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))\n sendfaces .= false\n\n for i in 1:length(basetopo.sendelems), j in 1:stacksize\n e = stacksize * (i - 1) + j\n for f in 1:(2 * (dim - 1))\n sendfaces[f, e] = basetopo.sendfaces[f, i]\n end\n end\n\n elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))\n\n for i in 1:length(basetopo.elems), j in 1:stacksize\n e = stacksize * (i - 1) + j\n\n for v in 1:(2^(dim - 1))\n for d in 1:(dim - 1)\n elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i]\n elemtocoord[d, 2^(dim - 1) + v, e] =\n basetopo.elemtocoord[d, v, i]\n end\n\n elemtocoord[dim, v, e] = stack[j]\n elemtocoord[dim, 2^(dim - 1) + v, e] = stack[j + 1]\n end\n end\n\n elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))\n elemtoface = similar(basetopo.elemtoface, nface, length(elems))\n elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))\n elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))\n\n for e in 1:(length(basetopo.elems) * stacksize), f in 1:nface\n elemtoelem[f, e] = e\n elemtoface[f, e] = f\n elemtoordr[f, e] = 1\n elemtobndy[f, e] = 0\n end\n\n for i in 1:length(basetopo.realelems), j in 1:stacksize\n e1 = stacksize * (i - 1) + j\n\n for f in 1:(2 * (dim - 1))\n e2 = stacksize * (basetopo.elemtoelem[f, i] - 1) + j\n\n elemtoelem[f, e1] = e2\n elemtoface[f, e1] = basetopo.elemtoface[f, i]\n\n # We assume a simple orientation right now\n @assert basetopo.elemtoordr[f, i] == 1\n elemtoordr[f, e1] = basetopo.elemtoordr[f, i]\n end\n\n et = stacksize * (i - 1) + j + 1\n eb = stacksize * (i - 1) + j - 1\n ft = 2 * (dim - 1) + 1\n fb = 2 * (dim - 1) + 2\n ot = 1\n ob = 1\n\n if j == stacksize\n et = periodicity[dim] ? stacksize * (i - 1) + 1 : e1\n ft = periodicity[dim] ? ft : 2 * (dim - 1) + 2\n end\n if j == 1\n eb = periodicity[dim] ? stacksize * (i - 1) + stacksize : e1\n fb = periodicity[dim] ? fb : 2 * (dim - 1) + 1\n end\n\n elemtoelem[2 * (dim - 1) + 1, e1] = eb\n elemtoelem[2 * (dim - 1) + 2, e1] = et\n elemtoface[2 * (dim - 1) + 1, e1] = fb\n elemtoface[2 * (dim - 1) + 2, e1] = ft\n elemtoordr[2 * (dim - 1) + 1, e1] = ob\n elemtoordr[2 * (dim - 1) + 2, e1] = ot\n end\n\n for i in 1:length(basetopo.elems), j in 1:stacksize\n e1 = stacksize * (i - 1) + j\n\n for f in 1:(2 * (dim - 1))\n elemtobndy[f, e1] = basetopo.elemtobndy[f, i]\n end\n\n bt = bb = 0\n\n if j == stacksize\n bt = periodicity[dim] ? bt : boundary[dim][2]\n end\n if j == 1\n bb = periodicity[dim] ? bb : boundary[dim][1]\n end\n\n elemtobndy[2 * (dim - 1) + 1, e1] = bb\n elemtobndy[2 * (dim - 1) + 2, e1] = bt\n end\n\n nabrtorank = basetopo.nabrtorank\n nabrtorecv = UnitRange{Int}[\n UnitRange(\n stacksize * (first(basetopo.nabrtorecv[n]) - 1) + 1,\n stacksize * last(basetopo.nabrtorecv[n]),\n ) for n in 1:length(nabrtorank)\n ]\n nabrtosend = UnitRange{Int}[\n UnitRange(\n stacksize * (first(basetopo.nabrtosend[n]) - 1) + 1,\n stacksize * last(basetopo.nabrtosend[n]),\n ) for n in 1:length(nabrtorank)\n ]\n\n bndytoelem, bndytoface = BrickMesh.enumerateboundaryfaces!(\n elemtoelem,\n elemtobndy,\n periodicity,\n boundary,\n )\n nb = length(bndytoelem)\n\n T = eltype(basetopo.elemtocoord)\n #----setting up DSS----------\n if basetopo.elemtouvert isa Array\n #---setup vertex DSS\n nvertby2 = div(nvert, 2)\n elemtouvert = similar(basetopo.elemtouvert, nvert, length(elems))\n # base level\n for i in 1:length(basetopo.elems)\n e = stacksize * (i - 1) + 1\n for vt in 1:nvertby2\n elemtouvert[vt, e] =\n (basetopo.elemtouvert[vt, i] - 1) * (stacksize + 1) + 1\n end\n end\n # interior levels\n for i in 1:length(basetopo.elems), j in 1:(stacksize - 1)\n e = stacksize * (i - 1) + j\n for vt in 1:nvertby2\n elemtouvert[nvertby2 + vt, e] = elemtouvert[vt, e] + 1\n elemtouvert[vt, e + 1] = elemtouvert[nvertby2 + vt, e]\n end\n end\n # top level\n if periodicity[dim]\n for i in 1:length(basetopo.elems)\n e = stacksize * i\n for vt in 1:nvertby2\n elemtouvert[nvertby2 + vt, e] =\n (basetopo.elemtouvert[vt, i] - 1) * (stacksize + 1) + 1\n end\n end\n else\n for i in 1:length(basetopo.elems)\n e = stacksize * i\n for vt in 1:nvertby2\n elemtouvert[nvertby2 + vt, e] = elemtouvert[vt, e] + 1\n end\n end\n end\n vtmax = maximum(unique(elemtouvert))\n vtconn = map(j -> zeros(Int, j), zeros(Int, vtmax))\n for el in elems, lvt in 1:nvert\n vt = elemtouvert[lvt, el]\n push!(vtconn[vt], lvt) # local vertex number\n push!(vtconn[vt], el) # local elem number\n end\n # building the vertex connectivity device array\n vtconnoff = zeros(Int, vtmax + 1)\n vtconnoff[1] = 1\n temp = Int64[]\n for vt in 1:vtmax\n nconn = length(vtconn[vt])\n if nconn > 2\n vtconnoff[vt + 1] = vtconnoff[vt] + nconn\n append!(temp, vtconn[vt])\n else\n vtconnoff[vt + 1] = vtconnoff[vt]\n end\n end\n vtconn = temp\n #---setup face DSS---------------------\n fcmarker = -ones(Int, nface, length(elems))\n fcno = 0\n for el in realelems\n for fc in 1:nface\n if elemtobndy[fc, el] == 0\n nabrel = elemtoelem[fc, el]\n nabrfc = elemtoface[fc, el]\n if fcmarker[fc, el] == -1 && fcmarker[nabrfc, nabrel] == -1\n fcno += 1\n fcmarker[fc, el] = fcno\n fcmarker[nabrfc, nabrel] = fcno\n end\n end\n end\n end\n fcconn = -ones(Int, fcno, 5)\n for el in realelems\n for fc in 1:nface\n fcno = fcmarker[fc, el]\n ordr = elemtoordr[fc, el]\n if fcno ≠ -1\n nabrel = elemtoelem[fc, el]\n nabrfc = elemtoface[fc, el]\n fcconn[fcno, 1] = fc\n fcconn[fcno, 2] = el\n fcconn[fcno, 3] = nabrfc\n fcconn[fcno, 4] = nabrel\n fcconn[fcno, 5] = ordr\n fcmarker[fc, el] = -1\n fcmarker[nabrfc, nabrel] = -1\n end\n end\n end\n #---setup edge DSS---------------------\n if dim == 3\n nedge = 12\n edgemask = [\n 1 3 5 7 1 2 5 6 1 2 3 4\n 2 4 6 8 3 4 7 8 5 6 7 8\n ]\n edges = Array{Int64}(undef, 6, nedge * length(elems))\n ledge = zeros(Int64, 2)\n uedge = Dict{Array{Int64, 1}, Int64}()\n orient = 1\n uedgno = Int64(0)\n ctr = 1\n for el in elems\n for edg in 1:nedge\n orient = 1\n for i in 1:2\n ledge[i] = elemtouvert[edgemask[i, edg], el]\n edges[i, ctr] = ledge[i] # edge vertices [1:2]\n end\n if ledge[1] > ledge[2]\n sort!(ledge)\n orient = 2\n end\n edges[3, ctr] = edg # local edge number\n edges[4, ctr] = orient # edge orientation\n edges[5, ctr] = el # element number\n if haskey(uedge, ledge[:])\n edges[6, ctr] = uedge[ledge[:]] # unique edge number\n else\n uedgno += 1\n uedge[ledge[:]] = uedgno\n edges[6, ctr] = uedgno\n end\n ctr += 1\n end\n end\n edgconn = map(j -> zeros(Int, j), zeros(Int, uedgno))\n ctr = 1\n for el in elems\n for edg in 1:nedge\n uedg = edges[6, ctr]\n push!(edgconn[uedg], edg) # local edge number\n push!(edgconn[uedg], edges[4, ctr]) # orientation\n push!(edgconn[uedg], el) # element number\n ctr += 1\n end\n end\n # remove edges belonging to a single element and edges\n # belonging exclusively to ghost elements\n # and build edge connectivity device array\n edgconnoff = Int64[]\n temp = Int64[]\n push!(edgconnoff, 1)\n for i in 1:uedgno\n elen = length(edgconn[i])\n encls = Int(elen / 3)\n edgmrk = true\n if encls > 1\n for j in 1:encls\n if edgconn[i][3 * j] ≤ nreal\n edgmrk = false\n end\n end\n end\n if edgmrk\n edgconn[i] = []\n else\n shift = edgconnoff[end]\n push!(edgconnoff, (shift + length(edgconn[i])))\n append!(temp, edgconn[i][:])\n end\n end\n edgconn = temp\n else\n edgconn = nothing\n edgconnoff = nothing\n end\n else\n elemtouvert = nothing\n vtconn = nothing\n vtconnoff = nothing\n fcconn = nothing\n edgconn = nothing\n edgconnoff = nothing\n end\n #---------------\n StackedBrickTopology{dim, T, nb}(\n BoxElementTopology{dim, T, nb}(\n mpicomm,\n elems,\n realelems,\n ghostelems,\n ghostfaces,\n sendelems,\n sendfaces,\n elemtocoord,\n elemtoelem,\n elemtoface,\n elemtoordr,\n elemtobndy,\n nabrtorank,\n nabrtorecv,\n nabrtosend,\n basetopo.origsendorder,\n bndytoelem,\n bndytoface,\n elemtouvert = elemtouvert,\n vtconn = vtconn,\n vtconnoff = vtconnoff,\n fcconn = fcconn,\n edgconn = edgconn,\n edgconnoff = edgconnoff,\n ),\n stacksize,\n periodicity[end],\n )\nend\n\n\"\"\"\n CubedShellTopology(mpicomm, Nelem, T) <: AbstractTopology{dim,T,nb}\n\nGenerate a cubed shell mesh with the number of elements along each dimension of\nthe cubes being `Nelem`. This topology actually creates a cube mesh, and the\nwarping should be done after the grid is created using the `cubed_sphere_warp`\nfunction. The coordinates of the points will be of type `T`.\n\nThe elements of the shell are partitioned equally across the MPI ranks based\non a space-filling curve.\n\nNote that this topology is logically 2-D but embedded in a 3-D space\n\n# Examples\n\nWe can build a cubed shell mesh with 10*10 elements on each cube face, i.e., the\ntotal number of elements is\n`10 * 10 * 6 = 600`, with\n```jldoctest brickmesh\nusing ClimateMachine.Topologies\nusing MPI\nMPI.Init()\ntopology = CubedShellTopology(MPI.COMM_SELF, 10, Float64)\n\n# Typically the warping would be done after the grid is created, but the cell\n# corners could be warped with...\n\n# Shell radius = 1\nx1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)\nfor n = 1:length(x1)\n x1[n], x2[n], x3[n] = Topologies.cubed_sphere_warp(EquiangularCubedSphere(), x1[n], x2[n], x3[n])\nend\n\n# in case a unitary equiangular cubed sphere is desired, or\n\n# Shell radius = 10\nx1, x2, x3 = ntuple(j->topology.elemtocoord[j, :, :], 3)\nfor n = 1:length(x1)\n x1[n], x2[n], x3[n] = Topologies.cubed_sphere_warp(EquidistantCubedSphere(), x1[n], x2[n], x3[n], 10)\nend\n\n# in case an equidistant cubed sphere of radius 10 is desired.\n```\n\"\"\"\nfunction CubedShellTopology(\n mpicomm,\n Neside,\n T;\n connectivity = :full,\n ghostsize = 1,\n)\n\n # We cannot handle anything else right now...\n @assert ghostsize == 1\n\n mpirank = MPI.Comm_rank(mpicomm)\n mpisize = MPI.Comm_size(mpicomm)\n\n topology = cubedshellmesh(Neside, part = mpirank + 1, numparts = mpisize)\n\n topology = BrickMesh.partition(mpicomm, topology...)\n origsendorder = topology[5]\n dim, nvert = 3, 4\n elemtovert = topology[1]\n nelem = size(elemtovert, 2)\n elemtocoord = Array{T}(undef, dim, nvert, nelem)\n ind2vert = CartesianIndices((Neside + 1, Neside + 1, Neside + 1))\n for e in 1:nelem\n for n in 1:nvert\n v = elemtovert[n, e]\n i, j, k = Tuple(ind2vert[v])\n elemtocoord[:, n, e] =\n (2 * [i - 1, j - 1, k - 1] .- Neside) / Neside\n end\n end\n\n topology =\n connectivity == :face ?\n BrickMesh.connectmesh(\n mpicomm,\n topology[1],\n elemtocoord,\n topology[3],\n topology[4],\n dim = 2,\n ) :\n BrickMesh.connectmeshfull(\n mpicomm,\n topology[1],\n elemtocoord,\n topology[3],\n topology[4],\n dim = 2,\n )\n\n CubedShellTopology{T}(BoxElementTopology{2, T, 0}(\n mpicomm,\n topology.elems,\n topology.realelems,\n topology.ghostelems,\n topology.ghostfaces,\n topology.sendelems,\n topology.sendfaces,\n topology.elemtocoord,\n topology.elemtoelem,\n topology.elemtoface,\n topology.elemtoordr,\n topology.elemtobndy,\n topology.nabrtorank,\n topology.nabrtorecv,\n topology.nabrtosend,\n origsendorder,\n (),\n (),\n elemtouvert = topology.elemtovert,\n ))\nend\n\n\"\"\"\n cubedshellmesh(T, Ne; part=1, numparts=1)\n\nGenerate a cubed mesh where each of the \"cubes\" has an `Ne X Ne` grid of\nelements.\n\nThe mesh can optionally be partitioned into `numparts` and this returns\npartition `part`. This is a simple Cartesian partition and further partitioning\n(e.g, based on a space-filling curve) should be done before the mesh is used for\ncomputation.\n\nThis mesh returns the cubed spehere in a flattened fashion for the vertex values,\nand a remapping is needed to embed the mesh in a 3-D space.\n\nThe mesh structures for the cubes is as follows:\n\n```\nx2\n ^\n |\n4Ne- +-------+\n | | |\n | | 6 |\n | | |\n3Ne- +-------+\n | | |\n | | 5 |\n | | |\n2Ne- +-------+\n | | |\n | | 4 |\n | | |\n Ne- +-------+-------+-------+\n | | | | |\n | | 1 | 2 | 3 |\n | | | | |\n 0- +-------+-------+-------+\n |\n +---|-------|-------|------|-> x1\n 0 Ne 2Ne 3Ne\n```\n\n\"\"\"\nfunction cubedshellmesh(Ne; part = 1, numparts = 1)\n dim = 2\n @assert 1 <= part <= numparts\n\n globalnelems = 6 * Ne^2\n\n # How many vertices and faces per element\n nvert = 2^dim # 4\n nface = 2dim # 4\n\n # linearly partition to figure out which elements we own\n elemlocal = BrickMesh.linearpartition(prod(globalnelems), part, numparts)\n\n # elemen to vertex maps which we own\n elemtovert = Array{Int}(undef, nvert, length(elemlocal))\n elemtocoord = Array{Int}(undef, dim, nvert, length(elemlocal))\n\n nelemcube = Ne^dim # Ne^2\n\n etoijb = CartesianIndices((Ne, Ne, 6))\n bx = [0 Ne 2Ne Ne Ne Ne]\n by = [0 0 0 Ne 2Ne 3Ne]\n\n vertmap = LinearIndices((Ne + 1, Ne + 1, Ne + 1))\n for (le, e) in enumerate(elemlocal)\n i, j, blck = Tuple(etoijb[e])\n elemtocoord[1, :, le] = bx[blck] .+ [i - 1 i i - 1 i]\n elemtocoord[2, :, le] = by[blck] .+ [j - 1 j - 1 j j]\n\n for n in 1:4\n ix = i + mod(n - 1, 2)\n jx = j + div(n - 1, 2)\n # set the vertices like they are the face vertices of a cube\n if blck == 1\n elemtovert[n, le] = vertmap[1, Ne + 2 - ix, jx]\n elseif blck == 2\n elemtovert[n, le] = vertmap[ix, 1, jx]\n elseif blck == 3\n elemtovert[n, le] = vertmap[Ne + 1, ix, jx]\n elseif blck == 4\n elemtovert[n, le] = vertmap[ix, jx, Ne + 1]\n elseif blck == 5\n elemtovert[n, le] = vertmap[ix, Ne + 1, Ne + 2 - jx]\n elseif blck == 6\n elemtovert[n, le] = vertmap[ix, Ne + 2 - jx, 1]\n end\n end\n end\n\n # no boundaries for a shell\n elemtobndy = zeros(Int, nface, length(elemlocal))\n\n # no faceconnections for a shell\n faceconnections = Array{Array{Int, 1}}(undef, 0)\n\n (elemtovert, elemtocoord, elemtobndy, faceconnections, collect(elemlocal))\nend\n\nabstract type AbstractCubedSphere end\nstruct EquiangularCubedSphere <: AbstractCubedSphere end\nstruct EquidistantCubedSphere <: AbstractCubedSphere end\nstruct ConformalCubedSphere <: AbstractCubedSphere end\n\n\"\"\"\n cubed_sphere_warp(::EquiangularCubedSphere, a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nGiven points `(a, b, c)` on the surface of a cube, warp the points out to a\nspherical shell of radius `R` based on the equiangular gnomonic grid proposed by\n[Ronchi1996](@cite)\n\"\"\"\nfunction cubed_sphere_warp(\n ::EquiangularCubedSphere,\n a,\n b,\n c,\n R = max(abs(a), abs(b), abs(c)),\n)\n\n function f(sR, ξ, η)\n X, Y = tan(π * ξ / 4), tan(π * η / 4)\n ζ1 = sR / sqrt(X^2 + Y^2 + 1)\n ζ2, ζ3 = X * ζ1, Y * ζ1\n ζ1, ζ2, ζ3\n end\n\n fdim = argmax(abs.((a, b, c)))\n if fdim == 1 && a < 0\n # (-R, *, *) : formulas for Face I from Ronchi, Iacono, Paolucci (1996)\n # but for us face II of the developed net of the cube\n x1, x2, x3 = f(-R, b / a, c / a)\n elseif fdim == 2 && b < 0\n # ( *,-R, *) : formulas for Face II from Ronchi, Iacono, Paolucci (1996)\n # but for us face III of the developed net of the cube\n x2, x1, x3 = f(-R, a / b, c / b)\n elseif fdim == 1 && a > 0\n # ( R, *, *) : formulas for Face III from Ronchi, Iacono, Paolucci (1996)\n # but for us face IV of the developed net of the cube\n x1, x2, x3 = f(R, b / a, c / a)\n elseif fdim == 2 && b > 0\n # ( *, R, *) : formulas for Face IV from Ronchi, Iacono, Paolucci (1996)\n # but for us face I of the developed net of the cube\n x2, x1, x3 = f(R, a / b, c / b)\n elseif fdim == 3 && c > 0\n # ( *, *, R) : formulas for Face V from Ronchi, Iacono, Paolucci (1996)\n # and the same for us on the developed net of the cube\n x3, x2, x1 = f(R, b / c, a / c)\n elseif fdim == 3 && c < 0\n # ( *, *,-R) : formulas for Face VI from Ronchi, Iacono, Paolucci (1996)\n # and the same for us on the developed net of the cube\n x3, x2, x1 = f(-R, b / c, a / c)\n else\n error(\"invalid case for cubed_sphere_warp(::EquiangularCubedSphere): $a, $b, $c\")\n end\n\n return x1, x2, x3\nend\n\n\"\"\"\n equiangular_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nA wrapper function for the cubed_sphere_warp function, when called with the\nEquiangularCubedSphere type\n\"\"\"\nequiangular_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c))) =\n cubed_sphere_warp(EquiangularCubedSphere(), a, b, c, R)\n\n\"\"\"\n cubed_sphere_unwarp(x1, x2, x3)\n\nThe inverse of [`cubed_sphere_warp`](@ref). This function projects\na given point `(x_1, x_2, x_3)` from the surface of a sphere onto a cube\n\"\"\"\nfunction cubed_sphere_unwarp(::EquiangularCubedSphere, x1, x2, x3)\n\n function g(R, X, Y)\n ξ = atan(X) * 4 / pi\n η = atan(Y) * 4 / pi\n R, R * ξ, R * η\n end\n\n R = hypot(x1, x2, x3)\n fdim = argmax(abs.((x1, x2, x3)))\n\n if fdim == 1 && x1 < 0\n # (-R, *, *) : formulas for Face I from Ronchi, Iacono, Paolucci (1996)\n # but for us face II of the developed net of the cube\n a, b, c = g(-R, x2 / x1, x3 / x1)\n elseif fdim == 2 && x2 < 0\n # ( *,-R, *) : formulas for Face II from Ronchi, Iacono, Paolucci (1996)\n # but for us face III of the developed net of the cube\n b, a, c = g(-R, x1 / x2, x3 / x2)\n elseif fdim == 1 && x1 > 0\n # ( R, *, *) : formulas for Face III from Ronchi, Iacono, Paolucci (1996)\n # but for us face IV of the developed net of the cube\n a, b, c = g(R, x2 / x1, x3 / x1)\n elseif fdim == 2 && x2 > 0\n # ( *, R, *) : formulas for Face IV from Ronchi, Iacono, Paolucci (1996)\n # but for us face I of the developed net of the cube\n b, a, c = g(R, x1 / x2, x3 / x2)\n elseif fdim == 3 && x3 > 0\n # ( *, *, R) : formulas for Face V from Ronchi, Iacono, Paolucci (1996)\n # and the same for us on the developed net of the cube\n c, b, a = g(R, x2 / x3, x1 / x3)\n elseif fdim == 3 && x3 < 0\n # ( *, *,-R) : formulas for Face VI from Ronchi, Iacono, Paolucci (1996)\n # and the same for us on the developed net of the cube\n c, b, a = g(-R, x2 / x3, x1 / x3)\n else\n error(\"invalid case for cubed_sphere_unwarp(::EquiangularCubedSphere): $a, $b, $c\")\n end\n\n return a, b, c\nend\n\n\"\"\"\n equiangular_cubed_sphere_unwarp(x1, x2, x3)\n\nA wrapper function for the cubed_sphere_unwarp function, when called with the\nEquiangularCubedSphere type\n\"\"\"\nequiangular_cubed_sphere_unwarp(x1, x2, x3) =\n cubed_sphere_unwarp(EquiangularCubedSphere(), x1, x2, x3)\n\n\"\"\"\n cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nGiven points `(a, b, c)` on the surface of a cube, warp the points out to a\nspherical shell of radius `R` based on the equidistant gnomonic grid outlined in\n[Rancic1996](@cite) and [Nair2005](@cite)\n\"\"\"\nfunction cubed_sphere_warp(\n ::EquidistantCubedSphere,\n a,\n b,\n c,\n R = max(abs(a), abs(b), abs(c)),\n)\n\n r = hypot(a, b, c)\n\n x1 = R * a / r\n x2 = R * b / r\n x3 = R * c / r\n\n return x1, x2, x3\nend\n\n\"\"\"\n equidistant_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nA wrapper function for the cubed_sphere_warp function, when called with the\nEquidistantCubedSphere type\n\"\"\"\nequidistant_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c))) =\n cubed_sphere_warp(EquidistantCubedSphere(), a, b, c, R)\n\n\"\"\"\n cubed_sphere_unwarp(x1, x2, x3)\n\nThe inverse of [`cubed_sphere_warp`](@ref). This function projects\na given point `(x_1, x_2, x_3)` from the surface of a sphere onto a cube\n\"\"\"\nfunction cubed_sphere_unwarp(::EquidistantCubedSphere, x1, x2, x3)\n\n r = hypot(1, x2 / x1, x3 / x1)\n R = hypot(x1, x2, x3)\n\n a = r * x1\n b = r * x2\n c = r * x3\n\n m = max(abs(a), abs(b), abs(c))\n\n return a * R / m, b * R / m, c * R / m\nend\n\n\"\"\"\n equidistant_cubed_sphere_unwarp(x1, x2, x3)\n\nA wrapper function for the cubed_sphere_unwarp function, when called with the\nEquidistantCubedSphere type\n\"\"\"\nequidistant_cubed_sphere_unwarp(x1, x2, x3) =\n cubed_sphere_unwarp(EquidistantCubedSphere(), x1, x2, x3)\n\n\"\"\"\n cubed_sphere_warp(::ConformalCubedSphere, a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nGiven points `(a, b, c)` on the surface of a cube, warp the points out to a\nspherical shell of radius `R` based on the conformal grid proposed by\n[Rancic1996](@cite)\n\"\"\"\nfunction cubed_sphere_warp(\n ::ConformalCubedSphere,\n a,\n b,\n c,\n R = max(abs(a), abs(b), abs(c)),\n)\n\n fdim = argmax(abs.((a, b, c)))\n M = max(abs.((a, b, c))...)\n if fdim == 1 && a < 0\n # left face\n x1, x2, x3 = conformal_cubed_sphere_mapping(-b / M, c / M)\n x1, x2, x3 = RotX(π / 2) * RotY(-π / 2) * [x1, x2, x3]\n elseif fdim == 2 && b < 0\n # front face\n x1, x2, x3 = conformal_cubed_sphere_mapping(a / M, c / M)\n x1, x2, x3 = RotX(π / 2) * [x1, x2, x3]\n elseif fdim == 1 && a > 0\n # right face\n x1, x2, x3 = conformal_cubed_sphere_mapping(b / M, c / M)\n x1, x2, x3 = RotX(π / 2) * RotY(π / 2) * [x1, x2, x3]\n elseif fdim == 2 && b > 0\n # back face\n x1, x2, x3 = conformal_cubed_sphere_mapping(a / M, -c / M)\n x1, x2, x3 = RotX(-π / 2) * [x1, x2, x3]\n elseif fdim == 3 && c > 0\n # top face\n x1, x2, x3 = conformal_cubed_sphere_mapping(a / M, b / M)\n elseif fdim == 3 && c < 0\n # bottom face\n x1, x2, x3 = conformal_cubed_sphere_mapping(a / M, -b / M)\n x1, x2, x3 = RotX(π) * [x1, x2, x3]\n else\n error(\"invalid case for cubed_sphere_warp(::ConformalCubedSphere): $a, $b, $c\")\n end\n\n return x1 * R, x2 * R, x3 * R\nend\n\n\"\"\"\n conformal_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\nA wrapper function for the cubed_sphere_warp function, when called with the\nConformalCubedSphere type\n\"\"\"\nconformal_cubed_sphere_warp(a, b, c, R = max(abs(a), abs(b), abs(c))) =\n cubed_sphere_warp(ConformalCubedSphere(), a, b, c, R)\n\n\"\"\"\n cubed_sphere_unwarp(::ConformalCubedSphere, x1, x2, x3)\n\nThe inverse of [`cubed_sphere_warp`](@ref). This function projects\na given point `(x_1, x_2, x_3)` from the surface of a sphere onto a cube\n[Rancic1996](@cite)\n\"\"\"\nfunction cubed_sphere_unwarp(::ConformalCubedSphere, x1, x2, x3)\n\n # Auxuliary function that flips coordinates, if needed, to prepare input\n # arguments in the correct quadrant for the `conformal_cubed_sphere_inverse_mapping`\n # function. Then, flips the output of `conformal_cubed_sphere_inverse_mapping`\n # back to original face and scales the coordinates so that result is on the cube\n function flip_unwarp_scale(x1, x2, x3)\n R = hypot(x1, x2, x3)\n flipx1, flipx2 = false, false\n if x1 < 0 # flip the point around x2 axis\n x1 = -x1\n flipx1 = true\n end\n if x2 < 0 # flip the point around x1 axis\n x2 = -x2\n flipx2 = true\n end\n a, b = conformal_cubed_sphere_inverse_mapping(x1 / R, x2 / R, x3 / R)\n if flipx1 == true\n a = -a\n end\n if flipx2 == true\n b = -b\n end\n\n # Rescale to desired length\n a *= R\n b *= R\n # Since we were trating coordinates on top face of the cube, the c\n # coordinate must have the top-face z value (z = R)\n c = R\n\n return a, b, c\n end\n\n fdim = argmax(abs.((x1, x2, x3)))\n if fdim == 1 && x1 < 0 # left face\n # rotate to align with top face\n x1, x2, x3 = RotY(π / 2) * RotX(-π / 2) * [x1, x2, x3]\n # call the unwarp function, with appropriate flipping and scaling\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n # rotate back\n a, b, c = RotX(π / 2) * RotY(-π / 2) * [a, b, c]\n elseif fdim == 2 && x2 < 0 # front face\n # rotate to align with top face\n x1, x2, x3 = RotX(-π / 2) * [x1, x2, x3]\n # call the unwarp function, with appropriate flipping and scaling\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n # rotate back\n a, b, c = RotX(π / 2) * [a, b, c]\n elseif fdim == 1 && x1 > 0 # right face\n # rotate to align with top face\n x1, x2, x3 = RotZ(-π / 2) * RotY(-π / 2) * [x1, x2, x3]\n # call the unwarp function, with appropriate flipping and scaling\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n # rotate back\n a, b, c = RotY(π / 2) * RotZ(π / 2) * [a, b, c]\n elseif fdim == 2 && x2 > 0 # back face\n # rotate to align with top face\n x1, x2, x3 = RotZ(π) * RotX(π / 2) * [x1, x2, x3]\n # call the unwarp function, with appropriate flipping and scaling\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n # rotate back\n a, b, c = RotX(-π / 2) * RotZ(-π) * [a, b, c]\n elseif fdim == 3 && x3 > 0 # top face\n # already on top face, no need to rotate\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n elseif fdim == 3 && x3 < 0 # bottom face\n # rotate to align with top face\n x1, x2, x3 = RotX(π) * [x1, x2, x3]\n # call the unwarp function, with appropriate flipping and scaling\n a, b, c = flip_unwarp_scale(x1, x2, x3)\n # rotate back\n a, b, c = RotX(-π) * [a, b, c]\n else\n error(\"invalid case for cubed_sphere_unwarp(::ConformalCubedSphere): $a, $b, $c\")\n end\n\n return a, b, c\nend\n\n\"\"\"\n conformal_cubed_sphere_unwarp(x1, x2, x3)\n\nA wrapper function for the cubed_sphere_unwarp function, when called with the\nConformalCubedSphere type\n\"\"\"\nconformal_cubed_sphere_unwarp(x1, x2, x3) =\n cubed_sphere_unwarp(ConformalCubedSphere(), x1, x2, x3)\n\n\"\"\"\n StackedCubedSphereTopology(mpicomm, Nhorz, Rrange;\n boundary=(1,1)) <: AbstractTopology{3}\n\nGenerate a stacked cubed sphere topology with `Nhorz` by `Nhorz` cells for each\nhorizontal face and `Rrange` is the radius edges of the stacked elements. This\ntopology actual creates a cube mesh, and the warping should be done after the\ngrid is created using the `cubed_sphere_warp` function. The coordinates of the\npoints will be of type `eltype(Rrange)`. The inner boundary condition type is\n`boundary[1]` and the outer boundary condition type is `boundary[2]`.\n\nThe elements are stacked such that the vertical elements are contiguous in the\nelement ordering.\n\nThe elements of the brick are partitioned equally across the MPI ranks based\non a space-filling curve. Further, stacks are not split at MPI boundaries.\n\n# Examples\n\nWe can build a cubed sphere mesh with 10 x 10 x 5 elements on each cube face,\ni.e., the total number of elements is `10 * 10 * 5 * 6 = 3000`, with\n```jldoctest brickmesh\nusing ClimateMachine.Topologies\nusing MPI\nMPI.Init()\nNhorz = 10\nNstack = 5\nRrange = Float64.(accumulate(+,1:Nstack+1))\ntopology = StackedCubedSphereTopology(MPI.COMM_SELF, Nhorz, Rrange)\n\nx1, x2, x3 = ntuple(j->reshape(topology.elemtocoord[j, :, :],\n 2, 2, 2, length(topology.elems)), 3)\nfor n = 1:length(x1)\n x1[n], x2[n], x3[n] = Topologies.cubed_sphere_warp(EquiangularCubedSphere(),x1[n], x2[n], x3[n])\nend\n```\nNote that the faces are listed in Cartesian order.\n\"\"\"\nfunction StackedCubedSphereTopology(\n mpicomm,\n Nhorz,\n Rrange;\n boundary = (1, 1),\n connectivity = :full,\n ghostsize = 1,\n)\n T = eltype(Rrange)\n\n basetopo = CubedShellTopology(\n mpicomm,\n Nhorz,\n T;\n connectivity = connectivity,\n ghostsize = ghostsize,\n )\n\n dim = 3\n nvert = 2^dim\n nface = 2dim\n stacksize = length(Rrange) - 1\n\n nreal = length(basetopo.realelems) * stacksize\n nghost = length(basetopo.ghostelems) * stacksize\n\n elems = 1:(nreal + nghost)\n realelems = 1:nreal\n ghostelems = nreal .+ (1:nghost)\n\n sendelems =\n similar(basetopo.sendelems, length(basetopo.sendelems) * stacksize)\n\n for i in 1:length(basetopo.sendelems), j in 1:stacksize\n sendelems[stacksize * (i - 1) + j] =\n stacksize * (basetopo.sendelems[i] - 1) + j\n end\n\n ghostfaces = similar(basetopo.ghostfaces, nface, length(ghostelems))\n ghostfaces .= false\n\n for i in 1:length(basetopo.ghostelems), j in 1:stacksize\n e = stacksize * (i - 1) + j\n for f in 1:(2 * (dim - 1))\n ghostfaces[f, e] = basetopo.ghostfaces[f, i]\n end\n end\n\n sendfaces = similar(basetopo.sendfaces, nface, length(sendelems))\n sendfaces .= false\n\n for i in 1:length(basetopo.sendelems), j in 1:stacksize\n e = stacksize * (i - 1) + j\n for f in 1:(2 * (dim - 1))\n sendfaces[f, e] = basetopo.sendfaces[f, i]\n end\n end\n\n elemtocoord = similar(basetopo.elemtocoord, dim, nvert, length(elems))\n\n for i in 1:length(basetopo.elems), j in 1:stacksize\n # i is base element\n # e is stacked element\n e = stacksize * (i - 1) + j\n\n\n # v is base vertex\n for v in 1:(2^(dim - 1))\n for d in 1:dim # dim here since shell is embedded in 3-D\n # v is lower stacked vertex\n elemtocoord[d, v, e] = basetopo.elemtocoord[d, v, i] * Rrange[j]\n # 2^(dim-1) + v is higher stacked vertex\n elemtocoord[d, 2^(dim - 1) + v, e] =\n basetopo.elemtocoord[d, v, i] * Rrange[j + 1]\n end\n end\n end\n\n elemtoelem = similar(basetopo.elemtoelem, nface, length(elems))\n elemtoface = similar(basetopo.elemtoface, nface, length(elems))\n elemtoordr = similar(basetopo.elemtoordr, nface, length(elems))\n elemtobndy = similar(basetopo.elemtobndy, nface, length(elems))\n\n for e in 1:(length(basetopo.elems) * stacksize), f in 1:nface\n elemtoelem[f, e] = e\n elemtoface[f, e] = f\n elemtoordr[f, e] = 1\n elemtobndy[f, e] = 0\n end\n\n for i in 1:length(basetopo.realelems), j in 1:stacksize\n e1 = stacksize * (i - 1) + j\n\n for f in 1:(2 * (dim - 1))\n e2 = stacksize * (basetopo.elemtoelem[f, i] - 1) + j\n\n elemtoelem[f, e1] = e2\n elemtoface[f, e1] = basetopo.elemtoface[f, i]\n\n # since the basetopo is 2-D we only need to worry about two\n # orientations\n @assert basetopo.elemtoordr[f, i] ∈ (1, 2)\n #=\n orientation 1:\n 2---3 2---3\n | | --> | |\n 0---1 0---1\n same:\n (a,b) --> (a,b)\n\n orientation 3:\n 2---3 3---2\n | | --> | |\n 0---1 1---0\n reverse first index:\n (a,b) --> (N+1-a,b)\n =#\n elemtoordr[f, e1] = basetopo.elemtoordr[f, i] == 1 ? 1 : 3\n end\n\n # If top or bottom of stack set neighbor to self on respective face\n elemtoelem[2 * (dim - 1) + 1, e1] =\n j == 1 ? e1 : stacksize * (i - 1) + j - 1\n elemtoelem[2 * (dim - 1) + 2, e1] =\n j == stacksize ? e1 : stacksize * (i - 1) + j + 1\n\n elemtoface[2 * (dim - 1) + 1, e1] =\n j == 1 ? 2 * (dim - 1) + 1 : 2 * (dim - 1) + 2\n elemtoface[2 * (dim - 1) + 2, e1] =\n j == stacksize ? 2 * (dim - 1) + 2 : 2 * (dim - 1) + 1\n\n elemtoordr[2 * (dim - 1) + 1, e1] = 1\n elemtoordr[2 * (dim - 1) + 2, e1] = 1\n end\n\n # Set the top and bottom boundary condition\n for i in 1:length(basetopo.elems)\n eb = stacksize * (i - 1) + 1\n et = stacksize * (i - 1) + stacksize\n\n elemtobndy[2 * (dim - 1) + 1, eb] = boundary[1]\n elemtobndy[2 * (dim - 1) + 2, et] = boundary[2]\n end\n\n nabrtorank = basetopo.nabrtorank\n nabrtorecv = UnitRange{Int}[\n UnitRange(\n stacksize * (first(basetopo.nabrtorecv[n]) - 1) + 1,\n stacksize * last(basetopo.nabrtorecv[n]),\n ) for n in 1:length(nabrtorank)\n ]\n nabrtosend = UnitRange{Int}[\n UnitRange(\n stacksize * (first(basetopo.nabrtosend[n]) - 1) + 1,\n stacksize * last(basetopo.nabrtosend[n]),\n ) for n in 1:length(nabrtorank)\n ]\n\n bndytoelem, bndytoface = BrickMesh.enumerateboundaryfaces!(\n elemtoelem,\n elemtobndy,\n (false,),\n (boundary,),\n )\n nb = length(bndytoelem)\n #----setting up DSS----------\n if basetopo.elemtouvert isa Array\n #---setup vertex DSS\n nvertby2 = div(nvert, 2)\n elemtouvert = similar(basetopo.elemtouvert, nvert, length(elems))\n # base level\n for i in 1:length(basetopo.elems)\n e = stacksize * (i - 1) + 1\n for vt in 1:nvertby2\n elemtouvert[vt, e] =\n (basetopo.elemtouvert[vt, i] - 1) * (stacksize + 1) + 1\n end\n end\n # interior levels\n for i in 1:length(basetopo.elems), j in 1:(stacksize - 1)\n e = stacksize * (i - 1) + j\n for vt in 1:nvertby2\n elemtouvert[nvertby2 + vt, e] = elemtouvert[vt, e] + 1\n elemtouvert[vt, e + 1] = elemtouvert[nvertby2 + vt, e]\n end\n end\n # top level\n for i in 1:length(basetopo.elems)\n e = stacksize * i\n for vt in 1:nvertby2\n elemtouvert[nvertby2 + vt, e] = elemtouvert[vt, e] + 1\n end\n end\n vtmax = maximum(unique(elemtouvert))\n vtconn = map(j -> zeros(Int, j), zeros(Int, vtmax))\n\n for el in elems, lvt in 1:nvert\n vt = elemtouvert[lvt, el]\n push!(vtconn[vt], lvt) # local vertex number\n push!(vtconn[vt], el) # local elem number\n end\n # building the vertex connectivity device array\n vtconnoff = zeros(Int, vtmax + 1)\n vtconnoff[1] = 1\n temp = Int64[]\n for vt in 1:vtmax\n nconn = length(vtconn[vt])\n if nconn > 2\n vtconnoff[vt + 1] = vtconnoff[vt] + nconn\n append!(temp, vtconn[vt])\n else\n vtconnoff[vt + 1] = vtconnoff[vt]\n end\n end\n vtconn = temp\n #---setup face DSS---------------------\n fcmarker = -ones(Int, nface, length(elems))\n fcno = 0\n for el in realelems\n for fc in 1:nface\n if elemtobndy[fc, el] == 0\n nabrel = elemtoelem[fc, el]\n nabrfc = elemtoface[fc, el]\n if fcmarker[fc, el] == -1 && fcmarker[nabrfc, nabrel] == -1\n fcno += 1\n fcmarker[fc, el] = fcno\n fcmarker[nabrfc, nabrel] = fcno\n end\n end\n end\n end\n fcconn = -ones(Int, fcno, 5)\n for el in realelems\n for fc in 1:nface\n fcno = fcmarker[fc, el]\n ordr = elemtoordr[fc, el]\n if fcno ≠ -1\n nabrel = elemtoelem[fc, el]\n nabrfc = elemtoface[fc, el]\n fcconn[fcno, 1] = fc\n fcconn[fcno, 2] = el\n fcconn[fcno, 3] = nabrfc\n fcconn[fcno, 4] = nabrel\n fcconn[fcno, 5] = ordr\n fcmarker[fc, el] = -1\n fcmarker[nabrfc, nabrel] = -1\n end\n end\n end\n #---setup edge DSS---------------------\n if dim == 3\n nedge = 12\n edgemask = [\n 1 3 5 7 1 2 5 6 1 2 3 4\n 2 4 6 8 3 4 7 8 5 6 7 8\n ]\n edges = Array{Int64}(undef, 6, nedge * length(elems))\n ledge = zeros(Int64, 2)\n uedge = Dict{Array{Int64, 1}, Int64}()\n orient = 1\n uedgno = Int64(0)\n ctr = 1\n for el in elems\n for edg in 1:nedge\n orient = 1\n for i in 1:2\n ledge[i] = elemtouvert[edgemask[i, edg], el]\n edges[i, ctr] = ledge[i] # edge vertices [1:2]\n end\n if ledge[1] > ledge[2]\n sort!(ledge)\n orient = 2\n end\n edges[3, ctr] = edg # local edge number\n edges[4, ctr] = orient # edge orientation\n edges[5, ctr] = el # element number\n if haskey(uedge, ledge[:])\n edges[6, ctr] = uedge[ledge[:]] # unique edge number\n else\n uedgno += 1\n uedge[ledge[:]] = uedgno\n edges[6, ctr] = uedgno\n end\n ctr += 1\n end\n end\n edgconn = map(j -> zeros(Int, j), zeros(Int, uedgno))\n ctr = 1\n for el in elems\n for edg in 1:nedge\n uedg = edges[6, ctr]\n push!(edgconn[uedg], edg) # local edge number\n push!(edgconn[uedg], edges[4, ctr]) # orientation\n push!(edgconn[uedg], el) # element number\n ctr += 1\n end\n end\n # remove edges belonging to a single element and edges\n # belonging exclusively to ghost elements\n # and build edge connectivity device array\n edgconnoff = Int64[]\n temp = Int64[]\n push!(edgconnoff, 1)\n for i in 1:uedgno\n elen = length(edgconn[i])\n encls = Int(elen / 3)\n edgmrk = true\n if encls > 1\n for j in 1:encls\n if edgconn[i][3 * j] ≤ nreal\n edgmrk = false\n end\n end\n end\n if edgmrk\n edgconn[i] = []\n else\n shift = edgconnoff[end]\n push!(edgconnoff, (shift + length(edgconn[i])))\n append!(temp, edgconn[i][:])\n end\n end\n edgconn = temp\n else\n edgconn = nothing\n edgconnoff = nothing\n end\n #-------------------------------------\n else\n elemtouvert = nothing\n vtconn = nothing\n vtconnoff = nothing\n fcconn = nothing\n edgconn = nothing\n edgconnoff = nothing\n end\n #-----------------------------\n StackedCubedSphereTopology{T, nb}(\n BoxElementTopology{3, T, nb}(\n mpicomm,\n elems,\n realelems,\n ghostelems,\n ghostfaces,\n sendelems,\n sendfaces,\n elemtocoord,\n elemtoelem,\n elemtoface,\n elemtoordr,\n elemtobndy,\n nabrtorank,\n nabrtorecv,\n nabrtosend,\n basetopo.origsendorder,\n bndytoelem,\n bndytoface,\n elemtouvert = elemtouvert,\n vtconn = vtconn,\n vtconnoff = vtconnoff,\n fcconn = fcconn,\n edgconn = edgconn,\n edgconnoff = edgconnoff,\n ),\n stacksize,\n )\nend\n\n\"\"\"\n grid1d(a, b[, stretch::AbstractGridStretching]; elemsize, nelem)\n\nDiscretize the 1D interval [`a`,`b`] into elements.\nExactly one of the following keyword arguments must be provided:\n- `elemsize`: the average element size, or\n- `nelem`: the number of elements.\n\nThe optional `stretch` argument allows stretching, otherwise the element sizes\nwill be uniform.\n\nReturns either a range object or a vector containing the element boundaries.\n\"\"\"\nfunction grid1d(a, b, stretch = nothing; elemsize = nothing, nelem = nothing)\n xor(nelem === nothing, elemsize === nothing) ||\n error(\"Either `elemsize` or `nelem` arguments must be provided\")\n if elemsize !== nothing\n nelem = round(Int, abs(b - a) / elemsize)\n end\n grid1d(a, b, stretch, nelem)\nend\nfunction grid1d(a, b, ::Nothing, nelem)\n range(a, stop = b, length = nelem + 1)\nend\n\n# TODO: document these\nabstract type AbstractGridStretching end\n\n\"\"\"\n SingleExponentialStretching(A)\n\nApply single-exponential stretching: `A > 0` will increase the density of points\nat the lower boundary, `A < 0` will increase the density at the upper boundary.\n\n# Reference\n* \"Handbook of Grid Generation\" J. F. Thompson, B. K. Soni, N. P. Weatherill\n (Editors) RCR Press 1999, §3.6.1 Single-Exponential Function\n\"\"\"\nstruct SingleExponentialStretching{T} <: AbstractGridStretching\n A::T\nend\nfunction grid1d(\n a::A,\n b::B,\n stretch::SingleExponentialStretching,\n nelem,\n) where {A, B}\n F = float(promote_type(A, B))\n s = range(zero(F), stop = one(F), length = nelem + 1)\n a .+ (b - a) .* expm1.(stretch.A .* s) ./ expm1(stretch.A)\nend\n\nstruct InteriorStretching{T} <: AbstractGridStretching\n attractor::T\nend\nfunction grid1d(a::A, b::B, stretch::InteriorStretching, nelem) where {A, B}\n F = float(promote_type(A, B))\n coe = F(2.5)\n s = range(zero(F), stop = one(F), length = nelem + 1)\n range(a, stop = b, length = nelem + 1) .+\n coe .* (stretch.attractor .- (b - a) .* s) .* (1 .- s) .* s\nend\n\nfunction basic_topology_info(topology::AbstractStackedTopology)\n nelem = length(topology.elems)\n nvertelem = topology.stacksize\n nhorzelem = div(nelem, nvertelem)\n nrealelem = length(topology.realelems)\n nhorzrealelem = div(nrealelem, nvertelem)\n\n return (\n nelem = nelem,\n nvertelem = nvertelem,\n nhorzelem = nhorzelem,\n nrealelem = nrealelem,\n nhorzrealelem = nhorzrealelem,\n )\nend\n\nfunction basic_topology_info(topology::AbstractTopology)\n return (\n nelem = length(topology.elems),\n nrealelem = length(topology.realelems),\n )\nend\n\n\n### Helper Functions for Topography Calculations\n\"\"\"\n compute_lat_long(X,Y,δ,faceid)\nHelper function to allow computation of latitute and longitude coordinates\ngiven the cubed sphere coordinates X, Y, δ, faceid\n\"\"\"\nfunction compute_lat_long(X, Y, δ, faceid)\n if faceid == 1\n λ = atan(X) # longitude\n ϕ = atan(cos(λ) * Y) # latitude\n elseif faceid == 2\n λ = atan(X) + π / 2\n ϕ = atan(Y * cos(atan(X)))\n elseif faceid == 3\n λ = atan(X) + π\n ϕ = atan(Y * cos(atan(X)))\n elseif faceid == 4\n λ = atan(X) + (3 / 2) * π\n ϕ = atan(Y * cos(atan(X)))\n elseif faceid == 5\n λ = atan(X, -Y) + π\n ϕ = atan(1 / sqrt(δ - 1))\n elseif faceid == 6\n λ = atan(X, Y)\n ϕ = -atan(1 / sqrt(δ - 1))\n end\n return λ, ϕ\nend\n\n\n\"\"\"\n AnalyticalTopography\nAbstract type to allow dispatch over different analytical topography prescriptions\nin experiments.\n\"\"\"\nabstract type AnalyticalTopography end\n\nfunction compute_analytical_topography(\n ::AnalyticalTopography,\n λ,\n ϕ,\n sR,\n r_inner,\n r_outer,\n)\n return sR\nend\n\n\"\"\"\n NoTopography <: AnalyticalTopography\nAllows definition of fallback methods in case cubed_sphere_topo_warp is used with\nno prescribed topography function.\n\"\"\"\nstruct NoTopography <: AnalyticalTopography end\n\n### DCMIP Mountain\n\"\"\"\n DCMIPMountain <: AnalyticalTopography\nTopography description based on standard DCMIP experiments.\n\"\"\"\nstruct DCMIPMountain <: AnalyticalTopography end\nfunction compute_analytical_topography(\n ::DCMIPMountain,\n λ,\n ϕ,\n sR,\n r_inner,\n r_outer,\n)\n #User specified warp parameters\n R_m = π * 3 / 4\n h0 = 2000\n ζ_m = π / 16\n φ_m = 0\n λ_m = π * 3 / 2\n r_m = acos(sin(φ_m) * sin(ϕ) + cos(φ_m) * cos(ϕ) * cos(λ - λ_m))\n # Define mesh decay profile\n Δ = (r_outer - abs(sR)) / (r_outer - r_inner)\n if r_m < R_m\n zs =\n 0.5 *\n h0 *\n (1 + cos(π * r_m / R_m)) *\n cos(π * r_m / ζ_m) *\n cos(π * r_m / ζ_m)\n else\n zs = 0.0\n end\n mR = sign(sR) * (abs(sR) + zs * Δ)\n return mR\nend\n\n\"\"\"\n cubed_sphere_topo_warp(a, b, c, R = max(abs(a), abs(b), abs(c));\n r_inner = _planet_radius,\n r_outer = _planet_radius + domain_height,\n topography = NoTopography())\n\nGiven points `(a, b, c)` on the surface of a cube, warp the points out to a\nspherical shell of radius `R` based on the equiangular gnomonic grid proposed by\n[Ronchi1996](@cite). Assumes a user specified modified radius using the\ncompute_analytical_topography function. Defaults to smooth cubed sphere unless otherwise specified\nvia the AnalyticalTopography type.\n\"\"\"\nfunction cubed_sphere_topo_warp(\n a,\n b,\n c,\n R = max(abs(a), abs(b), abs(c));\n r_inner = _planet_radius,\n r_outer = _planet_radius + domain_height,\n topography::AnalyticalTopography = NoTopography(),\n)\n\n function f(sR, ξ, η, faceid)\n X, Y = tan(π * ξ / 4), tan(π * η / 4)\n δ = 1 + X^2 + Y^2\n λ, ϕ = compute_lat_long(X, Y, δ, faceid)\n mR = compute_analytical_topography(\n topography,\n λ,\n ϕ,\n sR,\n r_inner,\n r_outer,\n )\n x1 = mR / sqrt(δ)\n x2, x3 = X * x1, Y * x1\n x1, x2, x3\n end\n fdim = argmax(abs.((a, b, c)))\n\n if fdim == 1 && a < 0\n faceid = 1\n # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(-R, b / a, c / a, faceid)\n elseif fdim == 2 && b < 0\n faceid = 2\n # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(-R, a / b, c / b, faceid)\n elseif fdim == 1 && a > 0\n faceid = 3\n # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(R, b / a, c / a, faceid)\n elseif fdim == 2 && b > 0\n faceid = 4\n # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(R, a / b, c / b, faceid)\n elseif fdim == 3 && c > 0\n faceid = 5\n # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(R, b / c, a / c, faceid)\n elseif fdim == 3 && c < 0\n faceid = 6\n # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(-R, b / c, a / c, faceid)\n else\n error(\"invalid case for cubed_sphere_warp(::EquiangularCubedSphere): $a, $b, $c\")\n end\n\n return x1, x2, x3\nend\n\nend\n" }, { "path": "src/Numerics/ODESolvers/AdditiveRungeKuttaMethod.jl", "content": "export AbstractAdditiveRungeKutta\nexport LowStorageVariant, NaiveVariant\nexport AdditiveRungeKutta\nexport ARK1ForwardBackwardEuler\nexport ARK2ImplicitExplicitMidpoint\nexport ARK2GiraldoKellyConstantinescu\nexport ARK548L2SA2KennedyCarpenter, ARK437L2SA1KennedyCarpenter\nexport Trap2LockWoodWeller\nexport DBM453VoglEtAl\n\n# Naive formulation that uses equation 3.8 from Giraldo, Kelly, and\n# Constantinescu (2013) directly. Seems to cut the number of solver iterations\n# by half but requires Nstages - 1 additional storage.\nstruct NaiveVariant end\nadditional_storage(::NaiveVariant, Q, Nstages) =\n (Lstages = ntuple(i -> similar(Q), Nstages),)\n\n# Formulation that does things exactly as in Giraldo, Kelly, and Constantinescu\n# (2013). Uses only one additional vector of storage regardless of the number\n# of stages.\nstruct LowStorageVariant end\nadditional_storage(::LowStorageVariant, Q, Nstages) = (Qtt = similar(Q),)\n\nabstract type AbstractAdditiveRungeKutta <: AbstractODESolver end\n\n\"\"\"\n AdditiveRungeKutta(f, l, backward_euler_solver, RKAe, RKAi, RKB, RKC, Q;\n split_explicit_implicit, variant, dt, t0 = 0)\n\nThis is a time stepping object for implicit-explicit time stepping of a\ndecomposed differential equation. When `split_explicit_implicit == false`\nthe equation is assumed to be decomposed as\n\n```math\n \\\\dot{Q} = [l(Q, t)] + [f(Q, t) - l(Q, t)]\n```\n\nwhere `Q` is the state, `f` is the full tendency and\n`l` is the chosen implicit operator. When `split_explicit_implicit == true`\nthe assumed decomposition is\n\n```math\n \\\\dot{Q} = [l(Q, t)] + [f(Q, t)]\n```\n\nwhere `f` is now only the nonlinear tendency. For both decompositions the implicit\noperator `l` is integrated implicitly whereas the remaining part is integrated\nexplicitly. Other arguments are the required time step size `dt` and the\noptional initial time `t0`. The resulting backward Euler type systems are solved\nusing the provided `backward_euler_solver`. This time stepping object is\nintended to be passed to the `solve!` command.\n\nThe constructor builds an additive Runge--Kutta scheme based on the provided\n`RKAe`, `RKAi`, `RKB` and `RKC` coefficient arrays. Additionally `variant`\nspecifies which of the analytically equivalent but numerically different\nformulations of the scheme is used.\n\nThe available concrete implementations are:\n\n - [`ARK1ForwardBackwardEuler`](@ref)\n - [`ARK2ImplicitExplicitMidpoint`](@ref)\n - [`ARK2GiraldoKellyConstantinescu`](@ref)\n - [`ARK548L2SA2KennedyCarpenter`](@ref)\n - [`ARK437L2SA1KennedyCarpenter`](@ref)\n - [`Trap2LockWoodWeller`](@ref)\n - [`DBM453VoglEtAl`](@ref)\n\"\"\"\nmutable struct AdditiveRungeKutta{\n T,\n RT,\n AT,\n V,\n VS,\n IST,\n Nstages,\n Nstages_sq,\n Nstagesm1,\n} <: AbstractAdditiveRungeKutta\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n rhs!::Any\n \"rhs linear operator\"\n rhs_implicit!::Any\n \"a dictionary of backward Euler solvers\"\n implicit_solvers::IST\n \"An integer which is updated to determine the appropriate cached implicit\n solver (used in substepping)\"\n substep_stage::Int\n \"Storage for solution during the AdditiveRungeKutta update\"\n Qstages::NTuple{Nstagesm1, AT}\n \"Storage for RHS during the AdditiveRungeKutta update\"\n Rstages::NTuple{Nstages, AT}\n \"Storage for the linear solver rhs vector\"\n Qhat::AT\n \"RK coefficient matrix A for the explicit scheme\"\n RKA_explicit::SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}\n \"RK coefficient matrix A for the implicit scheme\"\n RKA_implicit::SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}\n \"RK coefficient vector B for the explicit scheme (rhs add in scaling)\"\n RKB_explicit::SArray{Tuple{Nstages}, RT, 1, Nstages}\n \"RK coefficient vector B for the implicit scheme (rhs add in scaling)\"\n RKB_implicit::SArray{Tuple{Nstages}, RT, 1, Nstages}\n \"RK_explicit coefficient vector C for the explicit scheme (time scaling)\"\n RKC_explicit::SArray{Tuple{Nstages}, RT, 1, Nstages}\n \"RK_implicit coefficient vector C for the implicit scheme (time scaling)\"\n RKC_implicit::SArray{Tuple{Nstages}, RT, 1, Nstages}\n split_explicit_implicit::Bool\n \"Variant of the ARK scheme\"\n variant::V\n \"Storage dependent on the variant of the ARK scheme\"\n variant_storage::VS\n\n function AdditiveRungeKutta(\n rhs!,\n rhs_implicit!,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n ) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n Nstages = length(RKB_explicit)\n\n Qstages = ntuple(i -> similar(Q), Nstages - 1)\n Rstages = ntuple(i -> similar(Q), Nstages)\n Qhat = similar(Q)\n\n V = typeof(variant)\n variant_storage = additional_storage(variant, Q, Nstages)\n VS = typeof(variant_storage)\n\n implicit_solvers = Dict()\n\n rk_diag = unique(diag(RKA_implicit))\n # Remove all zero entries from `rk_diag`\n # so we build all unique implicit solvers (parameterized by the\n # corresponding RK coefficient)\n filter!(c -> !iszero(c), rk_diag)\n\n # LowStorageVariant ARK methods assume that both the explicit and\n # implicit B and C vectors are the same. Additionally, the diagonal\n # of the implicit Butcher table A is assumed to have the form:\n # [0, c, ... c ], where c is some non-zero constant.\n if variant isa LowStorageVariant\n @assert RKB_explicit == RKB_implicit\n @assert RKC_explicit == RKC_implicit\n # rk_diag here has been filtered of all non-unique and zero values.\n # So [0, c, ... c ] filters to [c]. We error if its length is not 1\n @assert length(rk_diag) == 1\n end\n\n if isempty(nsubsteps)\n for rk_coeff in rk_diag\n α = dt * rk_coeff\n besolver! = setup_backward_Euler_solver(\n backward_euler_solver,\n Q,\n α,\n rhs_implicit!,\n )\n @assert besolver! isa AbstractBackwardEulerSolver\n implicit_solvers[rk_coeff] = (besolver!,)\n end\n else\n nsteps = length(nsubsteps)\n for rk_coeff in rk_diag\n solvers = ntuple(\n i -> setup_backward_Euler_solver(\n backward_euler_solver,\n Q,\n dt * nsubsteps[i] * rk_coeff,\n rhs_implicit!,\n ),\n nsteps,\n )\n @assert(all(isa.(solvers, AbstractBackwardEulerSolver)))\n implicit_solvers[rk_coeff] = solvers\n end\n end\n\n IST = typeof(implicit_solvers)\n new{T, RT, AT, V, VS, IST, Nstages, Nstages^2, Nstages - 1}(\n RT(dt),\n RT(t0),\n 0,\n rhs!,\n rhs_implicit!,\n implicit_solvers,\n # By default (no substepping, this parameter is simply set to 1)\n 1,\n Qstages,\n Rstages,\n Qhat,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n variant_storage,\n )\n end\nend\n\nfunction AdditiveRungeKutta(\n ark,\n op::TimeScaledRHS{2, RT} where {RT},\n backward_euler_solver,\n Q::AT;\n dt = 0,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = true,\n variant = NaiveVariant(),\n) where {AT <: AbstractArray}\n return ark(\n op.rhs![1],\n op.rhs![2],\n backward_euler_solver,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n split_explicit_implicit = split_explicit_implicit,\n variant = variant,\n )\nend\n\n# this will only work for iterative solves\n# direct solvers use prefactorization\nfunction updatedt!(ark::AdditiveRungeKutta, dt)\n for (rk_coeff, implicit_solvers) in ark.implicit_solvers\n implicit_solver! = implicit_solvers[ark.substep_stage]\n @assert Δt_is_adjustable(implicit_solver!)\n # New coefficient\n α = dt * rk_coeff\n # Update with new dt and implicit coefficient\n ark.dt = dt\n update_backward_Euler_solver!(implicit_solver!, ark.Qstages[1], α)\n end\nend\n\nfunction dostep!(\n Q,\n ark::AdditiveRungeKutta,\n p,\n time,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n dostep!(Q, ark, ark.variant, p, time, slow_δ, slow_rv_dQ, slow_scaling)\nend\n\nfunction dostep!(\n Q,\n ark::AdditiveRungeKutta,\n p,\n time::Real,\n nsubsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n ark.substep_stage = iStage\n for i in 1:nsubsteps\n dostep!(Q, ark, ark.variant, p, time, slow_δ, slow_rv_dQ, slow_scaling)\n time += ark.dt\n end\nend\n\nfunction dostep!(\n Q,\n ark::AdditiveRungeKutta,\n variant::NaiveVariant,\n p,\n time::Real,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n dt = ark.dt\n\n RKA_explicit, RKA_implicit = ark.RKA_explicit, ark.RKA_implicit\n RKB_explicit, RKC_explicit = ark.RKB_explicit, ark.RKC_explicit\n RKB_implicit, RKC_implicit = ark.RKB_implicit, ark.RKC_implicit\n rhs!, rhs_implicit! = ark.rhs!, ark.rhs_implicit!\n Qstages, Rstages = (Q, ark.Qstages...), ark.Rstages\n Qhat = ark.Qhat\n split_explicit_implicit = ark.split_explicit_implicit\n Lstages = ark.variant_storage.Lstages\n\n rv_Q = realview(Q)\n rv_Qstages = realview.(Qstages)\n rv_Lstages = realview.(Lstages)\n rv_Rstages = realview.(Rstages)\n rv_Qhat = realview(Qhat)\n\n Nstages = length(RKB_explicit)\n\n groupsize = 256\n\n # calculate the rhs at first stage to initialize the stage loop\n rhs!(\n Rstages[1],\n Qstages[1],\n p,\n time + RKC_explicit[1] * dt,\n increment = false,\n )\n\n rhs_implicit!(\n Lstages[1],\n Qstages[1],\n p,\n time + RKC_implicit[1] * dt,\n increment = false,\n )\n\n # note that it is important that this loop does not modify Q!\n for istage in 2:Nstages\n stagetime_implicit = time + RKC_implicit[istage] * dt\n stagetime_explicit = time + RKC_explicit[istage] * dt\n\n # this kernel also initializes Qstages[istage] with an initial guess\n # for the linear solver\n event = Event(array_device(Q))\n event = stage_update!(array_device(Q), groupsize)(\n variant,\n rv_Q,\n rv_Qstages,\n rv_Lstages,\n rv_Rstages,\n rv_Qhat,\n RKA_explicit,\n RKA_implicit,\n dt,\n Val(istage),\n Val(split_explicit_implicit),\n slow_δ,\n slow_rv_dQ;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n\n # solves\n # Qs = Qhat + dt * RKA_implicit[istage, istage] * rhs_implicit!(Qs)\n rk_coeff = RKA_implicit[istage, istage]\n if !iszero(rk_coeff)\n α = rk_coeff * dt\n besolver! = ark.implicit_solvers[rk_coeff][ark.substep_stage]\n besolver!(Qstages[istage], Qhat, α, p, stagetime_implicit)\n end\n\n rhs!(\n Rstages[istage],\n Qstages[istage],\n p,\n stagetime_explicit,\n increment = false,\n )\n rhs_implicit!(\n Lstages[istage],\n Qstages[istage],\n p,\n stagetime_implicit,\n increment = false,\n )\n end\n\n # compose the final solution\n event = Event(array_device(Q))\n event = solution_update!(array_device(Q), groupsize)(\n variant,\n rv_Q,\n rv_Lstages,\n rv_Rstages,\n RKB_explicit,\n RKB_implicit,\n dt,\n Val(Nstages),\n Val(split_explicit_implicit),\n slow_δ,\n slow_rv_dQ,\n slow_scaling;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\nend\n\nfunction dostep!(\n Q,\n ark::AdditiveRungeKutta,\n variant::LowStorageVariant,\n p,\n time::Real,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n dt = ark.dt\n\n RKA_explicit, RKA_implicit = ark.RKA_explicit, ark.RKA_implicit\n # LowStorageVariant ARK methods assumes that the implicit\n # Butcher table has an SDIRK form; meaning explicit first step (no\n # implicit solve at the first stage) and all non-zero diaognal\n # coefficients are the same.\n rk_coeff = RKA_implicit[2, 2]\n besolver! = ark.implicit_solvers[rk_coeff][ark.substep_stage]\n # NOTE: Using low-storage variant assumes that the butcher tables\n # for both the explicit and implicit parts have the same B and C\n # vectors\n RKB, RKC = ark.RKB_explicit, ark.RKC_explicit\n rhs!, rhs_implicit! = ark.rhs!, ark.rhs_implicit!\n Qstages, Rstages = (Q, ark.Qstages...), ark.Rstages\n Qhat = ark.Qhat\n split_explicit_implicit = ark.split_explicit_implicit\n Qtt = ark.variant_storage.Qtt\n\n rv_Q = realview(Q)\n rv_Qstages = realview.(Qstages)\n rv_Rstages = realview.(Rstages)\n rv_Qhat = realview(Qhat)\n rv_Qtt = realview(Qtt)\n\n Nstages = length(RKB)\n\n groupsize = 256\n\n # calculate the rhs at first stage to initialize the stage loop\n rhs!(Rstages[1], Qstages[1], p, time + RKC[1] * dt, increment = false)\n\n # note that it is important that this loop does not modify Q!\n for istage in 2:Nstages\n stagetime = time + RKC[istage] * dt\n\n # this kernel also initializes Qtt for the linear solver\n event = Event(array_device(Q))\n event = stage_update!(array_device(Q), groupsize)(\n variant,\n rv_Q,\n rv_Qstages,\n rv_Rstages,\n rv_Qhat,\n rv_Qtt,\n RKA_explicit,\n RKA_implicit,\n dt,\n Val(istage),\n Val(split_explicit_implicit),\n slow_δ,\n slow_rv_dQ;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n\n # solves\n # Q_tt = Qhat + dt * RKA_implicit[istage, istage] * rhs_implicit!(Q_tt)\n α = dt * RKA_implicit[istage, istage]\n besolver!(Qtt, Qhat, α, p, stagetime)\n\n # update Qstages\n Qstages[istage] .+= Qtt\n\n rhs!(Rstages[istage], Qstages[istage], p, stagetime, increment = false)\n end\n\n if split_explicit_implicit\n for istage in 1:Nstages\n stagetime = time + RKC[istage] * dt\n rhs_implicit!(\n Rstages[istage],\n Qstages[istage],\n p,\n stagetime,\n increment = true,\n )\n end\n end\n\n # compose the final solution\n event = Event(array_device(Q))\n event = solution_update!(array_device(Q), groupsize)(\n variant,\n rv_Q,\n rv_Rstages,\n RKB,\n dt,\n Val(Nstages),\n slow_δ,\n slow_rv_dQ,\n slow_scaling;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\nend\n\n@kernel function stage_update!(\n ::NaiveVariant,\n Q,\n Qstages,\n Lstages,\n Rstages,\n Qhat,\n RKA_explicit,\n RKA_implicit,\n dt,\n ::Val{is},\n ::Val{split_explicit_implicit},\n slow_δ,\n slow_dQ,\n) where {is, split_explicit_implicit}\n i = @index(Global, Linear)\n @inbounds begin\n Qhat_i = Q[i]\n Qstages_is_i = Q[i]\n\n if slow_δ !== nothing\n Rstages[is - 1][i] += slow_δ * slow_dQ[i]\n end\n\n @unroll for js in 1:(is - 1)\n R_explicit = dt * RKA_explicit[is, js] * Rstages[js][i]\n L_explicit = dt * RKA_explicit[is, js] * Lstages[js][i]\n L_implicit = dt * RKA_implicit[is, js] * Lstages[js][i]\n Qhat_i += (R_explicit + L_implicit)\n Qstages_is_i += R_explicit\n if split_explicit_implicit\n Qstages_is_i += L_explicit\n else\n Qhat_i -= L_explicit\n end\n end\n Qstages[is][i] = Qstages_is_i\n Qhat[i] = Qhat_i\n end\nend\n\n@kernel function stage_update!(\n ::LowStorageVariant,\n Q,\n Qstages,\n Rstages,\n Qhat,\n Qtt,\n RKA_explicit,\n RKA_implicit,\n dt,\n ::Val{is},\n ::Val{split_explicit_implicit},\n slow_δ,\n slow_dQ,\n) where {is, split_explicit_implicit}\n i = @index(Global, Linear)\n @inbounds begin\n Qhat_i = Q[i]\n Qstages_is_i = -zero(eltype(Q))\n\n if slow_δ !== nothing\n Rstages[is - 1][i] += slow_δ * slow_dQ[i]\n end\n\n @unroll for js in 1:(is - 1)\n if split_explicit_implicit\n rkcoeff = RKA_implicit[is, js] / RKA_implicit[is, is]\n else\n rkcoeff =\n (RKA_implicit[is, js] - RKA_explicit[is, js]) /\n RKA_implicit[is, is]\n end\n commonterm = rkcoeff * Qstages[js][i]\n Qhat_i += commonterm + dt * RKA_explicit[is, js] * Rstages[js][i]\n Qstages_is_i -= commonterm\n end\n Qstages[is][i] = Qstages_is_i\n Qhat[i] = Qhat_i\n Qtt[i] = Qhat_i\n end\nend\n\n@kernel function solution_update!(\n ::NaiveVariant,\n Q,\n Lstages,\n Rstages,\n RKB_explicit,\n RKB_implicit,\n dt,\n ::Val{Nstages},\n ::Val{split_explicit_implicit},\n slow_δ,\n slow_dQ,\n slow_scaling,\n) where {Nstages, split_explicit_implicit}\n i = @index(Global, Linear)\n @inbounds begin\n if slow_δ !== nothing\n Rstages[Nstages][i] += slow_δ * slow_dQ[i]\n end\n if slow_scaling !== nothing\n slow_dQ[i] *= slow_scaling\n end\n\n @unroll for is in 1:Nstages\n Q[i] += RKB_explicit[is] * dt * Rstages[is][i]\n if split_explicit_implicit\n Q[i] += RKB_implicit[is] * dt * Lstages[is][i]\n end\n end\n end\nend\n\n@kernel function solution_update!(\n ::LowStorageVariant,\n Q,\n Rstages,\n RKB,\n dt,\n ::Val{Nstages},\n slow_δ,\n slow_dQ,\n slow_scaling,\n) where {Nstages}\n i = @index(Global, Linear)\n @inbounds begin\n if slow_δ !== nothing\n Rstages[Nstages][i] += slow_δ * slow_dQ[i]\n end\n if slow_scaling !== nothing\n slow_dQ[i] *= slow_scaling\n end\n\n @unroll for is in 1:Nstages\n Q[i] += RKB[is] * dt * Rstages[is][i]\n end\n end\nend\n\n\"\"\"\n ARK1ForwardBackwardEuler(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThis uses a first-order-accurate two-stage additive Runge--Kutta scheme\nby combining a forward Euler explicit step with a backward Euler implicit\ncorrection.\n\n### References\n @article{Ascher1997,\n title = {Implicit-explicit Runge-Kutta methods for time-dependent\n partial differential equations},\n author = {Uri M. Ascher and Steven J. Ruuth and Raymond J. Spiteri},\n volume = {25},\n number = {2-3},\n pages = {151--167},\n year = {1997},\n journal = {Applied Numerical Mathematics},\n publisher = {Elsevier {BV}}\n }\n\"\"\"\nfunction ARK1ForwardBackwardEuler(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n RKA_explicit = [\n RT(0) RT(0)\n RT(1) RT(0)\n ]\n RKA_implicit = [\n RT(0) RT(0)\n RT(0) RT(1)\n ]\n\n RKB_explicit = [RT(0), RT(1)]\n RKC_explicit = [RT(0), RT(1)]\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n Nstages = length(RKB_explicit)\n\n AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n ARK2ImplicitExplicitMidpoint(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThis uses a second-order-accurate two-stage additive Runge--Kutta scheme\nby combining the implicit and explicit midpoint methods.\n\n### References\n @article{Ascher1997,\n title = {Implicit-explicit Runge-Kutta methods for time-dependent\n partial differential equations},\n author = {Uri M. Ascher and Steven J. Ruuth and Raymond J. Spiteri},\n volume = {25},\n number = {2-3},\n pages = {151--167},\n year = {1997},\n journal = {Applied Numerical Mathematics},\n publisher = {Elsevier {BV}}\n }\n\"\"\"\nfunction ARK2ImplicitExplicitMidpoint(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n RKA_explicit = [\n RT(0) RT(0)\n RT(1 / 2) RT(0)\n ]\n RKA_implicit = [\n RT(0) RT(0)\n RT(0) RT(1 / 2)\n ]\n\n RKB_explicit = [RT(0), RT(1)]\n RKC_explicit = [RT(0), RT(1 / 2)]\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n Nstages = length(RKB_explicit)\n\n AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n ARK2GiraldoKellyConstantinescu(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant, paperversion)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\n`paperversion=true` uses the coefficients from the paper, `paperversion=false`\nuses coefficients that make the scheme (much) more stable but less accurate\n\nThis uses the second-order-accurate 3-stage additive Runge--Kutta scheme of\nGiraldo, Kelly and Constantinescu (2013).\n\n### References\n - [Giraldo2013](@cite)\n\"\"\"\nfunction ARK2GiraldoKellyConstantinescu(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n paperversion = false,\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n a32 = RT(paperversion ? (3 + 2 * sqrt(2)) / 6 : 1 // 2)\n RKA_explicit = [\n RT(0) RT(0) RT(0)\n RT(2 - sqrt(2)) RT(0) RT(0)\n RT(1 - a32) RT(a32) RT(0)\n ]\n\n RKA_implicit = [\n RT(0) RT(0) RT(0)\n RT(1 - 1 / sqrt(2)) RT(1 - 1 / sqrt(2)) RT(0)\n RT(1 / (2 * sqrt(2))) RT(1 / (2 * sqrt(2))) RT(1 - 1 / sqrt(2))\n ]\n\n RKB_explicit =\n [RT(1 / (2 * sqrt(2))), RT(1 / (2 * sqrt(2))), RT(1 - 1 / sqrt(2))]\n RKC_explicit = [RT(0), RT(2 - sqrt(2)), RT(1)]\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n Nstages = length(RKB_explicit)\n\n AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n Trap2LockWoodWeller(F, L, backward_euler_solver, Q; dt, t0, nsubsteps,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThe time integrator scheme used is Trap2(2,3,2) with δ_s = 1, δ_f = 0, from\nthe following reference\n\n### References\n @article{Ascher1997,\n title = {Numerical analyses of Runge–Kutta implicit–explicit schemes\n for horizontally explicit, vertically implicit solutions of\n atmospheric models},\n author = {S.-J. Lock and N. Wood and H. Weller},\n volume = {140},\n number = {682},\n pages = {1654-1669},\n year = {2014},\n journal = {Quarterly Journal of the Royal Meteorological Society},\n publisher = {{RMetS}}\n }\n\"\"\"\nfunction Trap2LockWoodWeller(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = NaiveVariant(),\n δ_s = 1,\n δ_f = 0,\n α = 0,\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n # In this scheme B and C vectors do not coincide,\n # hence we can't use the LowStorageVariant optimization\n @assert variant isa NaiveVariant\n\n T = eltype(Q)\n RT = real(T)\n\n #! format: off\n RKA_explicit = [\n RT(0) RT(0) RT(0) RT(0)\n RT(δ_s) RT(0) RT(0) RT(0)\n RT(1 / 2) RT(1 / 2) RT(0) RT(0)\n RT(1 / 2) RT(0) RT(1 / 2) RT(0)\n ]\n\n RKA_implicit = [\n RT(0) RT(0) RT(0) RT(0)\n RT(δ_f * (1 - α) / 2) RT(δ_f * (1 + α) / 2) RT(0) RT(0)\n RT(1 / 2) RT(0) RT(1 / 2) RT(0)\n RT(1 / 2) RT(0) RT(0) RT(1 / 2)\n ]\n #! format: on\n\n RKB_explicit = [RT(1 / 2), RT(0), RT(1 / 2), RT(0)]\n RKB_implicit = [RT(1 / 2), RT(0), RT(0), RT(1 / 2)]\n RKC_explicit = [RT(0), RT(δ_s), RT(1), RT(1)]\n RKC_implicit = [RT(0), RT(δ_f), RT(1), RT(1)]\n\n Nstages = length(RKB_explicit)\n\n AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n ARK548L2SA2KennedyCarpenter(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThis uses the fifth-order-accurate 8-stage additive Runge--Kutta scheme of\nKennedy and Carpenter (2013).\n\n### References\n - [Kennedy2019](@cite)\n\"\"\"\nfunction ARK548L2SA2KennedyCarpenter(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n Nstages = 8\n gamma = RT(2 // 9)\n\n # declared as Arrays for mutability, later these will be converted to static\n # arrays\n RKA_explicit = zeros(RT, Nstages, Nstages)\n RKA_implicit = zeros(RT, Nstages, Nstages)\n RKB_explicit = zeros(RT, Nstages)\n RKC_explicit = zeros(RT, Nstages)\n\n # the main diagonal\n for is in 2:Nstages\n RKA_implicit[is, is] = gamma\n end\n\n RKA_implicit[3, 2] = RT(2366667076620 // 8822750406821)\n RKA_implicit[4, 2] = RT(-257962897183 // 4451812247028)\n RKA_implicit[4, 3] = RT(128530224461 // 14379561246022)\n RKA_implicit[5, 2] = RT(-486229321650 // 11227943450093)\n RKA_implicit[5, 3] = RT(-225633144460 // 6633558740617)\n RKA_implicit[5, 4] = RT(1741320951451 // 6824444397158)\n RKA_implicit[6, 2] = RT(621307788657 // 4714163060173)\n RKA_implicit[6, 3] = RT(-125196015625 // 3866852212004)\n RKA_implicit[6, 4] = RT(940440206406 // 7593089888465)\n RKA_implicit[6, 5] = RT(961109811699 // 6734810228204)\n RKA_implicit[7, 2] = RT(2036305566805 // 6583108094622)\n RKA_implicit[7, 3] = RT(-3039402635899 // 4450598839912)\n RKA_implicit[7, 4] = RT(-1829510709469 // 31102090912115)\n RKA_implicit[7, 5] = RT(-286320471013 // 6931253422520)\n RKA_implicit[7, 6] = RT(8651533662697 // 9642993110008)\n\n RKA_explicit[3, 1] = RT(1 // 9)\n RKA_explicit[3, 2] = RT(1183333538310 // 1827251437969)\n RKA_explicit[4, 1] = RT(895379019517 // 9750411845327)\n RKA_explicit[4, 2] = RT(477606656805 // 13473228687314)\n RKA_explicit[4, 3] = RT(-112564739183 // 9373365219272)\n RKA_explicit[5, 1] = RT(-4458043123994 // 13015289567637)\n RKA_explicit[5, 2] = RT(-2500665203865 // 9342069639922)\n RKA_explicit[5, 3] = RT(983347055801 // 8893519644487)\n RKA_explicit[5, 4] = RT(2185051477207 // 2551468980502)\n RKA_explicit[6, 1] = RT(-167316361917 // 17121522574472)\n RKA_explicit[6, 2] = RT(1605541814917 // 7619724128744)\n RKA_explicit[6, 3] = RT(991021770328 // 13052792161721)\n RKA_explicit[6, 4] = RT(2342280609577 // 11279663441611)\n RKA_explicit[6, 5] = RT(3012424348531 // 12792462456678)\n RKA_explicit[7, 1] = RT(6680998715867 // 14310383562358)\n RKA_explicit[7, 2] = RT(5029118570809 // 3897454228471)\n RKA_explicit[7, 3] = RT(2415062538259 // 6382199904604)\n RKA_explicit[7, 4] = RT(-3924368632305 // 6964820224454)\n RKA_explicit[7, 5] = RT(-4331110370267 // 15021686902756)\n RKA_explicit[7, 6] = RT(-3944303808049 // 11994238218192)\n RKA_explicit[8, 1] = RT(2193717860234 // 3570523412979)\n RKA_explicit[8, 2] = RKA_explicit[8, 1]\n RKA_explicit[8, 3] = RT(5952760925747 // 18750164281544)\n RKA_explicit[8, 4] = RT(-4412967128996 // 6196664114337)\n RKA_explicit[8, 5] = RT(4151782504231 // 36106512998704)\n RKA_explicit[8, 6] = RT(572599549169 // 6265429158920)\n RKA_explicit[8, 7] = RT(-457874356192 // 11306498036315)\n\n RKB_explicit[2] = 0\n RKB_explicit[3] = RT(3517720773327 // 20256071687669)\n RKB_explicit[4] = RT(4569610470461 // 17934693873752)\n RKB_explicit[5] = RT(2819471173109 // 11655438449929)\n RKB_explicit[6] = RT(3296210113763 // 10722700128969)\n RKB_explicit[7] = RT(-1142099968913 // 5710983926999)\n\n RKC_explicit[2] = RT(4 // 9)\n RKC_explicit[3] = RT(6456083330201 // 8509243623797)\n RKC_explicit[4] = RT(1632083962415 // 14158861528103)\n RKC_explicit[5] = RT(6365430648612 // 17842476412687)\n RKC_explicit[6] = RT(18 // 25)\n RKC_explicit[7] = RT(191 // 200)\n\n for is in 2:Nstages\n RKA_implicit[is, 1] = RKA_implicit[is, 2]\n end\n\n for is in 1:(Nstages - 1)\n RKA_implicit[Nstages, is] = RKB_explicit[is]\n end\n\n RKB_explicit[1] = RKB_explicit[2]\n RKB_explicit[8] = gamma\n\n RKA_explicit[2, 1] = RKC_explicit[2]\n RKA_explicit[Nstages, 1] = RKA_explicit[Nstages, 2]\n\n RKC_explicit[1] = 0\n RKC_explicit[Nstages] = 1\n\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n ark = AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n ARK437L2SA1KennedyCarpenter(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThis uses the fourth-order-accurate 7-stage additive Runge--Kutta scheme of\nKennedy and Carpenter (2013).\n\n### References\n - [Kennedy2019](@cite)\n\"\"\"\nfunction ARK437L2SA1KennedyCarpenter(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n Nstages = 7\n gamma = RT(1235 // 10000)\n\n # declared as Arrays for mutability, later these will be converted to static\n # arrays\n RKA_explicit = zeros(RT, Nstages, Nstages)\n RKA_implicit = zeros(RT, Nstages, Nstages)\n RKB_explicit = zeros(RT, Nstages)\n RKC_explicit = zeros(RT, Nstages)\n\n # the main diagonal\n for is in 2:Nstages\n RKA_implicit[is, is] = gamma\n end\n\n RKA_implicit[3, 2] = RT(624185399699 // 4186980696204)\n RKA_implicit[4, 2] = RT(1258591069120 // 10082082980243)\n RKA_implicit[4, 3] = RT(-322722984531 // 8455138723562)\n RKA_implicit[5, 2] = RT(-436103496990 // 5971407786587)\n RKA_implicit[5, 3] = RT(-2689175662187 // 11046760208243)\n RKA_implicit[5, 4] = RT(4431412449334 // 12995360898505)\n RKA_implicit[6, 2] = RT(-2207373168298 // 14430576638973)\n RKA_implicit[6, 3] = RT(242511121179 // 3358618340039)\n RKA_implicit[6, 4] = RT(3145666661981 // 7780404714551)\n RKA_implicit[6, 5] = RT(5882073923981 // 14490790706663)\n RKA_implicit[7, 2] = 0\n RKA_implicit[7, 3] = RT(9164257142617 // 17756377923965)\n RKA_implicit[7, 4] = RT(-10812980402763 // 74029279521829)\n RKA_implicit[7, 5] = RT(1335994250573 // 5691609445217)\n RKA_implicit[7, 6] = RT(2273837961795 // 8368240463276)\n\n RKA_explicit[3, 1] = RT(247 // 4000)\n RKA_explicit[3, 2] = RT(2694949928731 // 7487940209513)\n RKA_explicit[4, 1] = RT(464650059369 // 8764239774964)\n RKA_explicit[4, 2] = RT(878889893998 // 2444806327765)\n RKA_explicit[4, 3] = RT(-952945855348 // 12294611323341)\n RKA_explicit[5, 1] = RT(476636172619 // 8159180917465)\n RKA_explicit[5, 2] = RT(-1271469283451 // 7793814740893)\n RKA_explicit[5, 3] = RT(-859560642026 // 4356155882851)\n RKA_explicit[5, 4] = RT(1723805262919 // 4571918432560)\n RKA_explicit[6, 1] = RT(6338158500785 // 11769362343261)\n RKA_explicit[6, 2] = RT(-4970555480458 // 10924838743837)\n RKA_explicit[6, 3] = RT(3326578051521 // 2647936831840)\n RKA_explicit[6, 4] = RT(-880713585975 // 1841400956686)\n RKA_explicit[6, 5] = RT(-1428733748635 // 8843423958496)\n RKA_explicit[7, 2] = RT(760814592956 // 3276306540349)\n RKA_explicit[7, 3] = RT(-47223648122716 // 6934462133451)\n RKA_explicit[7, 4] = RT(71187472546993 // 9669769126921)\n RKA_explicit[7, 5] = RT(-13330509492149 // 9695768672337)\n RKA_explicit[7, 6] = RT(11565764226357 // 8513123442827)\n\n RKB_explicit[2] = 0\n RKB_explicit[3] = RT(9164257142617 // 17756377923965)\n RKB_explicit[4] = RT(-10812980402763 // 74029279521829)\n RKB_explicit[5] = RT(1335994250573 // 5691609445217)\n RKB_explicit[6] = RT(2273837961795 // 8368240463276)\n RKB_explicit[7] = RT(247 // 2000)\n\n RKC_explicit[2] = RT(247 // 1000)\n RKC_explicit[3] = RT(4276536705230 // 10142255878289)\n RKC_explicit[4] = RT(67 // 200)\n RKC_explicit[5] = RT(3 // 40)\n RKC_explicit[6] = RT(7 // 10)\n\n for is in 2:Nstages\n RKA_implicit[is, 1] = RKA_implicit[is, 2]\n end\n\n for is in 1:(Nstages - 1)\n RKA_implicit[Nstages, is] = RKB_explicit[is]\n end\n\n RKB_explicit[1] = RKB_explicit[2]\n\n RKA_explicit[2, 1] = RKC_explicit[2]\n RKA_explicit[Nstages, 1] = RKA_explicit[Nstages, 2]\n\n RKC_explicit[1] = 0\n RKC_explicit[Nstages] = 1\n\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n ark = AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n\n\"\"\"\n DBM453VoglEtAl(f, l, backward_euler_solver, Q; dt, t0,\n split_explicit_implicit, variant)\n\nThis function returns an [`AdditiveRungeKutta`](@ref) time stepping object,\nsee the documentation of [`AdditiveRungeKutta`](@ref) for arguments definitions.\nThis time stepping object is intended to be passed to the `solve!` command.\n\nThis uses the third-order-accurate 5-stage additive Runge--Kutta scheme of\nVogl et al. (2019).\n\n### References\n - [Vogl2019](@cite)\n\"\"\"\nfunction DBM453VoglEtAl(\n F,\n L,\n backward_euler_solver,\n Q::AT;\n dt = nothing,\n t0 = 0,\n nsubsteps = [],\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n) where {AT <: AbstractArray}\n\n @assert dt !== nothing\n\n T = eltype(Q)\n RT = real(T)\n\n Nstages = 5\n gamma = RT(0.32591194130117247)\n\n # declared as Arrays for mutability, later these will be converted to static\n # arrays\n RKA_explicit = zeros(RT, Nstages, Nstages)\n RKA_implicit = zeros(RT, Nstages, Nstages)\n RKB_explicit = zeros(RT, Nstages)\n RKC_explicit = zeros(RT, Nstages)\n\n # the main diagonal\n for is in 2:Nstages\n RKA_implicit[is, is] = gamma\n end\n\n RKA_implicit[2, 1] = -0.22284985318525410\n RKA_implicit[3, 1] = -0.46801347074080545\n RKA_implicit[3, 2] = 0.86349284225716961\n RKA_implicit[4, 1] = -0.46509906651927421\n RKA_implicit[4, 2] = 0.81063103116959553\n RKA_implicit[4, 3] = 0.61036726756832357\n RKA_implicit[5, 1] = 0.87795339639076675\n RKA_implicit[5, 2] = -0.72692641526151547\n RKA_implicit[5, 3] = 0.75204137157372720\n RKA_implicit[5, 4] = -0.22898029400415088\n\n RKA_explicit[2, 1] = 0.10306208811591838\n RKA_explicit[3, 1] = -0.94124866143519894\n RKA_explicit[3, 2] = 1.66263997425273560\n RKA_explicit[4, 1] = -1.36709752014377650\n RKA_explicit[4, 2] = 1.38158529110168730\n RKA_explicit[4, 3] = 1.26732340256190650\n RKA_explicit[5, 1] = -0.81287582068772448\n RKA_explicit[5, 2] = 0.81223739060505738\n RKA_explicit[5, 3] = 0.90644429603699305\n RKA_explicit[5, 4] = 0.094194134045674111\n\n RKB_explicit[1] = 0.87795339639076672\n RKB_explicit[2] = -0.72692641526151549\n RKB_explicit[3] = 0.7520413715737272\n RKB_explicit[4] = -0.22898029400415090\n RKB_explicit[5] = 0.32591194130117247\n\n RKC_explicit[1] = 0\n RKC_explicit[2] = 0.1030620881159184\n RKC_explicit[3] = 0.72139131281753662\n RKC_explicit[4] = 1.28181117351981733\n RKC_explicit[5] = 1\n\n # For this ARK method, both RK methods share the same\n # B and C vectors in the Butcher table\n RKB_implicit = RKB_explicit\n RKC_implicit = RKC_explicit\n\n ark = AdditiveRungeKutta(\n F,\n L,\n backward_euler_solver,\n RKA_explicit,\n RKA_implicit,\n RKB_explicit,\n RKB_implicit,\n RKC_explicit,\n RKC_implicit,\n split_explicit_implicit,\n variant,\n Q;\n dt = dt,\n t0 = t0,\n nsubsteps = nsubsteps,\n )\nend\n" }, { "path": "src/Numerics/ODESolvers/BackwardEulerSolvers.jl", "content": "export LinearBackwardEulerSolver, AbstractBackwardEulerSolver\nexport NonLinearBackwardEulerSolver\n\nabstract type AbstractImplicitOperator end\n\n\"\"\"\n op! = EulerOperator(f!, ϵ)\n\nConstruct a linear operator which performs an explicit Euler step ``Q + α\nf(Q)``, where `f!` and `op!` both operate inplace, with extra arguments passed\nthrough, i.e.\n```\nop!(LQ, Q, args...)\n```\nis equivalent to\n```\nf!(dQ, Q, args...)\nLQ .= Q .+ ϵ .* dQ\n```\n\"\"\"\nmutable struct EulerOperator{F, FT} <: AbstractImplicitOperator\n f!::F\n ϵ::FT\nend\n\nfunction (op::EulerOperator)(LQ, Q, args...)\n op.f!(LQ, Q, args..., increment = false)\n @. LQ = Q + op.ϵ * LQ\nend\n\n\"\"\"\n AbstractBackwardEulerSolver\n\nAn abstract backward Euler method\n\"\"\"\nabstract type AbstractBackwardEulerSolver end\n\n\"\"\"\n (be::AbstractBackwardEulerSolver)(Q, Qhat, α, param, time)\n\nEach concrete implementations of `AbstractBackwardEulerSolver` should provide a\ncallable version which solves the following system for `Q`\n```\n Q = Qhat + α f(Q, param, time)\n```\nwhere `f` is the ODE tendency function, `param` are the ODE parameters, and\n`time` is the current ODE time. The arguments `Q` should be modified in place\nand should not be assumed to be initialized to any value.\n\"\"\"\n(be::AbstractBackwardEulerSolver)(Q, Qhat, α, p, t) =\n throw(MethodError(be, (Q, Qhat, α, p, t)))\n\n\"\"\"\n Δt_is_adjustable(::AbstractBackwardEulerSolver)\n\nReturn `Bool` for whether this backward Euler solver can be updated. default is\n`false`.\n\"\"\"\nΔt_is_adjustable(::AbstractBackwardEulerSolver) = false\n\n\"\"\"\n update_backward_Euler_solver!(::AbstractBackwardEulerSolver, α)\n\nUpdate the given backward Euler solver for the parameter `α`; see\n['AbstractBackwardEulerSolver'](@ref). Default behavior is no change to the\nsolver.\n\"\"\"\nupdate_backward_Euler_solver!(::AbstractBackwardEulerSolver, Q, α) = nothing\n\n\"\"\"\n setup_backward_Euler_solver(solver, Q, α, tendency!)\n\nReturns a concrete implementation of an `AbstractBackwardEulerSolver` that will\nsolve for `Q` in systems of the form of\n```\n Q = Qhat + α f(Q, param, time)\n```\nwhere `tendency!` is the in-place tendency function. Not the array `Q` is just\npassed in for type information, e.g., `Q` the same `Q` will not be used for all\ncalls to the solver.\n\"\"\"\nsetup_backward_Euler_solver(solver::AbstractBackwardEulerSolver, _...) = solver\n\n\"\"\"\n LinearBackwardEulerSolver(::AbstractSystemSolver; isadjustable = false)\n\nHelper type for specifying building a backward Euler solver with a linear\nsolver. If `isadjustable == true` then the solver can be updated with a new\ntime step size.\n\"\"\"\nstruct LinearBackwardEulerSolver{LS}\n solver::LS\n isadjustable::Bool\n preconditioner_update_freq::Int\n LinearBackwardEulerSolver(\n solver;\n isadjustable = false,\n preconditioner_update_freq = -1,\n ) = new{typeof(solver)}(solver, isadjustable, preconditioner_update_freq)\nend\n\n\"\"\"\n LinBESolver\n\nConcrete implementation of an `AbstractBackwardEulerSolver` to use linear\nsolvers of type `AbstractSystemSolver`. See helper type\n[`LinearBackwardEulerSolver`](@ref)\n```\n Q = Qhat + α f(Q, param, time)\n```\n\"\"\"\nmutable struct LinBESolver{FT, F, LS} <: AbstractBackwardEulerSolver\n α::FT\n f_imp!::F\n solver::LS\n isadjustable::Bool\n # used only for iterative solver\n preconditioner::AbstractPreconditioner\n # used only for direct solver\n factors::Any\nend\n\nΔt_is_adjustable(lin::LinBESolver) = lin.isadjustable\n\nfunction setup_backward_Euler_solver(\n lin::LinearBackwardEulerSolver,\n Q,\n α,\n f_imp!,\n)\n FT = eltype(α)\n rhs! = EulerOperator(f_imp!, -α)\n\n factors = prefactorize(rhs!, lin.solver, Q, nothing, FT(NaN))\n\n # when direct solver is applied preconditioner_update_freq <= 0\n @assert(\n typeof(lin.solver) <: AbstractIterativeSystemSolver ||\n lin.preconditioner_update_freq <= 0\n )\n\n preconditioner_update_freq = lin.preconditioner_update_freq\n # construct an empty preconditioner\n preconditioner = (\n preconditioner_update_freq > 0 ?\n ColumnwiseLUPreconditioner(f_imp!, Q, preconditioner_update_freq) :\n NoPreconditioner()\n )\n\n LinBESolver(\n α,\n f_imp!,\n lin.solver,\n lin.isadjustable,\n preconditioner,\n factors,\n )\nend\n\nfunction update_backward_Euler_solver!(lin::LinBESolver, Q, α)\n lin.α = α\n FT = eltype(Q)\n # for direct solver, update factors\n # for iterative solver, set factors to Nothing (TODO optimize)\n lin.factors = prefactorize(\n EulerOperator(lin.f_imp!, -α),\n lin.solver,\n Q,\n nothing,\n FT(NaN),\n )\nend\n\nfunction (lin::LinBESolver)(Q, Qhat, α, p, t)\n\n # If α is not the same as the already assembled operator\n # with its previous value of α (lin.α), then we need to\n # first check that the solver CAN be rebuilt (lin.isadjustable)\n # followed by recalling the set up routine by updating the\n # coefficient with the new version of α\n if lin.α != α\n @assert lin.isadjustable\n update_backward_Euler_solver!(lin, Q, α)\n end\n\n rhs! = EulerOperator(lin.f_imp!, -α)\n\n if typeof(lin.solver) <: AbstractIterativeSystemSolver\n FT = eltype(α)\n preconditioner_update!(rhs!, rhs!.f!, lin.preconditioner, p, t)\n linearsolve!(rhs!, lin.preconditioner, lin.solver, Q, Qhat, p, t)\n preconditioner_counter_update!(lin.preconditioner)\n else\n linearsolve!(rhs!, lin.factors, lin.solver, Q, Qhat, p, t)\n end\nend\n\n\"\"\"\n struct NonLinearBackwardEulerSolver{NLS}\n nlsolver::NLS\n isadjustable::Bool\n preconditioner_update_freq::Int64\n end\n\nHelper type for specifying building a nonlinear backward Euler solver with a nonlinear\nsolver.\n\n# Arguments\n- `nlsolver`: iterative nonlinear solver, i.e., JacobianFreeNewtonKrylovSolver\n- `isadjustable`: TODO not used, might use for updating preconditioner\n- `preconditioner_update_freq`: relavent to Jacobian free -1: no preconditioner;\n positive number, update every freq times\n\"\"\"\nstruct NonLinearBackwardEulerSolver{NLS}\n nlsolver::NLS\n isadjustable::Bool\n # preconditioner_update_freq, -1: no preconditioner;\n # positive number, update every freq times\n preconditioner_update_freq::Int\n function NonLinearBackwardEulerSolver(\n nlsolver;\n isadjustable = false,\n preconditioner_update_freq = -1,\n )\n NLS = typeof(nlsolver)\n return new{NLS}(nlsolver, isadjustable, preconditioner_update_freq)\n end\nend\n\n\n\"\"\"\n LinBESolver\n\nConcrete implementation of an `AbstractBackwardEulerSolver` to use nonlinear\nsolvers of type `NLS`. See helper type\n[`NonLinearBackwardEulerSolver`](@ref)\n```\n Q = Qhat + α f_imp(Q, param, time)\n```\n\"\"\"\nmutable struct NonLinBESolver{FT, F, NLS} <: AbstractBackwardEulerSolver\n # Solve Q - α f_imp(Q) = Qrhs\n α::FT\n # implcit operator\n f_imp!::F\n # jacobian action, which approximates drhs!/dQ⋅ΔQ , here rhs!(Q) = Q - α f_imp(Q)\n jvp!::JacobianAction\n # nonlinear solver\n nlsolver::NLS\n # whether adjust the time step or not\n isadjustable::Bool\n # preconditioner, approximation of drhs!/dQ\n preconditioner::AbstractPreconditioner\n\nend\n\nΔt_is_adjustable(nlsolver::NonLinBESolver) = nlsolver.isadjustable\n\n\"\"\"\n setup_backward_Euler_solver(solver::NonLinearBackwardEulerSolver, Q, α, tendency!)\n\nReturns a concrete implementation of an `AbstractBackwardEulerSolver` that will\nsolve for `Q` in nonlinear systems of the form of\n```\n Q = Qhat + α f(Q, param, time)\n```\nCreate an empty JacobianAction\n\nCreate an empty preconditioner if preconditioner_update_freq > 0\n\"\"\"\nfunction setup_backward_Euler_solver(\n nlbesolver::NonLinearBackwardEulerSolver,\n Q,\n α,\n f_imp!,\n)\n # Create an empty JacobianAction (without operator)\n jvp! = JacobianAction(nothing, Q, nlbesolver.nlsolver.ϵ)\n\n # Create an empty preconditioner if preconditioner_update_freq > 0\n preconditioner_update_freq = nlbesolver.preconditioner_update_freq\n # construct an empty preconditioner\n preconditioner = (\n preconditioner_update_freq > 0 ?\n ColumnwiseLUPreconditioner(f_imp!, Q, preconditioner_update_freq) :\n NoPreconditioner()\n )\n NonLinBESolver(\n α,\n f_imp!,\n jvp!,\n nlbesolver.nlsolver,\n nlbesolver.isadjustable,\n preconditioner,\n )\nend\n\n\"\"\"\nNonlinear solve\n\nUpdate rhs! with α\n\nUpdate the rhs! in the jacobian action jvp!\n\"\"\"\nfunction (nlbesolver::NonLinBESolver)(Q, Qhat, α, p, t)\n\n rhs! = EulerOperator(nlbesolver.f_imp!, -α)\n\n nlbesolver.jvp!.rhs! = rhs!\n\n nonlinearsolve!(\n rhs!,\n nlbesolver.jvp!,\n nlbesolver.preconditioner,\n nlbesolver.nlsolver,\n Q,\n Qhat,\n p,\n t;\n max_newton_iters = nlbesolver.nlsolver.M,\n )\n\nend\n" }, { "path": "src/Numerics/ODESolvers/DifferentialEquations.jl", "content": "import DiffEqBase\nexport DiffEqJLSolver, DiffEqJLIMEXSolver\n\nabstract type AbstractDiffEqJLSolver <: AbstractODESolver end\n\n\"\"\"\n DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nvia a DifferentialEquations.jl DEAlgorithm, which includes support\nfor OrdinaryDiffEq.jl, Sundials.jl, and more.\n\"\"\"\nmutable struct DiffEqJLSolver{I} <: AbstractDiffEqJLSolver\n integ::I\n steps::Int\n\n function DiffEqJLSolver(\n rhs!,\n alg,\n Q,\n args...;\n t0 = 0,\n p = nothing,\n kwargs...,\n )\n prob = DiffEqBase.ODEProblem(\n (du, u, p, t) -> rhs!(du, u, p, t; increment = false),\n Q,\n (float(t0), typemax(typeof(float(t0)))),\n p,\n )\n integ = DiffEqBase.init(\n prob,\n alg,\n args...;\n adaptive = false,\n save_everystep = false,\n save_start = false,\n save_end = false,\n kwargs...,\n )\n new{typeof(integ)}(integ, 0)\n end\nend\n\n\"\"\"\n DiffEqJLSolver(f, RKA, RKB, RKC, Q; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f_I(Q, t) + f_E(Q, t)\n```\n\nvia a DifferentialEquations.jl DEAlgorithm, which includes support\nfor OrdinaryDiffEq.jl, Sundials.jl, and more.\n\"\"\"\nmutable struct DiffEqJLIMEXSolver{I} <: AbstractDiffEqJLSolver\n integ::I\n steps::Int\n\n function DiffEqJLIMEXSolver(\n rhs!,\n rhs_implicit!,\n alg,\n Q,\n args...;\n t0 = 0,\n p = nothing,\n kwargs...,\n )\n prob = DiffEqBase.SplitODEProblem(\n (du, u, p, t) -> rhs_implicit!(du, u, p, t; increment = false),\n (du, u, p, t) -> rhs!(du, u, p, t; increment = false),\n Q,\n (float(t0), typemax(typeof(float(t0)))),\n p,\n )\n integ = DiffEqBase.init(\n prob,\n alg,\n args...;\n adaptive = false,\n save_everystep = false,\n save_start = false,\n save_end = false,\n kwargs...,\n )\n\n new{typeof(integ)}(integ, 0)\n end\nend\n\ngettime(solver::AbstractDiffEqJLSolver) = solver.integ.t\ngetdt(solver::AbstractDiffEqJLSolver) = solver.integ.dt\nupdatedt!(solver::AbstractDiffEqJLSolver, dt) =\n DiffEqBase.set_proposed_dt!(solver.integ, dt)\nupdatetime!(solver::AbstractDiffEqJLSolver, t) =\n DiffEqBase.set_t!(solver.integ, t)\nisadjustable(solver::AbstractDiffEqJLSolver) = true # Is this isadaptive? Or something different?\n\n\"\"\"\n ODESolvers.general_dostep!(Q, solver::AbstractODESolver, p,\n timeend::Real, adjustfinalstep::Bool)\n\nUse the solver to step `Q` forward in time from the current time, to the time\n`timeend`. If `adjustfinalstep == true` then `dt` is adjusted so that the step\ndoes not take the solution beyond the `timeend`.\n\"\"\"\nfunction general_dostep!(\n Q,\n solver::AbstractDiffEqJLSolver,\n p,\n timeend::Real;\n adjustfinalstep::Bool,\n)\n integ = solver.integ\n\n if first(integ.opts.tstops) !== timeend\n DiffEqBase.add_tstop!(integ, timeend)\n end\n dostep!(Q, solver, p, time)\n solver.integ.t\nend\n\nfunction dostep!(\n Q,\n solver::AbstractDiffEqJLSolver,\n p,\n time,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n in_slow_scaling = nothing,\n)\n\n integ = solver.integ\n integ.p = p # Can this change?\n\n rv_Q = realview(Q)\n if integ.u != Q\n integ.u .= Q\n DiffEqBase.u_modified!(integ, true)\n # Will time always be correct?\n end\n\n DiffEqBase.step!(integ)\n rv_Q .= solver.integ.u\nend\n\nfunction DiffEqJLConstructor(alg)\n constructor =\n (F, Q; dt = 0, t0 = 0) -> DiffEqJLSolver(F, alg, Q; t0 = t0, dt = dt)\n return constructor\nend\n" }, { "path": "src/Numerics/ODESolvers/GenericCallbacks.jl", "content": "\"\"\"\n GenericCallbacks\n\nThis module defines interfaces and wrappers for callbacks to be used with an\n`AbstractODESolver`.\n\nA callback `cb` defines three methods:\n\n- `GenericCallbacks.init!(cb, solver, Q, param, t)`, to be called at solver\n initialization.\n\n- `GenericCallbacks.call!(cb, solver, Q, param, t)`, to be called after each\n time step: the return value dictates what action should be taken:\n\n * `0` or `nothing`: continue time stepping as usual\n * `1`: stop time stepping after all callbacks have been executed\n * `2`: stop time stepping immediately\n\n- `GenericCallbacks.fini!(cb, solver, Q, param, t)`, to be called at solver\n finish.\n\nAdditionally, _wrapper_ callbacks can be used to execute the callbacks under\ncertain conditions:\n\n - [`AtInit`](@ref)\n - [`AtInitAndFini`](@ref)\n - [`EveryXWallTimeSeconds`](@ref)\n - [`EveryXSimulationTime`](@ref)\n - [`EveryXSimulationSteps`](@ref)\n\nFor convenience, the following objects can also be used as callbacks:\n\n- A `Function` object `f`, `init!` and `fini!` are no-ops, and `call!` will\n call `f()`, and ignore the return value.\n- A `Tuple` object will call `init!`, `call!` and `fini!` on each element\n of the tuple.\n\"\"\"\nmodule GenericCallbacks\n\nexport AtInit,\n AtInitAndFini,\n EveryXWallTimeSeconds,\n EveryXSimulationTime,\n EveryXSimulationSteps\n\nusing MPI\n\ninit!(f::Function, solver, Q, param, t) = nothing\nfunction call!(f::Function, solver, Q, param, t)\n f()\n return nothing\nend\nfini!(f::Function, solver, Q, param, t) = nothing\n\nfunction init!(callbacks::Tuple, solver, Q, param, t)\n for cb in callbacks\n GenericCallbacks.init!(cb, solver, Q, param, t)\n end\nend\nfunction call!(callbacks::Tuple, solver, Q, param, t)\n val = 0\n for cb in callbacks\n val_i = GenericCallbacks.call!(cb, solver, Q, param, t)\n val_i = (val_i === nothing) ? 0 : val_i\n val = max(val, val_i)\n if val == 2\n return val\n end\n end\n return val\nend\nfunction fini!(callbacks::Tuple, solver, Q, param, t)\n for cb in callbacks\n GenericCallbacks.fini!(cb, solver, Q, param, t)\n end\nend\n\nabstract type AbstractCallback end\n\n\"\"\"\n AtInit(callback) <: AbstractCallback\n\nA wrapper callback to execute `callback` at initialization as well as\nafter each interval.\n\"\"\"\nstruct AtInit <: AbstractCallback\n callback::Any\nend\nfunction init!(cb::AtInit, solver, Q, param, t)\n init!(cb.callback, solver, Q, param, t)\n call!(cb.callback, solver, Q, param, t)\nend\nfunction call!(cb::AtInit, solver, Q, param, t)\n call!(cb.callback, solver, Q, param, t)\nend\nfunction fini!(cb::AtInit, solver, Q, param, t)\n fini!(cb.callback, solver, Q, param, t)\nend\n\n\"\"\"\n AtInitAndFini(callback) <: AbstractCallback\n\nA wrapper callback to execute `callback` at initialization and at\nfinish as well as after each interval.\n\"\"\"\nstruct AtInitAndFini <: AbstractCallback\n callback::Any\nend\nfunction init!(cb::AtInitAndFini, solver, Q, param, t)\n init!(cb.callback, solver, Q, param, t)\n call!(cb.callback, solver, Q, param, t)\nend\nfunction call!(cb::AtInitAndFini, solver, Q, param, t)\n call!(cb.callback, solver, Q, param, t)\nend\nfunction fini!(cb::AtInitAndFini, solver, Q, param, t)\n call!(cb.callback, solver, Q, param, t)\n fini!(cb.callback, solver, Q, param, t)\nend\n\n\"\"\"\n EveryXWallTimeSeconds(callback, Δtime, mpicomm)\n\nA wrapper callback to execute `callback` every `Δtime` wallclock time seconds.\n`mpicomm` is used to syncronize runtime across MPI ranks.\n\"\"\"\nmutable struct EveryXWallTimeSeconds <: AbstractCallback\n \"callback to wrap\"\n callback::Any\n \"wall time seconds between callbacks\"\n Δtime::Real\n \"MPI communicator\"\n mpicomm::MPI.Comm\n \"time of the last callback\"\n lastcbtime_ns::UInt64\n function EveryXWallTimeSeconds(callback, Δtime, mpicomm)\n lastcbtime_ns = zero(UInt64)\n new(callback, Δtime, mpicomm, lastcbtime_ns)\n end\nend\n\nfunction init!(cb::EveryXWallTimeSeconds, solver, Q, param, t)\n cb.lastcbtime_ns = time_ns()\n init!(cb.callback, solver, Q, param, t)\nend\nfunction call!(cb::EveryXWallTimeSeconds, solver, Q, param, t)\n # Check whether we should do a callback\n currtime_ns = time_ns()\n runtime = (currtime_ns - cb.lastcbtime_ns) * 1e-9\n runtime = MPI.Allreduce(runtime, max, cb.mpicomm)\n if runtime < cb.Δtime\n return 0\n else\n # Compute the next time to do a callback\n cb.lastcbtime_ns = currtime_ns\n return call!(cb.callback, solver, Q, param, t)\n end\nend\nfunction fini!(cb::EveryXWallTimeSeconds, solver, Q, param, t)\n fini!(cb.callback, solver, Q, param, t)\nend\n\n\n\"\"\"\n EveryXSimulationTime(f, Δtime)\n\nA wrapper callback to execute `callback` every `time` simulation time seconds.\n\"\"\"\nmutable struct EveryXSimulationTime <: AbstractCallback\n \"callback to wrap\"\n callback::Any\n \"simulation time seconds between callbacks\"\n Δtime::Real\n \"time of the last callback\"\n lastcbtime::Real\n function EveryXSimulationTime(callback, Δtime)\n new(callback, Δtime, 0)\n end\nend\n\nfunction init!(cb::EveryXSimulationTime, solver, Q, param, t)\n cb.lastcbtime = t\n init!(cb.callback, solver, Q, param, t)\nend\nfunction call!(cb::EveryXSimulationTime, solver, Q, param, t)\n # Check whether we should do a callback\n if (t - cb.lastcbtime) < cb.Δtime\n return 0\n else\n # Compute the next time to do a callback\n cb.lastcbtime = t\n return call!(cb.callback, solver, Q, param, t)\n end\nend\nfunction fini!(cb::EveryXSimulationTime, solver, Q, param, t)\n fini!(cb.callback, solver, Q, param, t)\nend\n\n\n\"\"\"\n EveryXSimulationSteps(callback, Δsteps)\n\nA wrapper callback to execute `callback` every `nsteps` of the time stepper.\n\"\"\"\nmutable struct EveryXSimulationSteps <: AbstractCallback\n \"callback to wrap\"\n callback::Any\n \"number of steps between callbacks\"\n Δsteps::Int\n \"number of steps since last callback\"\n steps::Int\n function EveryXSimulationSteps(callback, Δsteps)\n new(callback, Δsteps, 0)\n end\nend\n\nfunction init!(cb::EveryXSimulationSteps, solver, Q, param, t)\n cb.steps = 0\n init!(cb.callback, solver, Q, param, t)\nend\nfunction call!(cb::EveryXSimulationSteps, solver, Q, param, t)\n cb.steps += 1\n if cb.steps < cb.Δsteps\n return 0\n else\n cb.steps = 0\n return call!(cb.callback, solver, Q, param, t)\n end\nend\nfunction fini!(cb::EveryXSimulationSteps, solver, Q, param, t)\n fini!(cb.callback, solver, Q, param, t)\nend\n\nend\n" }, { "path": "src/Numerics/ODESolvers/LowStorageRungeKutta3NMethod.jl", "content": "export LowStorageRungeKutta3N\nexport LS3NRK44Classic, LS3NRK33Heuns\n\n\"\"\"\n LowStorageRungeKutta3N(f, RKA, RKB, RKC, RKW, Q; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThe constructor builds a low-storage Runge--Kutta scheme using 3N storage\nbased on the provided `RKA`, `RKB` and `RKC` coefficient arrays.\n `RKC` (vector of length the number of stages `ns`) set nodal points position;\n `RKA` and `RKB` (size: ns x 2) set weight for tendency and stage-state;\n `RKW` (unused) provides RK weight (last row in Butcher's tableau).\n\nThe 3-N storage formulation from Fyfe (1966) is applicable to any 4-stage,\nfourth-order RK scheme. It is implemented here as:\n\n```math\n\\\\hspace{-20mm} for ~~ j ~~ in ~ [1:ns]: \\\\hspace{10mm}\n t_j = t^n + \\\\Delta t ~ rkC_j\n```\n```math\n dQ_j = dQ^*_j + f(Q_j,t_j)\n```\n```math\n Q_{j+1} = Q_{j} + \\\\Delta t \\\\{ rkB_{j,1} ~ dQ_j + rkB_{j,2} ~ dR_j \\\\}\n```\n```math\n dR_{j+1} = dR_j + rkA_{j+1,2} ~ dQ_j\n```\n```math\n dQ^*_{j+1} = rkA_{j+1,1} ~ dQ_j\n```\n\nThe available concrete implementations are:\n\n - [`LS3NRK44Classic`](@ref)\n - [`LS3NRK33Heuns`](@ref)\n\n### References\n\n @article{Fyfe1966,\n title = {Economical Evaluation of Runge-Kutta Formulae},\n author = {Fyfe, David J.},\n journal = {Mathematics of Computation},\n volume = {20},\n pages = {392--398},\n year = {1966}\n }\n\"\"\"\nmutable struct LowStorageRungeKutta3N{T, RT, AT, Nstages} <: AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n rhs!::Any\n \"Storage for RHS during the `LowStorageRungeKutta3N` update\"\n dQ::AT\n \"Secondary Storage for RHS during the `LowStorageRungeKutta3N` update\"\n dR::AT\n \"low storage RK coefficient array A (rhs scaling)\"\n RKA::Array{RT, 2}\n \"low storage RK coefficient array B (rhs add in scaling)\"\n RKB::Array{RT, 2}\n \"low storage RK coefficient vector C (time scaling)\"\n RKC::Array{RT, 1}\n \"RK weight coefficient vector W (last row in Butcher's tableau)\"\n RKW::Array{RT, 1}\n\n function LowStorageRungeKutta3N(\n rhs!,\n RKA,\n RKB,\n RKC,\n RKW,\n Q::AT;\n dt = 0,\n t0 = 0,\n ) where {AT <: AbstractArray}\n\n T = eltype(Q)\n RT = real(T)\n\n dQ = similar(Q)\n dR = similar(Q)\n fill!(dQ, 0)\n fill!(dR, 0)\n\n new{T, RT, AT, length(RKC)}(\n RT(dt),\n RT(t0),\n 0,\n rhs!,\n dQ,\n dR,\n RKA,\n RKB,\n RKC,\n RKW,\n )\n end\nend\n\n\"\"\"\n dostep!(Q, lsrk3n::LowStorageRungeKutta3N, p, time::Real, nsubsteps::Int,\n iStage::Int, [slow_δ, slow_rv_dQ, slow_scaling])\n\nWrapper function to use the 3N low storage Runge--Kutta method `lsrk3n` as the fast\nsolver for a Multirate Infinitesimal Step method by calling dostep!(Q,\nlsrk3n::LowStorageRungeKutta3N, p, time::Real, [slow_δ, slow_rv_dQ, slow_scaling])\nnsubsteps times.\n\"\"\"\nfunction dostep!(\n Q,\n lsrk3n::LowStorageRungeKutta3N,\n p,\n time::Real,\n nsubsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n for i in 1:nsubsteps\n dostep!(Q, lsrk3n, p, time, slow_δ, slow_rv_dQ, slow_scaling)\n time += lsrk3n.dt\n end\nend\n\n\"\"\"\n dostep!(Q, lsrk3n::LowStorageRungeKutta3N, p, time::Real,\n [slow_δ, slow_rv_dQ, slow_scaling])\n\nUse the 3N low storage Runge--Kutta method `lsrk3n` to step `Q` forward in time\nfrom the current time `time` to final time `time + getdt(lsrk3n)`.\n\nIf the optional parameter `slow_δ !== nothing` then `slow_rv_dQ * slow_δ` is\nadded as an additional ODE right-hand side source. If the optional parameter\n`slow_scaling !== nothing` then after the final stage update the scaling\n`slow_rv_dQ *= slow_scaling` is performed.\n\"\"\"\nfunction dostep!(\n Q,\n lsrk3n::LowStorageRungeKutta3N,\n p,\n time,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n in_slow_scaling = nothing,\n)\n dt = lsrk3n.dt\n\n RKA, RKB, RKC = lsrk3n.RKA, lsrk3n.RKB, lsrk3n.RKC\n rhs!, dQ, dR = lsrk3n.rhs!, lsrk3n.dQ, lsrk3n.dR\n\n rv_Q = realview(Q)\n rv_dQ = realview(dQ)\n rv_dR = realview(dR)\n groupsize = 256\n\n rv_dR .= -0\n for s in 1:length(RKC)\n rhs!(dQ, Q, p, time + RKC[s] * dt, increment = true)\n\n slow_scaling = nothing\n if s == length(RKC)\n slow_scaling = in_slow_scaling\n end\n # update solution and scale RHS\n event = Event(array_device(Q))\n event = update!(array_device(Q), groupsize)(\n rv_dQ,\n rv_dR,\n rv_Q,\n RKA[s % length(RKC) + 1, 1],\n RKA[s % length(RKC) + 1, 2],\n RKB[s, 1],\n RKB[s, 2],\n dt,\n slow_δ,\n slow_rv_dQ,\n slow_scaling;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\nend\n\n@kernel function update!(\n dQ,\n dR,\n Q,\n rka1,\n rka2,\n rkb1,\n rkb2,\n dt,\n slow_δ,\n slow_dQ,\n slow_scaling,\n)\n i = @index(Global, Linear)\n @inbounds begin\n if slow_δ !== nothing\n dQ[i] += slow_δ * slow_dQ[i]\n end\n Q[i] += rkb1 * dt * dQ[i] + rkb2 * dt * dR[i]\n dR[i] += rka2 * dQ[i]\n dQ[i] *= rka1\n if slow_scaling !== nothing\n slow_dQ[i] *= slow_scaling\n end\n end\nend\n\n\"\"\"\n LS3NRK44Classic(f, Q; dt, t0 = 0)\n\nThis function returns a [`LowStorageRungeKutta3N`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the classic 4-stage, fourth-order Runge--Kutta scheme\nin the low-storage implementation of Blum (1962)\n\n### References\n @article {Blum1962,\n title = {A Modification of the Runge-Kutta Fourth-Order Method}\n author = {Blum, E. K.},\n journal = {Mathematics of Computation},\n volume = {16},\n pages = {176-187},\n year = {1962}\n }\n\"\"\"\nfunction LS3NRK44Classic(F, Q::AT; dt = 0, t0 = 0) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n\n RKA = [\n RT(0) RT(0)\n RT(0) RT(1)\n RT(-1 // 2) RT(0)\n RT(2) RT(-6)\n ]\n\n RKB = [\n RT(1 // 2) RT(0)\n RT(1 // 2) RT(-1 // 2)\n RT(1) RT(0)\n RT(1 // 6) RT(1 // 6)\n ]\n\n RKC = [RT(0), RT(1 // 2), RT(1 // 2), RT(1)]\n\n RKW = [RT(1 // 6), RT(1 // 3), RT(1 // 3), RT(1 // 6)]\n\n LowStorageRungeKutta3N(F, RKA, RKB, RKC, RKW, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n LS3NRK33Heuns(f, Q; dt, t0 = 0)\n\nThis function returns a [`LowStorageRungeKutta3N`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis method uses the 3-stage, third-order Heun's Runge--Kutta scheme.\n\n### References\n @article {Heun1900,\n title = {Neue Methoden zur approximativen Integration der\n Differentialgleichungen einer unabh\\\"{a}ngigen Ver\\\"{a}nderlichen}\n author = {Heun, Karl},\n journal = {Z. Math. Phys},\n volume = {45},\n pages = {23--38},\n year = {1900}\n }\n\"\"\"\nfunction LS3NRK33Heuns(\n F,\n Q::AT;\n dt = nothing,\n t0 = 0,\n) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n\n RKA = [\n RT(0) RT(0)\n RT(0) RT(1)\n RT(-1) RT(1 // 3)\n ]\n\n RKB = [\n RT(1 // 3) RT(0)\n RT(2 // 3) RT(-1 // 3)\n RT(3 // 4) RT(1 // 4)\n ]\n\n RKC = [RT(0), RT(1 // 3), RT(2 // 3)]\n\n RKW = [RT(1 // 4), RT(0), RT(3 // 4)]\n\n LowStorageRungeKutta3N(F, RKA, RKB, RKC, RKW, Q; dt = dt, t0 = t0)\nend\n" }, { "path": "src/Numerics/ODESolvers/LowStorageRungeKuttaMethod.jl", "content": "\nexport LowStorageRungeKutta2N\nexport LSRK54CarpenterKennedy, LSRK144NiegemannDiehlBusch, LSRKEulerMethod\n\n\"\"\"\n LowStorageRungeKutta2N(f, RKA, RKB, RKC, Q; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThe constructor builds a low-storage Runge-Kutta scheme using 2N\nstorage based on the provided `RKA`, `RKB` and `RKC` coefficient arrays.\n\nThe available concrete implementations are:\n\n - [`LSRK54CarpenterKennedy`](@ref)\n - [`LSRK144NiegemannDiehlBusch`](@ref)\n\"\"\"\nmutable struct LowStorageRungeKutta2N{T, RT, AT, Nstages} <: AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n rhs!::Any\n \"Storage for RHS during the LowStorageRungeKutta update\"\n dQ::AT\n \"low storage RK coefficient vector A (rhs scaling)\"\n RKA::NTuple{Nstages, RT}\n \"low storage RK coefficient vector B (rhs add in scaling)\"\n RKB::NTuple{Nstages, RT}\n \"low storage RK coefficient vector C (time scaling)\"\n RKC::NTuple{Nstages, RT}\n\n function LowStorageRungeKutta2N(\n rhs!,\n RKA,\n RKB,\n RKC,\n Q::AT;\n dt = 0,\n t0 = 0,\n ) where {AT <: AbstractArray}\n\n T = eltype(Q)\n RT = real(T)\n\n dQ = similar(Q)\n fill!(dQ, 0)\n\n new{T, RT, AT, length(RKA)}(RT(dt), RT(t0), 0, rhs!, dQ, RKA, RKB, RKC)\n end\nend\n\n\"\"\"\n dostep!(Q, lsrk::LowStorageRungeKutta2N, p, time::Real, nsubsteps::Int,\n iStage::Int, [slow_δ, slow_rv_dQ, slow_scaling])\n\nWrapper function to use the 2N low storage Runge--Kutta method `lsrk` as the fast\nsolver for a Multirate Infinitesimal Step method by calling dostep!(Q,\nlsrk::LowStorageRungeKutta2N, p, time::Real, [slow_δ, slow_rv_dQ, slow_scaling])\nnsubsteps times.\n\"\"\"\nfunction dostep!(\n Q,\n lsrk::LowStorageRungeKutta2N,\n p,\n time::Real,\n nsubsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n for i in 1:nsubsteps\n dostep!(Q, lsrk, p, time, slow_δ, slow_rv_dQ, slow_scaling)\n time += lsrk.dt\n end\nend\n\n\"\"\"\n dostep!(Q, lsrk::LowStorageRungeKutta2N, p, time::Real,\n [slow_δ, slow_rv_dQ, slow_scaling])\n\nUse the 2N low storage Runge--Kutta method `lsrk` to step `Q` forward in time\nfrom the current time `time` to final time `time + getdt(lsrk)`.\n\nIf the optional parameter `slow_δ !== nothing` then `slow_rv_dQ * slow_δ` is\nadded as an additionall ODE right-hand side source. If the optional parameter\n`slow_scaling !== nothing` then after the final stage update the scaling\n`slow_rv_dQ *= slow_scaling` is performed.\n\"\"\"\nfunction dostep!(\n Q,\n lsrk::LowStorageRungeKutta2N,\n p,\n time,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n in_slow_scaling = nothing,\n)\n dt = lsrk.dt\n\n RKA, RKB, RKC = lsrk.RKA, lsrk.RKB, lsrk.RKC\n rhs!, dQ = lsrk.rhs!, lsrk.dQ\n\n rv_Q = realview(Q)\n rv_dQ = realview(dQ)\n\n groupsize = 256\n\n for s in 1:length(RKA)\n rhs!(dQ, Q, p, time + RKC[s] * dt, increment = true)\n\n slow_scaling = nothing\n if s == length(RKA)\n slow_scaling = in_slow_scaling\n end\n # update solution and scale RHS\n event = Event(array_device(Q))\n event = update!(array_device(Q), groupsize)(\n rv_dQ,\n rv_Q,\n RKA[s % length(RKA) + 1],\n RKB[s],\n dt,\n slow_δ,\n slow_rv_dQ,\n slow_scaling;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\nend\n\n@kernel function update!(dQ, Q, rka, rkb, dt, slow_δ, slow_dQ, slow_scaling)\n i = @index(Global, Linear)\n @inbounds begin\n if slow_δ !== nothing\n dQ[i] += slow_δ * slow_dQ[i]\n end\n Q[i] += rkb * dt * dQ[i]\n dQ[i] *= rka\n if slow_scaling !== nothing\n slow_dQ[i] *= slow_scaling\n end\n end\nend\n\n\"\"\"\n dostep!(Q, lsrk::LowStorageRungeKutta2N, p::MRIParam, time::Real,\n dt::Real)\n\nUse the 2N low storage Runge--Kutta method `lsrk` to step `Q` forward in time\nfrom the current time `time` to final time `time + dt`.\n\nIf the optional parameter `slow_δ !== nothing` then `slow_rv_dQ * slow_δ` is\nadded as an additionall ODE right-hand side source. If the optional parameter\n`slow_scaling !== nothing` then after the final stage update the scaling\n`slow_rv_dQ *= slow_scaling` is performed.\n\"\"\"\nfunction dostep!(Q, lsrk::LowStorageRungeKutta2N, mrip::MRIParam, time::Real)\n dt = lsrk.dt\n\n RKA, RKB, RKC = lsrk.RKA, lsrk.RKB, lsrk.RKC\n rhs!, dQ = lsrk.rhs!, lsrk.dQ\n\n rv_Q = realview(Q)\n rv_dQ = realview(dQ)\n\n groupsize = 256\n\n for s in 1:length(RKA)\n stage_time = time + RKC[s] * dt\n rhs!(dQ, Q, mrip.p, stage_time, increment = true)\n\n # update solution and scale RHS\n τ = (stage_time - mrip.ts) / mrip.Δts\n event = Event(array_device(Q))\n event = lsrk_mri_update!(array_device(Q), groupsize)(\n rv_dQ,\n rv_Q,\n RKA[s % length(RKA) + 1],\n RKB[s],\n τ,\n dt,\n mrip.γs,\n mrip.Rs;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\nend\n\n@kernel function lsrk_mri_update!(dQ, Q, rka, rkb, τ, dt, γs, Rs)\n i = @index(Global, Linear)\n @inbounds begin\n NΓ = length(γs)\n Ns = length(γs[1])\n dqi = dQ[i]\n\n for s in 1:Ns\n ri = Rs[s][i]\n sc = γs[NΓ][s]\n for k in (NΓ - 1):-1:1\n sc = sc * τ + γs[k][s]\n end\n dqi += sc * ri\n end\n\n Q[i] += rkb * dt * dqi\n dQ[i] = rka * dqi\n end\nend\n\n\n\"\"\"\n LSRKEulerMethod(f, Q; dt, t0 = 0)\n\nThis function returns a [`LowStorageRungeKutta2N`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis method uses the LSRK2N framework to implement a simple Eulerian forward time stepping scheme for the use of debugging.\n\n### References\n\n\"\"\"\nfunction LSRKEulerMethod(\n F,\n Q::AT;\n dt = nothing,\n t0 = 0,\n) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n\n RKA = (RT(0),)\n\n RKB = (RT(1),)\n\n RKC = (RT(0),)\n\n LowStorageRungeKutta2N(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n LSRK54CarpenterKennedy(f, Q; dt, t0 = 0)\n\nThis function returns a [`LowStorageRungeKutta2N`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the fourth-order, low-storage, Runge--Kutta scheme of Carpenter\nand Kennedy (1994) (in their notation (5,4) 2N-Storage RK scheme).\n\n### References\n\n @TECHREPORT{CarpenterKennedy1994,\n author = {M.~H. Carpenter and C.~A. Kennedy},\n title = {Fourth-order {2N-storage} {Runge-Kutta} schemes},\n institution = {National Aeronautics and Space Administration},\n year = {1994},\n number = {NASA TM-109112},\n address = {Langley Research Center, Hampton, VA},\n }\n\"\"\"\nfunction LSRK54CarpenterKennedy(\n F,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n\n RKA = (\n RT(0),\n RT(-567301805773 // 1357537059087),\n RT(-2404267990393 // 2016746695238),\n RT(-3550918686646 // 2091501179385),\n RT(-1275806237668 // 842570457699),\n )\n\n RKB = (\n RT(1432997174477 // 9575080441755),\n RT(5161836677717 // 13612068292357),\n RT(1720146321549 // 2090206949498),\n RT(3134564353537 // 4481467310338),\n RT(2277821191437 // 14882151754819),\n )\n\n RKC = (\n RT(0),\n RT(1432997174477 // 9575080441755),\n RT(2526269341429 // 6820363962896),\n RT(2006345519317 // 3224310063776),\n RT(2802321613138 // 2924317926251),\n )\n\n LowStorageRungeKutta2N(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n LSRK144NiegemannDiehlBusch((f, Q; dt, t0 = 0)\n\nThis function returns a [`LowStorageRungeKutta2N`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the fourth-order, 14-stage, low-storage, Runge--Kutta scheme of\nNiegemann, Diehl, and Busch (2012) with optimized stability region\n\n### References\n - [Niegemann2012](@cite)\n\"\"\"\nfunction LSRK144NiegemannDiehlBusch(\n F,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n\n RKA = (\n RT(0),\n RT(-0.7188012108672410),\n RT(-0.7785331173421570),\n RT(-0.0053282796654044),\n RT(-0.8552979934029281),\n RT(-3.9564138245774565),\n RT(-1.5780575380587385),\n RT(-2.0837094552574054),\n RT(-0.7483334182761610),\n RT(-0.7032861106563359),\n RT(0.0013917096117681),\n RT(-0.0932075369637460),\n RT(-0.9514200470875948),\n RT(-7.1151571693922548),\n )\n\n RKB = (\n RT(0.0367762454319673),\n RT(0.3136296607553959),\n RT(0.1531848691869027),\n RT(0.0030097086818182),\n RT(0.3326293790646110),\n RT(0.2440251405350864),\n RT(0.3718879239592277),\n RT(0.6204126221582444),\n RT(0.1524043173028741),\n RT(0.0760894927419266),\n RT(0.0077604214040978),\n RT(0.0024647284755382),\n RT(0.0780348340049386),\n RT(5.5059777270269628),\n )\n\n RKC = (\n RT(0),\n RT(0.0367762454319673),\n RT(0.1249685262725025),\n RT(0.2446177702277698),\n RT(0.2476149531070420),\n RT(0.2969311120382472),\n RT(0.3978149645802642),\n RT(0.5270854589440328),\n RT(0.6981269994175695),\n RT(0.8190890835352128),\n RT(0.8527059887098624),\n RT(0.8604711817462826),\n RT(0.8627060376969976),\n RT(0.8734213127600976),\n )\n\n LowStorageRungeKutta2N(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n" }, { "path": "src/Numerics/ODESolvers/MultirateInfinitesimalGARKDecoupledImplicit.jl", "content": "export MRIGARKDecoupledImplicit\nexport MRIGARKIRK21aSandu,\n MRIGARKESDIRK34aSandu,\n MRIGARKESDIRK46aSandu,\n MRIGARKESDIRK23LSA,\n MRIGARKESDIRK24LSA\n\n\"\"\"\n MRIGARKDecoupledImplicit(f!, backward_euler_solver, fastsolver, Γs, γ̂s, Q,\n Δt, t0)\n\nConstruct a decoupled implicit MultiRate Infinitesimal General-structure\nAdditive Runge--Kutta (MRI-GARK) scheme to solve\n\n```math\n \\\\dot{y} = f(y, t) + g(y, t)\n```\n\nwhere `f` is the slow tendency function and `g` is the fast tendency function;\nsee Sandu (2019).\n\nThe fast tendency is integrated using the `fastsolver` and the slow tendency\nusing the MRI-GARK scheme. Since this is a decoupled, implicit MRI-GARK there is no implicit coupling between the fast and slow tendencies.\n\nThe `backward_euler_solver` should be of type `AbstractBackwardEulerSolver` or\n`LinearBackwardEulerSolver`, and is used to perform the backward Euler solves\nfor `y` given the slow tendency function, namely\n\n```math\n y = z + α f(y, t; p)\n```\n\nCurrently only [`LowStorageRungeKutta2N`](@ref) schemes are supported for\n`fastsolver`\n\nThe coefficients defined by `γ̂s` can be used for an embedded scheme (only the\nlast stage is different).\n\nThe available concrete implementations are:\n\n - [`MRIGARKIRK21aSandu`](@ref)\n - [`MRIGARKESDIRK34aSandu`](@ref)\n - [`MRIGARKESDIRK46aSandu`](@ref)\n\n### References\n - [Sandu2019](@cite)\n\"\"\"\nmutable struct MRIGARKDecoupledImplicit{\n T,\n RT,\n AT,\n Nstages,\n NΓ,\n FS,\n Nx,\n Ny,\n Nx_Ny,\n BE,\n} <: AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n slowrhs!::Any\n \"backwark Euler solver\"\n besolver!::BE\n \"Storage for RHS during the `MRIGARKDecoupledImplicit` update\"\n Rstages::NTuple{Nstages, AT}\n \"Storage for the implicit solver data vector\"\n Qhat::AT\n \"RK coefficient matrices for coupling coefficients\"\n Γs::NTuple{NΓ, SArray{Tuple{Nx, Ny}, RT, 2, Nx_Ny}}\n \"RK coefficient matrices for embedded scheme\"\n γ̂s::NTuple{NΓ, SArray{NTuple{1, Ny}, RT, 1, Ny}}\n \"RK coefficient vector C (time scaling)\"\n Δc::SArray{NTuple{1, Nstages}, RT, 1, Nstages}\n \"fast solver\"\n fastsolver::FS\n\n function MRIGARKDecoupledImplicit(\n slowrhs!,\n backward_euler_solver,\n fastsolver,\n Γs,\n γ̂s,\n Q::AT,\n dt,\n t0,\n ) where {AT <: AbstractArray}\n NΓ = length(Γs)\n\n T = eltype(Q)\n RT = real(T)\n\n # Compute the Δc coefficients (only explicit values kept)\n Δc = sum(Γs[1], dims = 2)\n\n # Couple of sanity checks on the assumptions of coefficients being of\n # the decoupled implicit structure of Sandu (2019)\n @assert all(isapprox.(Δc[2:2:end], 0; atol = 2 * eps(RT)))\n\n Δc = Δc[1:2:(end - 1)]\n\n # number of slow RHS values we need to keep\n Nstages = length(Δc)\n\n # Couple more sanity checks on the decoupled implicit structure\n @assert Nstages == size(Γs[1], 2) - 1\n @assert Nstages == div(size(Γs[1], 1), 2)\n\n # Scale in the Δc to the Γ and γ̂, and convert to real type\n Γs = ntuple(k -> RT.(Γs[k]), NΓ)\n γ̂s = ntuple(k -> RT.(γ̂s[k]), NΓ)\n\n # Convert to real type\n Δc = RT.(Δc)\n\n # create storage for the stage values\n Rstages = ntuple(i -> similar(Q), Nstages)\n Qhat = similar(Q)\n\n FS = typeof(fastsolver)\n Nx, Ny = size(Γs[1])\n\n # Set up the backward Euler solver with the initial value of α\n α = dt * Γs[1][2, 2]\n besolver! =\n setup_backward_Euler_solver(backward_euler_solver, Q, α, slowrhs!)\n @assert besolver! isa AbstractBackwardEulerSolver\n BE = typeof(besolver!)\n\n new{T, RT, AT, Nstages, NΓ, FS, Nx, Ny, Nx * Ny, BE}(\n RT(dt),\n RT(t0),\n 0,\n slowrhs!,\n besolver!,\n Rstages,\n Qhat,\n Γs,\n γ̂s,\n Δc,\n fastsolver,\n )\n end\nend\n\nfunction updatedt!(mrigark::MRIGARKDecoupledImplicit, dt)\n @assert Δt_is_adjustable(mrigark.besolver!)\n α = dt * mrigark.Γs[1][2, 2]\n update_backward_Euler_solver!(mrigark.besolver!, mrigark.Qhat, α)\n mrigark.dt = dt\nend\n\nfunction dostep!(Q, mrigark::MRIGARKDecoupledImplicit, param, time::Real)\n dt = mrigark.dt\n fast = mrigark.fastsolver\n\n Rs = mrigark.Rstages\n Δc = mrigark.Δc\n\n Nstages = length(Δc)\n groupsize = 256\n\n slowrhs! = mrigark.slowrhs!\n Γs = mrigark.Γs\n NΓ = length(Γs)\n\n ts = time\n groupsize = 256\n\n besolver! = mrigark.besolver!\n Qhat = mrigark.Qhat\n # Since decoupled implicit methods are being used, there is an purely\n # explicit stage followed by an implicit correction stage, hence the divide\n # by two\n for s in 1:Nstages\n # Stage dt\n dts = Δc[s] * dt\n stage_end_time = ts + dts\n\n # initialize the slow tendency stage value\n slowrhs!(Rs[s], Q, param, ts, increment = false)\n\n # # advance fast solution to time stage_end_time\n γs = ntuple(k -> ntuple(j -> Γs[k][2s - 1, j] / Δc[s], s), NΓ)\n mriparam = MRIParam(param, γs, realview.(Rs[1:s]), ts, dts)\n updatetime!(mrigark.fastsolver, ts)\n solve!(Q, mrigark.fastsolver, mriparam; timeend = stage_end_time)\n\n # correct with implicit slow solve\n # Qhat = Q + ∑_j Σ_k Γ_{sjk} dt Rs[j] / k\n # (Divide by k arises from the integration γ_{ij}(τ) in Sandu (2019);\n # see Equation (2.2b) and Definition 2.2\n γs = ntuple(k -> ntuple(j -> dt * Γs[k][2s, j] / k, s), NΓ)\n event = Event(array_device(Q))\n event = mri_create_Qhat!(array_device(Q), groupsize)(\n realview(Qhat),\n realview(Q),\n γs,\n mriparam.Rs;\n ndrange = length(realview(Q)),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n\n # Solve: Q = Qhat + α fslow(Q, stage_end_time)\n α = dt * Γs[1][2s, s + 1]\n besolver!(Q, Qhat, α, param, stage_end_time)\n\n # update time\n ts += dts\n end\nend\n\n# Compute: Qhat = Q + ∑_j Σ_k Γ_{sjk} dt Rs[j] / k\n@kernel function mri_create_Qhat!(Qhat, Q, γs, Rs)\n i = @index(Global, Linear)\n @inbounds begin\n NΓ = length(γs)\n Ns = length(γs[1])\n qhat = Q[i]\n\n for s in 1:Ns\n ri = Rs[s][i]\n sc = γs[1][s]\n for k in 2:NΓ\n sc += γs[k][s]\n end\n qhat += sc * ri\n end\n Qhat[i] = qhat\n end\nend\n\n\"\"\"\n MRIGARKIRK21aSandu(f!, fastsolver, Q; dt, t0 = 0)\n\nThe 2rd order, 2 stage implicit scheme from Sandu (2019).\n\"\"\"\nfunction MRIGARKIRK21aSandu(\n slowrhs!,\n backward_euler_solver,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n)\n #! format: off\n Γ0 = [ 1 // 1 0 // 1\n -1 // 2 1 // 2 ]\n γ̂0 = [-1 // 2 1 // 2 ]\n #! format: on\n MRIGARKDecoupledImplicit(\n slowrhs!,\n backward_euler_solver,\n fastsolver,\n (Γ0,),\n (γ̂0,),\n Q,\n dt,\n t0,\n )\nend\n\n\"\"\"\n MRIGARKESDIRK34aSandu(f!, fastsolver, Q; dt, t0=0)\n\nThe 3rd order, 4 stage decoupled implicit scheme from Sandu (2019).\n\"\"\"\nfunction MRIGARKESDIRK34aSandu(\n slowrhs!,\n backward_euler_solver,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n)\n T = real(eltype(Q))\n μ = acot(2 * sqrt(T(2))) / 3\n λ = 1 - cos(μ) / sqrt(T(2)) + sqrt(T(3 // 2)) * sin(μ)\n @assert isapprox(-1 + 9λ - 18 * λ^2 + 6 * λ^3, 0, atol = 2 * eps(T))\n\n #! format: off\n Γ0 = [\n T(1 // 3) 0 0 0\n -λ λ 0 0\n (3-10λ) / (24λ-6) (5-18λ) / (6-24λ) 0 0\n (-24λ^2+6λ+1) / (6-24λ) (-48λ^2+12λ+1) / (24λ-6) λ 0\n (3-16λ) / (12-48λ) (48λ^2-21λ+2) / (12λ-3) (3-16λ) / 4 0\n -λ 0 0 λ\n ]\n γ̂0 = [ 0 0 0 0]\n #! format: on\n MRIGARKDecoupledImplicit(\n slowrhs!,\n backward_euler_solver,\n fastsolver,\n (Γ0,),\n (γ̂0,),\n Q,\n dt,\n t0,\n )\nend\n\n\"\"\"\n MRIGARKESDIRK46aSandu(f!, fastsolver, Q; dt, t0=0)\n\nThe 4th order, 6 stage decoupled implicit scheme from Sandu (2019).\n\"\"\"\nfunction MRIGARKESDIRK46aSandu(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n)\n T = real(eltype(Q))\n μ = acot(2 * sqrt(T(2))) / 3\n λ = 1 - cos(μ) / sqrt(T(2)) + sqrt(T(3 // 2)) * sin(μ)\n @assert isapprox(-1 + 9λ - 18 * λ^2 + 6 * λ^3, 0, atol = 2 * eps(T))\n\n #! format: off\n Γ0 = [\n 1 // 5 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1\n -1 // 4 1 // 4 0 // 1 0 // 1 0 // 1 0 // 1\n 1771023115159 // 1929363690800 -1385150376999 // 1929363690800 0 // 1 0 // 1 0 // 1 0 // 1\n 914009 // 345800 -1000459 // 345800 1 // 4 0 // 1 0 // 1 0 // 1\n 18386293581909 // 36657910125200 5506531089 // 80566835440 -178423463189 // 482340922700 0 // 1 0 // 1 0 // 1\n 36036097 // 8299200 4621 // 118560 -38434367 // 8299200 1 // 4 0 // 1 0 // 1\n -247809665162987 // 146631640500800 10604946373579 // 14663164050080 10838126175385 // 5865265620032 -24966656214317 // 36657910125200 0 // 1 0 // 1\n 38519701 // 11618880 10517363 // 9682400 -23284701 // 19364800 -10018609 // 2904720 1 // 4 0 // 1\n -52907807977903 // 33838070884800 74846944529257 // 73315820250400 365022522318171 // 146631640500800 -20513210406809 // 109973730375600 -2918009798 // 1870301537 0 // 1\n 19 // 100 -73 // 300 127 // 300 127 // 300 -313 // 300 1 // 4\n ]\n\n Γ1 = [\n 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1\n 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1\n -1674554930619 // 964681845400 1674554930619 // 964681845400 0 // 1 0 // 1 0 // 1 0 // 1\n -1007739 // 172900 1007739 // 172900 0 // 1 0 // 1 0 // 1 0 // 1\n -8450070574289 // 18328955062600 -39429409169 // 40283417720 173621393067 // 120585230675 0 // 1 0 // 1 0 // 1\n -122894383 // 16598400 14501 // 237120 121879313 // 16598400 0 // 1 0 // 1 0 // 1\n 32410002731287 // 15434909526400 -46499276605921 // 29326328100160 -34914135774643 // 11730531240064 45128506783177 // 18328955062600 0 // 1 0 // 1\n -128357303 // 23237760 -35433927 // 19364800 71038479 // 38729600 8015933 // 1452360 0 // 1 0 // 1\n 136721604296777 // 67676141769600 -349632444539303 // 146631640500800 -1292744859249609 // 293263281001600 8356250416309 // 54986865187800 17282943803 // 3740603074 0 // 1\n 3 // 25 -29 // 300 71 // 300 71 // 300 -149 // 300 0 // 1\n ]\n\n γ̂0 = [-1 // 4 5595 // 8804 -2445 // 8804 -4225 // 8804 2205 // 4402 -567 // 4402]\n γ̂1 = [ 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1 ]\n #! format: on\n MRIGARKDecoupledImplicit(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n (Γ0, Γ1),\n (γ̂0, γ̂1),\n Q,\n dt,\n t0,\n )\nend\n\n\"\"\"\n MRIGARKESDIRK23LSA(f!, fastsolver, Q; dt, t0 = 0, δ = 0\n\nA 2nd order, 3 stage decoupled implicit scheme. It is based on L-Stable,\nstiffly-accurate ESDIRK scheme of Bank et al (1985); see also Kennedy and\nCarpenter (2016).\n\nThe free parameter `δ` can take any values for accuracy.\n\n### References\n - [Bank1985](@cite)\n - [KennedyCarpenter2016](@cite)\n\"\"\"\nfunction MRIGARKESDIRK23LSA(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n δ = 0,\n)\n T = real(eltype(Q))\n rt2 = sqrt(T(2))\n\n #! format: off\n Γ0 = [\n 2 - rt2 0 0\n (1 - rt2) / rt2 (rt2 - 1) / rt2 0\n δ rt2 - 1 - δ 0\n (3 - 2rt2 * (1 + δ)) / 2rt2 (δ * 2rt2 - 1) / 2rt2 (rt2 - 1) / rt2\n ]\n\n γ̂0 = [0 0 0]\n #! format: on\n Δc = sum(Γ0, dims = 2)\n\n # Check that the explicit steps match the Δc values:\n @assert Γ0[1, 1] ≈ Δc[1]\n @assert Γ0[3, 1] + Γ0[3, 2] ≈ Δc[3]\n\n # Check the implicit stages have no Δc\n @assert isapprox(Γ0[2, 1] + Γ0[2, 2], 0, atol = eps(T))\n @assert isapprox(Γ0[4, 1] + Γ0[4, 2] + Γ0[4, 3], 0, atol = eps(T))\n\n # Check consistency with the original scheme\n @assert Γ0[1, 1] + Γ0[2, 1] ≈ 1 - 1 / rt2\n @assert Γ0[2, 2] ≈ 1 - 1 / rt2\n\n @assert Γ0[1, 1] + Γ0[2, 1] + Γ0[3, 1] + Γ0[4, 1] ≈ 1 / (2 * rt2)\n @assert Γ0[2, 2] + Γ0[3, 2] + Γ0[4, 2] ≈ 1 / (2 * rt2)\n @assert Γ0[4, 3] ≈ 1 - 1 / rt2\n\n\n MRIGARKDecoupledImplicit(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n (Γ0,),\n (γ̂0,),\n Q,\n dt,\n t0,\n )\nend\n\n\"\"\"\n MRIGARKESDIRK24LSA(f!,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n γ = 0.2,\n c3 = (2γ + 1) / 2,\n a32 = 0.2,\n α = -0.1,\n β1 = c3 / 10,\n β2 = c3 / 10,\n )\n\nA 2nd order, 4 stage decoupled implicit scheme. It is based on an L-Stable,\nstiffly-accurate ESDIRK.\n\"\"\"\nfunction MRIGARKESDIRK24LSA(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n Q;\n dt,\n t0 = 0,\n γ = 0.2,\n c3 = (2γ + 1) / 2,\n a32 = 0.2,\n α = -0.1,\n β1 = c3 / 10,\n β2 = c3 / 10,\n)\n T = real(eltype(Q))\n\n # Check L-Stability constraint; bound comes from Kennedy and Carpenter\n # (2016) Table 5.\n @assert 0.1804253064293985641345831 ≤ γ < 1 // 2\n\n # check the stage times are increasing\n @assert 2γ < c3 < 1\n\n # Original RK scheme\n # Enforce L-Stability\n b3 = (2 * (1 - γ)^2 - 1) / 4 / a32\n\n # Enforce 2nd order accuracy\n b2 = (1 - 2γ - 2b3 * c3) / 4γ\n\n A = [\n 0 0 0 0\n γ γ 0 0\n c3 - a32-γ a32 γ 0\n 1 - b2 - b3-γ b2 b3 γ\n ]\n c = sum(A, dims = 2)\n b = A[end, :]\n\n # Check 2nd order accuracy\n @assert sum(b) ≈ 1\n @assert 2 * sum(A' * b) ≈ 1\n\n # Setup the GARK Tableau\n Δc = [c[2], 0, c[3] - c[2], 0, c[4] - c[3], 0]\n\n Γ0 = zeros(T, 6, 4)\n Γ0[1, 1] = Δc[1]\n\n Γ0[2, 1] = A[2, 1] - Γ0[1, 1]\n Γ0[2, 2] = A[2, 2]\n\n Γ0[3, 1] = α\n Γ0[3, 2] = Δc[3] - Γ0[3, 1]\n\n Γ0[4, 1] = A[3, 1] - Γ0[1, 1] - Γ0[2, 1] - Γ0[3, 1]\n Γ0[4, 2] = A[3, 2] - Γ0[1, 2] - Γ0[2, 2] - Γ0[3, 2]\n Γ0[4, 3] = A[3, 3]\n\n Γ0[5, 1] = β1\n Γ0[5, 2] = β2\n Γ0[5, 3] = Δc[5] - Γ0[5, 1] - Γ0[5, 2]\n\n Γ0[6, 1] = A[4, 1] - Γ0[1, 1] - Γ0[2, 1] - Γ0[3, 1] - Γ0[4, 1] - Γ0[5, 1]\n Γ0[6, 2] = A[4, 2] - Γ0[1, 2] - Γ0[2, 2] - Γ0[3, 2] - Γ0[4, 2] - Γ0[5, 2]\n Γ0[6, 3] = A[4, 3] - Γ0[1, 3] - Γ0[2, 3] - Γ0[3, 3] - Γ0[4, 3] - Γ0[5, 3]\n Γ0[6, 4] = A[4, 4]\n\n γ̂0 = [0 0 0 0]\n\n # Check consistency with original scheme\n @assert all(A ≈ [0 0 0 0; accumulate(+, Γ0, dims = 1)[2:2:end, :]])\n @assert all(Δc ≈ sum(Γ0, dims = 2))\n\n MRIGARKDecoupledImplicit(\n slowrhs!,\n implicitsolve!,\n fastsolver,\n (Γ0,),\n (γ̂0,),\n Q,\n dt,\n t0,\n )\nend\n" }, { "path": "src/Numerics/ODESolvers/MultirateInfinitesimalGARKExplicit.jl", "content": "export MRIGARKExplicit\nexport MRIGARKERK33aSandu, MRIGARKERK45aSandu\n\n\"\"\"\n MRIParam(p, γs, Rs, ts, Δts)\n\nConstruct a type for passing the data around for the `MRIGARKExplicit` explicit\ntime stepper to follow on methods. `p` is the original user defined ODE\nparameters, `γs` and `Rs` are the MRI parameters and stage values, respectively.\n`ts` and `Δts` are the stage time and stage time step.\n\"\"\"\nstruct MRIParam{P, T, AT, N, M}\n p::P\n γs::NTuple{M, SArray{NTuple{1, N}, T, 1, N}}\n Rs::NTuple{N, AT}\n ts::T\n Δts::T\n function MRIParam(\n p::P,\n γs::NTuple{M},\n Rs::NTuple{N, AT},\n ts,\n Δts,\n ) where {P, M, N, AT}\n T = eltype(γs[1])\n new{P, T, AT, N, M}(p, γs, Rs, ts, Δts)\n end\nend\n\n# We overload get property to access the original param\nfunction Base.getproperty(mriparam::MRIParam, s::Symbol)\n if s === :p\n p = getfield(mriparam, :p)\n return p isa MRIParam ? p.p : p\n else\n getfield(mriparam, s)\n end\nend\n\n\"\"\"\n MRIGARKExplicit(f!, fastsolver, Γs, γ̂s, Q, Δt, t0)\n\nConstruct an explicit MultiRate Infinitesimal General-structure Additive\nRunge--Kutta (MRI-GARK) scheme to solve\n\n```math\n \\\\dot{y} = f_{slow}(y, t) + f_{fast}(y, t)\n```\n\nwhere `f_{slow}` is the slow tendency function and `f_{fast}` is the fast tendency function;\nsee Sandu (2019).\n\nThe fast tendency is integrated using the `fastsolver` and the slow tendency\nusing the MRI-GARK scheme. Namely, at each stage the scheme solves\n\n```math\n\\\\begin{aligned}\n v(T_i) &= Y_i \\\\\\\\\n \\\\dot{v} &= f(v, t) + \\\\sum_{j=1}^{i} \\\\bar{γ}_{ij}(t) R_j \\\\\\\\\n \\\\bar{γ}_{ijk}(t) &= \\\\sum_{k=0}^{NΓ-1} γ_{ijk} τ(t)^k / Δc_s \\\\\\\\\n τ(t) &= (t - t_s) / Δt \\\\\\\\\n Y_{i+1} &= v(T_i + c_s * Δt)\n\\\\end{aligned}\n```\n\nwhere ``Y_1 = y_n`` and ``y_{n+1} = Y_{Nstages+1}``.\n\nHere ``R_j = g(Y_j, t_0 + c_j * Δt)`` is the tendency for stage ``j``,\n``γ_{ijk}`` are the GARK coupling coefficients,\n``NΓ`` is the number of sets of GARK coupling coefficients there are\n``Δc_s = \\\\sum_{j=1}^{Nstages} γ_{sj1} = c_{s+1} - c_s`` is the scaling\nincrement between stage times. The ODE for ``v(t)`` is solved using the\n`fastsolver`. Note that this form of the scheme is based on Definition 2.2 of\nSandu (2019), but ODE for ``v(t)`` is written to go from ``t_s`` to\n``T_i + c_s * Δt`` as opposed to ``0`` to ``1``.\n\nCurrently only [`LowStorageRungeKutta2N`](@ref) schemes are supported for\n`fastsolver`\n\nThe coefficients defined by `γ̂s` can be used for an embedded scheme (only the\nlast stage is different).\n\nThe available concrete implementations are:\n\n - [`MRIGARKERK33aSandu`](@ref)\n - [`MRIGARKERK45aSandu`](@ref)\n\n### References\n - [Sandu2019](@cite)\n\"\"\"\nmutable struct MRIGARKExplicit{T, RT, AT, Nstages, NΓ, FS, Nstages_sq} <:\n AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n slowrhs!::Any\n \"Storage for RHS during the `MRIGARKExplicit` update\"\n Rstages::NTuple{Nstages, AT}\n \"RK coefficient matrices for coupling coefficients\"\n Γs::NTuple{NΓ, SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}}\n \"RK coefficient matrices for embedded scheme\"\n γ̂s::NTuple{NΓ, SArray{NTuple{1, Nstages}, RT, 1, Nstages}}\n \"RK coefficient vector C (time scaling)\"\n Δc::SArray{NTuple{1, Nstages}, RT, 1, Nstages}\n \"fast solver\"\n fastsolver::FS\n\n function MRIGARKExplicit(\n slowrhs!,\n fastsolver,\n Γs,\n γ̂s,\n Q::AT,\n dt,\n t0,\n ) where {AT <: AbstractArray}\n NΓ = length(Γs)\n Nstages = size(Γs[1], 1)\n T = eltype(Q)\n RT = real(T)\n\n # Compute the Δc coefficients\n Δc = sum(Γs[1], dims = 2)[:]\n\n # Scale in the Δc to the Γ and γ̂, and convert to real type\n Γs = ntuple(k -> RT.(Γs[k] ./ Δc), NΓ)\n γ̂s = ntuple(k -> RT.(γ̂s[k] / Δc[Nstages]), NΓ)\n\n # Convert to real type\n Δc = RT.(Δc)\n\n # create storage for the stage values\n Rstages = ntuple(i -> similar(Q), Nstages)\n\n FS = typeof(fastsolver)\n new{T, RT, AT, Nstages, NΓ, FS, Nstages^2}(\n RT(dt),\n RT(t0),\n 0,\n slowrhs!,\n Rstages,\n Γs,\n γ̂s,\n Δc,\n fastsolver,\n )\n end\nend\n\nfunction dostep!(Q, mrigark::MRIGARKExplicit, param, time::Real)\n dt = mrigark.dt\n fast = mrigark.fastsolver\n\n Rs = mrigark.Rstages\n Δc = mrigark.Δc\n Nstages = length(Δc)\n slowrhs! = mrigark.slowrhs!\n Γs = mrigark.Γs\n NΓ = length(Γs)\n\n ts = time\n groupsize = 256\n for s in 1:Nstages\n # Stage dt\n dts = Δc[s] * dt\n\n p = param isa MRIParam ? param.p : param\n slowrhs!(Rs[s], Q, p, ts, increment = false)\n if param isa MRIParam\n # fraction of the step slower stage increment we are on\n τ = (ts - param.ts) / param.Δts\n event = Event(array_device(Q))\n event = mri_update_rate!(array_device(Q), groupsize)(\n realview(Rs[s]),\n τ,\n param.γs,\n param.Rs;\n ndrange = length(realview(Rs[s])),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n γs = ntuple(k -> ntuple(j -> Γs[k][s, j], s), NΓ)\n mriparam = MRIParam(param, γs, realview.(Rs[1:s]), ts, dts)\n updatetime!(mrigark.fastsolver, ts)\n solve!(Q, mrigark.fastsolver, mriparam; timeend = ts + dts)\n\n # update time\n ts += dts\n end\nend\n\n@kernel function mri_update_rate!(dQ, τ, γs, Rs)\n i = @index(Global, Linear)\n @inbounds begin\n NΓ = length(γs)\n Ns = length(γs[1])\n dqi = dQ[i]\n\n for s in 1:Ns\n ri = Rs[s][i]\n sc = γs[NΓ][s]\n for k in (NΓ - 1):-1:1\n sc = sc * τ + γs[k][s]\n end\n dqi += sc * ri\n end\n\n dQ[i] = dqi\n end\nend\n\n\"\"\"\n MRIGARKERK33aSandu(f!, fastsolver, Q; dt, t0 = 0, δ = -1 // 2)\n\nThe 3rd order, 3 stage scheme from Sandu (2019). The parameter `δ` defaults to\nthe value suggested by Sandu, but can be varied.\n\"\"\"\nfunction MRIGARKERK33aSandu(slowrhs!, fastsolver, Q; dt, t0 = 0, δ = -1 // 2)\n T = eltype(Q)\n RT = real(T)\n #! format: off\n Γ0 = [\n 1 // 3 0 // 1 0 // 1\n (-6δ - 7) // 12 (6δ + 11) // 12 0 // 1\n 0 // 1 (6δ - 5) // 12 (3 - 2δ) // 4\n ]\n γ̂0 = [ 1 // 12 -1 // 3 7 // 12]\n\n Γ1 = [\n 0 // 1 0 // 1 0 // 1\n (2δ + 1) // 2 -(2δ + 1) // 2 0 // 1\n 1 // 2 -(2δ + 1) // 2 δ // 1\n ]\n γ̂1 = [ 0 // 1 0 // 1 0 // 1]\n #! format: on\n MRIGARKExplicit(slowrhs!, fastsolver, (Γ0, Γ1), (γ̂0, γ̂1), Q, dt, t0)\nend\n\n\"\"\"\n MRIGARKERK45aSandu(f!, fastsolver, Q; dt, t0 = 0)\n\nThe 4th order, 5 stage scheme from Sandu (2019).\n\"\"\"\nfunction MRIGARKERK45aSandu(slowrhs!, fastsolver, Q; dt, t0 = 0)\n T = eltype(Q)\n RT = real(T)\n #! format: off\n Γ0 = [\n 1 // 5 0 // 1 0 // 1 0 // 1 0 // 1\n -53 // 16 281 // 80 0 // 1 0 // 1 0 // 1\n -36562993 // 71394880 34903117 // 17848720 -88770499 // 71394880 0 // 1 0 // 1\n -7631593 // 71394880 -166232021 // 35697440 6068517 // 1519040 8644289 // 8924360 0 // 1\n 277061 // 303808 -209323 // 1139280 -1360217 // 1139280 -148789 // 56964 147889 // 45120\n ]\n γ̂0 = [-1482837 // 759520 175781 // 71205 -790577 // 1139280 -6379 // 56964 47 // 96]\n Γ1 = [\n 0 // 1 0 // 1 0 // 1 0 // 1 0 // 1\n 503 // 80 -503 // 80 0 // 1 0 // 1 0 // 1\n -1365537 // 35697440 4963773 // 7139488 -1465833 // 2231090 0 // 1 0 // 1\n 66974357 // 35697440 21445367 // 7139488 -3 // 1 -8388609 // 4462180 0 // 1\n -18227 // 7520 2 // 1 1 // 1 5 // 1 -41933 // 7520\n ]\n γ̂1 = [ 6213 // 1880 -6213 // 1880 0 // 1 0 // 1 0 // 1]\n #! format: on\n MRIGARKExplicit(slowrhs!, fastsolver, (Γ0, Γ1), (γ̂0, γ̂1), Q, dt, t0)\nend\n" }, { "path": "src/Numerics/ODESolvers/MultirateInfinitesimalStepMethod.jl", "content": "\nexport MultirateInfinitesimalStep,\n TimeScaledRHS,\n MISRK1,\n MIS2,\n MISRK2a,\n MISRK2b,\n MIS3C,\n MISRK3,\n MIS4,\n MIS4a,\n MISKWRK43,\n TVDMISA,\n TVDMISB,\n getnsubsteps\n\n\"\"\"\n TimeScaledRHS(a, b, rhs!)\n\nWhen evaluate at time `t`, evaluates `rhs!` at time `a + bt`.\n\"\"\"\nmutable struct TimeScaledRHS{N, RT}\n a::RT\n b::RT\n rhs!::Any\n function TimeScaledRHS(a, b, rhs!)\n RT = typeof(a)\n if isa(rhs!, Tuple)\n N = length(rhs!)\n else\n N = 1\n end\n new{N, RT}(a, b, rhs!)\n end\nend\n\nfunction (o::TimeScaledRHS{1, RT} where {RT})(dQ, Q, params, tau; increment)\n o.rhs!(dQ, Q, params, o.a + o.b * tau; increment = increment)\nend\n\nfunction (o::TimeScaledRHS{2, RT} where {RT})(dQ, Q, params, tau, i; increment)\n o.rhs![i](dQ, Q, params, o.a + o.b * tau; increment = increment)\nend\n\nmutable struct OffsetRHS{AT}\n offset::AT\n rhs!::Any\n function OffsetRHS(offset, rhs!)\n AT = typeof(offset)\n new{AT}(offset, rhs!)\n end\nend\n\nfunction (o::OffsetRHS{AT} where {AT})(dQ, Q, params, tau; increment)\n o.rhs!(dQ, Q, params, tau; increment = increment)\n dQ .+= o.offset\nend\n\n\"\"\"\n MultirateInfinitesimalStep(slowrhs!, fastrhs!, fastmethod,\n α, β, γ,\n Q::AT; dt=0, t0=0) where {AT<:AbstractArray}\n\nThis is a time stepping object for explicitly time stepping the partitioned differential\nequation given by right-hand-side functions `f_fast` and `f_slow` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f_{fast}(Q, t) + f_{slow}(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThe constructor builds a multirate infinitesimal step Runge-Kutta scheme\nbased on the provided `α`, `β` and `γ` tableaux and `fastmethod` for solving\nthe fast modes.\n\nThe available concrete implementations are:\n\n - [`MISRK1`](@ref)\n - [`MIS2`](@ref)\n - [`MISRK2a`](@ref)\n - [`MISRK2b`](@ref)\n - [`MIS3C`](@ref)\n - [`MISRK3`](@ref)\n - [`MIS4`](@ref)\n - [`MIS4a`](@ref)\n - [`MISKWRK43`](@ref)\n - [`TVDMISA`](@ref)\n - [`TVDMISB`](@ref)\n\n### References\n - [KnothWensch2014](@cite)\n - [WickerSkamarock2002](@cite)\n - [KnothWolke1998](@cite)\n\"\"\"\nmutable struct MultirateInfinitesimalStep{\n T,\n RT,\n AT,\n FS,\n Nstages,\n Nstagesm1,\n Nstagesm2,\n Nstages_sq,\n} <: AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"storage for y_n\"\n yn::AT\n \"Storage for ``Y_nj - y_n``\"\n ΔYnj::NTuple{Nstagesm2, AT}\n \"Storage for ``f(Y_nj)``\"\n fYnj::NTuple{Nstagesm1, AT}\n \"Storage for offset\"\n offset::AT\n \"slow rhs function\"\n slowrhs!::Any\n \"RHS for fast solver\"\n tsfastrhs!::TimeScaledRHS{N, RT} where {N}\n \"fast rhs method\"\n fastsolver::FS\n \"number of substeps per stage\"\n nsubsteps::Int\n α::SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}\n β::SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}\n γ::SArray{NTuple{2, Nstages}, RT, 2, Nstages_sq}\n d::SArray{NTuple{1, Nstages}, RT, 1, Nstages}\n c::SArray{NTuple{1, Nstages}, RT, 1, Nstages}\n c̃::SArray{NTuple{1, Nstages}, RT, 1, Nstages}\n function MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q::AT;\n dt = 0,\n t0 = 0,\n ) where {AT <: AbstractArray}\n\n T = eltype(Q)\n RT = real(T)\n\n Nstages = size(α, 1)\n\n yn = similar(Q)\n ΔYnj = ntuple(_ -> similar(Q), Nstages - 2)\n fYnj = ntuple(_ -> similar(Q), Nstages - 1)\n offset = similar(Q)\n tsfastrhs! = TimeScaledRHS(RT(0), RT(0), fastrhs!)\n fastsolver = fastmethod(tsfastrhs!, Q)\n\n d = sum(β, dims = 2)\n\n c = similar(d)\n for i in eachindex(c)\n c[i] = d[i]\n if i > 1\n c[i] += sum(j -> (α[i, j] + γ[i, j]) * c[j], 1:(i - 1))\n end\n\n # When d[i] = 0, we do not perform fast substepping, therefore\n # we do not need to scale the β, γ coefficients\n if !(abs(d[i]) < 1.e-10)\n β[i, :] ./= d[i]\n γ[i, :] ./= d[i]\n end\n end\n c̃ = α * c\n\n new{\n T,\n RT,\n AT,\n typeof(fastsolver),\n Nstages,\n Nstages - 1,\n Nstages - 2,\n Nstages^2,\n }(\n RT(dt),\n RT(t0),\n 0,\n yn,\n ΔYnj,\n fYnj,\n offset,\n slowrhs!,\n tsfastrhs!,\n fastsolver,\n nsubsteps,\n α,\n β,\n γ,\n d,\n c,\n c̃,\n )\n end\nend\n\nfunction MultirateInfinitesimalStep(\n mis,\n op::TimeScaledRHS{2, RT} where {RT},\n fastmethod,\n Q = nothing;\n dt = 0,\n t0 = 0,\n nsubsteps = 1,\n) where {AT <: AbstractArray}\n\n return mis(\n op.rhs![1],\n op.rhs![2],\n fastmethod,\n nsubsteps,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction dostep!(\n Q,\n mis::MultirateInfinitesimalStep,\n p,\n time::Real,\n nsubsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n if isa(mis.slowrhs!, OffsetRHS{AT} where {AT})\n mis.slowrhs!.offset = slow_rv_dQ\n else\n mis.slowrhs! = OffsetRHS(slow_rv_dQ, mis.slowrhs!)\n end\n for i in 1:nsubsteps\n dostep!(Q, mis, p, time)\n time += mis.fastsolver.dt\n end\nend\n\nfunction dostep!(Q, mis::MultirateInfinitesimalStep, p, time)\n dt = mis.dt\n FT = eltype(dt)\n α = mis.α\n β = mis.β\n γ = mis.γ\n yn = mis.yn\n ΔYnj = mis.ΔYnj\n fYnj = mis.fYnj\n offset = mis.offset\n d = mis.d\n c = mis.c\n c̃ = mis.c̃\n slowrhs! = mis.slowrhs!\n fastsolver = mis.fastsolver\n fastrhs! = mis.tsfastrhs!\n nsubsteps = mis.nsubsteps\n\n nstages = size(α, 1)\n\n copyto!(yn, Q) # first stage\n for i in 2:nstages\n slowrhs!(fYnj[i - 1], Q, p, time + c[i - 1] * dt, increment = false)\n\n groupsize = 256\n event = Event(array_device(Q))\n event = update!(array_device(Q), groupsize)(\n realview(Q),\n realview(offset),\n Val(i),\n realview(yn),\n map(realview, ΔYnj[1:(i - 2)]),\n map(realview, fYnj[1:(i - 1)]),\n α[i, :],\n β[i, :],\n γ[i, :],\n dt;\n ndrange = length(realview(Q)),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n\n # When d[i] = 0, we do not perform fast substepping;\n # instead we just update the slow tendency\n if iszero(d[i])\n Q .+= dt .* offset\n else\n fastrhs!.a = time + c̃[i] * dt\n fastrhs!.b = (c[i] - c̃[i]) / d[i]\n\n τ = zero(FT)\n nsubstepsLoc = ceil(Int, nsubsteps * d[i])\n dτ = d[i] * dt / nsubstepsLoc\n updatetime!(fastsolver, τ)\n updatedt!(fastsolver, dτ)\n # TODO: we want to be able to write\n # solve!(Q, fastsolver, p; numberofsteps = mis.nsubsteps) #(1c)\n # especially if we want to use StormerVerlet, but need some way to pass in `offset`\n dostep!(\n Q,\n fastsolver,\n p,\n τ,\n nsubstepsLoc,\n i,\n FT(1),\n realview(offset),\n nothing,\n ) #(1c)\n end\n end\nend\n\n@kernel function update!(\n Q,\n offset,\n ::Val{i},\n yn,\n ΔYnj,\n fYnj,\n αi,\n βi,\n γi,\n dt,\n) where {i}\n e = @index(Global, Linear)\n @inbounds begin\n if i > 2\n ΔYnj[i - 2][e] = Q[e] - yn[e] # is 0 for i == 2\n end\n Q[e] = yn[e] # (1a)\n offset[e] = (βi[1]) .* fYnj[1][e] # (1b)\n @unroll for j in 2:(i - 1)\n Q[e] += αi[j] .* ΔYnj[j - 1][e] # (1a cont.)\n offset[e] += (γi[j] / dt) * ΔYnj[j - 1][e] + βi[j] * fYnj[j][e] # (1b cont.)\n end\n end\nend\n\n\"\"\"\n MISRK1(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MISRK1` method is a 1st-order accurate MIS method based\non the RK1 (explicit Euler) method.\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction MISRK1(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n α = zeros(2, 2)\n β = beta(MISRK1, RT)\n γ = zeros(2, 2)\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MISRK1), RT::DataType)\n β = [\n 0 0\n 1 0\n ]\nend\n\n\"\"\"\n MIS2(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MIS2` method is a 2nd-order accurate, 3-stage MIS method whose\nconstruction is summarized in Table 1 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction MIS2(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.536946566710) 0 0\n 0 RT(0.480892968551) RT(0.500561163566) 0\n ]\n\n β = beta(MIS2, RT)\n\n γ = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.652465126004) 0 0\n 0 RT(-0.0732769849457) RT(0.144902430420) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MIS2), RT::DataType)\n β = [\n 0 0 0 0\n RT(0.126848494553) 0 0 0\n RT(-0.784838278826) RT(1.37442675268) 0 0\n RT(-0.0456727081749) RT(-0.00875082271190) RT(0.524775788629) 0\n ]\nend\n\n\"\"\"\n MISRK2a(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MISRK2a` method is a 2nd-order accurate, 2-stage MIS method\nbased on the approach detailed by Wicker and Skamarock in\n[WickerSkamarock2002](@cite).\n\n### References\n - [WickerSkamarock2002](@cite)\n\"\"\"\nfunction MISRK2a(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0\n 0 0 0\n 0 1 0\n ]\n\n β = beta(MISRK2a, RT)\n\n γ = zeros(3, 3)\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MISRK2a), RT::DataType)\n β = [\n 0 0 0\n RT(0.5) 0 0\n RT(-0.5) 1 0\n ]\nend\n\n\"\"\"\n MISRK2b(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MISRK2b` method is a 2nd-order accurate, 2-stage MIS method and a\nvariant of the `MISRK2a` method based on the approach detailed by Wicker\nand Skamarock in [WickerSkamarock2002](@cite).\n\n### References\n - [WickerSkamarock2002](@cite)\n\"\"\"\nfunction MISRK2b(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0\n 0 0 0\n 0 1 0\n ]\n\n β = beta(MISRK2b, RT)\n\n γ = zeros(3, 3)\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MISRK2b), RT::DataType)\n β = [\n 0 0 0\n 1 0 0\n RT(-0.5) RT(0.5) 0\n ]\nend\n\n\"\"\"\n MIS3C(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MIS3C` method is a 3rd-order accurate, 3-stage MIS method whose\nconstruction is summarized in Table 2 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction MIS3C(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.589557277145) 0 0\n 0 RT(0.544036601551) RT(0.565511042564) 0\n ]\n\n β = beta(MIS3C, RT)\n\n γ = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.142798786398) 0 0\n 0 RT(-0.0428918957402) RT(0.0202720980282) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MIS3C), RT::DataType)\n β = [\n 0 0 0 0\n RT(0.397525189225) 0 0 0\n RT(-0.227036463644) RT(0.624528794618) 0 0\n RT(-0.00295238076840) RT(-0.270971764284) RT(0.671323159437) 0\n ]\nend\n\n\"\"\"\n MISRK3(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MISRK3` method is a 3rd-order accurate, 3-stage MIS method\nbased on the approach detailed by Wicker and Skamarock in\n[WickerSkamarock2002](@cite).\n\n### References\n - [WickerSkamarock2002](@cite)\n\"\"\"\nfunction MISRK3(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n α = zeros(4, 4)\n β = beta(MISRK3, RT)\n γ = zeros(4, 4)\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MISRK3), RT::DataType)\n β = [\n 0 0 0 0\n RT(0.3333333333333333) 0 0 0\n 0 RT(0.5) 0 0\n 0 0 1 0\n ]\nend\n\n\"\"\"\n MIS4(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MIS4` method is a 3rd-order accurate, 4-stage MIS method whose\nconstruction is summarized in Table 3 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction MIS4(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0 0\n 0 0 0 0 0\n 0 RT(0.914092810304) 0 0 0\n 0 RT(1.14274417397) RT(-0.295211246188) 0 0\n 0 RT(0.112965282231) RT(0.337369411296) RT(0.503747183119) 0\n ]\n\n β = beta(MIS4, RT)\n\n γ = [\n 0 0 0 0 0\n 0 0 0 0 0\n 0 RT(0.678951983291) 0 0 0\n 0 RT(-1.38974164070) RT(0.503864576302) 0 0\n 0 RT(-0.375328608282) RT(0.320925021109) RT(-0.158259688945) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MIS4), RT::DataType)\n β = [\n 0 0 0 0 0\n RT(0.136296478423) 0 0 0 0\n RT(0.280462398979) RT(-0.0160351333596) 0 0 0\n RT(0.904713355208) RT(-1.04011183154) RT(0.652337563489) 0 0\n RT(0.0671969845546) RT(-0.365621862610) RT(-0.154861470835) RT(0.970362444469) 0\n ]\nend\n\n\"\"\"\n MIS4a(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MIS4a` method is a 3rd-order accurate, 4-stage MIS method whose\nconstruction is summarized in Table 4 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction MIS4a(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0 0\n 0 0 0 0 0\n 0 RT(0.52349249922385610) 0 0 0\n 0 RT(1.1683374366893629) RT(-0.75762080241712637) 0 0\n 0 RT(-0.036477233846797109) RT(0.56936148730740477) RT(0.47746263002599681) 0\n ]\n\n # β[5,1] in the paper is incorrect\n # the correct value is used below (from authors)\n β = beta(MIS4a, RT)\n\n γ = [\n 0 0 0 0 0\n 0 0 0 0 0\n 0 RT(0.13145089796226542) 0 0 0\n 0 RT(-0.36855857648747881) RT(0.33159232636600550) 0 0\n 0 RT(-0.065767130537473045) RT(0.040591093109036858) RT(0.064902111640806712) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MIS4a), RT::DataType)\n β = [\n 0 0 0 0 0\n RT(0.38758444641450318) 0 0 0 0\n RT(-0.025318448354142823) RT(0.38668943087310403) 0 0 0\n RT(0.20899983523553325) RT(-0.45856648476371231) RT(0.43423187573425748) 0 0\n RT(-0.10048822195663100) RT(-0.46186171956333327) RT(0.83045062122462809) RT(0.27014914900250392) 0\n ]\nend\n\n\"\"\"\n MISKWRK43(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `MISKWRK43` method is a 3rd-order accurate, 4-stage MIS method. It is the\nMIS analog of an RK43 method, based on the approach detailed in\n[KnothWolke1998](@cite).\n\n### References\n - [KnothWolke1998](@cite)\n\"\"\"\nfunction MISKWRK43(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0 0\n 0 0 0 0 0\n 0 1 0 0 0\n 0 0 1 0 0\n 0 0 0 1 0\n ]\n\n β = beta(MISKWRK43, RT)\n\n γ = zeros(5, 5)\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(MISKWRK43), RT::DataType)\n β = [\n 0 0 0 0 0\n 0.5 0 0 0 0\n -RT(2 // 3) RT(2 // 3) 0 0 0\n RT(0.5) -1 1 0 0\n -RT(1 // 6) RT(2 // 3) -RT(2 // 3) RT(1 // 6) 0\n ]\nend\n\n\"\"\"\n TVDMISA(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `TVDMISA` method is a 3rd-order accurate, 3-stage MIS method whose\nconstruction is summarized in Table 6 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction TVDMISA(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.1946360605647457) 0 0\n 0 RT(0.3971200136786614) RT(0.2609434606211801) 0\n ]\n\n β = beta(TVDMISA, RT)\n\n γ = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.5624048933209129) 0 0\n 0 RT(0.4408467475713277) RT(-0.2459300561692391) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(TVDMISA), RT::DataType)\n β = [\n 0 0 0 0\n RT(2 // 3) 0 0 0\n RT(-0.28247174703488398) RT(4 // 9) 0 0\n RT(-0.31198081960042401) RT(0.18082737579913699) RT(9 // 16) 0\n ]\nend\n\n\"\"\"\n TVDMISB(slowrhs!, fastrhs!, fastmethod, nsubsteps, Q; dt = 0, t0 = 0)\n\nThe `TVDMISB` method is a 3rd-order accurate, 3-stage MIS method whose\nconstruction is summarized in Table 7 of [KnothWensch2014](@cite).\n\n### References\n - [KnothWensch2014](@cite)\n\"\"\"\nfunction TVDMISB(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n FT = eltype(Q)\n RT = real(FT)\n\n α = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.42668232863311001) 0 0\n 0 RT(0.26570779016173801) RT(0.41489966891866698) 0\n ]\n\n β = beta(TVDMISB, RT)\n\n γ = [\n 0 0 0 0\n 0 0 0 0\n 0 RT(0.28904389120139701) 0 0\n 0 RT(0.45113560071334202) RT(-0.25006656847591002) 0\n ]\n\n MultirateInfinitesimalStep(\n slowrhs!,\n fastrhs!,\n fastmethod,\n nsubsteps,\n α,\n β,\n γ,\n Q;\n dt = dt,\n t0 = t0,\n )\nend\n\nfunction beta(::typeof(TVDMISB), RT::DataType)\n β = [\n 0 0 0 0\n RT(2 // 3) 0 0 0\n RT(-0.25492859100078202) RT(4 // 9) 0 0\n RT(-0.26452517179288798) RT(0.11424084424766399) RT(9 // 16) 0\n ]\nend\n\nfunction getnsubsteps(mis, ns::Int, RT::DataType)\n d = sum(beta(mis, RT), dims = 2)\n return d ./ ns\nend\n" }, { "path": "src/Numerics/ODESolvers/MultirateRungeKuttaMethod.jl", "content": "\nexport MultirateRungeKutta\n\nLSRK2N = LowStorageRungeKutta2N\n\n\"\"\"\n MultirateRungeKutta(slow_solver, fast_solver; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f_fast(Q, t) + f_slow(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThe constructor builds a multirate Runge-Kutta scheme using two different RK\nsolvers. This is based on\n\nCurrently only the low storage RK methods can be used as slow solvers\n\n### References\n - [Schlegel2012](@cite)\n\"\"\"\nmutable struct MultirateRungeKutta{SS, FS, RT} <: AbstractODESolver\n \"slow solver\"\n slow_solver::SS\n \"fast solver\"\n fast_solver::FS\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n\n function MultirateRungeKutta(\n slow_solver::LSRK2N,\n fast_solver,\n Q = nothing;\n dt = getdt(slow_solver),\n t0 = slow_solver.t,\n ) where {AT <: AbstractArray}\n SS = typeof(slow_solver)\n FS = typeof(fast_solver)\n RT = real(eltype(slow_solver.dQ))\n new{SS, FS, RT}(slow_solver, fast_solver, RT(dt), RT(t0), 0)\n end\nend\n\nfunction MultirateRungeKutta(\n solvers::Tuple,\n Q = nothing;\n dt = getdt(solvers[1]),\n t0 = solvers[1].t,\n) where {AT <: AbstractArray}\n if length(solvers) < 2\n error(\"Must specify atleast two solvers\")\n elseif length(solvers) == 2\n fast_solver = solvers[2]\n else\n fast_solver = MultirateRungeKutta(solvers[2:end], Q; dt = dt, t0 = t0)\n end\n\n slow_solver = solvers[1]\n\n MultirateRungeKutta(slow_solver, fast_solver, Q; dt = dt, t0 = t0)\nend\n\nfunction MultirateRungeKutta(\n mrk,\n op::TimeScaledRHS{2, RT} where {RT},\n Q = nothing;\n dt = 0,\n t0 = 0,\n steps = 0,\n) where {AT <: AbstractArray}\n\n slow_solver = mrk(op.rhs![1], Q, dt = dt, t0 = t0)\n fast_solver = mrk(op.rhs![2], Q, dt = dt, t0 = t0)\n\n MultirateRungeKutta(slow_solver, fast_solver, Q; dt = dt, t0 = t0)\nend\n\nfunction dostep!(\n Q,\n mrrk::MultirateRungeKutta{SS},\n param,\n time::Real,\n nsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n) where {SS <: LSRK2N}\n for i in 1:nsteps\n dostep!(Q, mrrk, param, time, slow_δ, slow_rv_dQ, slow_scaling)\n time += mrrk.fast_solver.dt\n end\nend\n\nfunction dostep!(\n Q,\n mrrk::MultirateRungeKutta{SS},\n param,\n time,\n in_slow_δ = nothing,\n in_slow_rv_dQ = nothing,\n in_slow_scaling = nothing,\n) where {SS <: LSRK2N}\n dt = mrrk.dt\n\n slow = mrrk.slow_solver\n fast = mrrk.fast_solver\n\n slow_rv_dQ = realview(slow.dQ)\n\n groupsize = 256\n\n fast_dt_in = getdt(fast)\n\n for slow_s in 1:length(slow.RKA)\n # Currnent slow state time\n slow_stage_time = time + slow.RKC[slow_s] * dt\n\n # Evaluate the slow mode\n slow.rhs!(slow.dQ, Q, param, slow_stage_time, increment = true)\n\n if in_slow_δ !== nothing\n slow_scaling = nothing\n if slow_s == length(slow.RKA)\n slow_scaling = in_slow_scaling\n end\n # update solution and scale RHS\n event = Event(array_device(Q))\n event = update!(array_device(Q), groupsize)(\n slow_rv_dQ,\n in_slow_rv_dQ,\n in_slow_δ,\n slow_scaling;\n ndrange = length(realview(Q)),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n # Fractional time for slow stage\n if slow_s == length(slow.RKA)\n γ = 1 - slow.RKC[slow_s]\n else\n γ = slow.RKC[slow_s + 1] - slow.RKC[slow_s]\n end\n\n # RKB for the slow with fractional time factor remove (since full\n # integration of fast will result in scaling by γ)\n slow_δ = slow.RKB[slow_s] / (γ)\n\n # RKB for the slow with fractional time factor remove (since full\n # integration of fast will result in scaling by γ)\n nsubsteps = fast_dt_in > 0 ? ceil(Int, γ * dt / fast_dt_in) : 1\n fast_dt = γ * dt / nsubsteps\n\n updatedt!(fast, fast_dt)\n\n for substep in 1:nsubsteps\n slow_rka = nothing\n if substep == nsubsteps\n slow_rka = slow.RKA[slow_s % length(slow.RKA) + 1]\n end\n fast_time = slow_stage_time + (substep - 1) * fast_dt\n dostep!(Q, fast, param, fast_time, slow_δ, slow_rv_dQ, slow_rka)\n end\n end\n updatedt!(fast, fast_dt_in)\nend\n\n@kernel function update!(fast_dQ, slow_dQ, δ, slow_rka = nothing)\n i = @index(Global, Linear)\n @inbounds begin\n fast_dQ[i] += δ * slow_dQ[i]\n if slow_rka !== nothing\n slow_dQ[i] *= slow_rka\n end\n end\nend\n" }, { "path": "src/Numerics/ODESolvers/ODESolvers.jl", "content": "\"\"\"\n ODESolvers\n\nOrdinary differential equation solvers\n\"\"\"\nmodule ODESolvers\n\nusing LinearAlgebra\nusing KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\nusing StaticArrays\nusing ..SystemSolvers\nusing ..MPIStateArrays: array_device, realview\nusing ..GenericCallbacks\n\nexport AbstractODESolver, solve!, updatedt!, gettime, getsteps\n\nabstract type AbstractODESolver end\n\n\"\"\"\n gettime(solver::AbstractODESolver)\n\nReturns the current simulation time of the ODE solver `solver`\n\"\"\"\ngettime(solver::AbstractODESolver) = solver.t\n\n\"\"\"\n getdt(solver::AbstractODESolver)\n\nReturns the current simulation time step of the ODE solver `solver`\n\"\"\"\ngetdt(solver::AbstractODESolver) = solver.dt\n\n\"\"\"\n getsteps(solver::AbstractODESolver)\n\nReturns the number of completed time steps of the ODE solver `solver`\n\"\"\"\ngetsteps(solver::AbstractODESolver) = solver.steps\n\n\"\"\"\n ODESolvers.general_dostep!(Q, solver::AbstractODESolver, p,\n timeend::Real, adjustfinalstep::Bool)\n\nUse the solver to step `Q` forward in time from the current time, to the time\n`timeend`. If `adjustfinalstep == true` then `dt` is adjusted so that the step\ndoes not take the solution beyond the `timeend`.\n\"\"\"\nfunction general_dostep!(\n Q,\n solver::AbstractODESolver,\n p,\n timeend::Real;\n adjustfinalstep::Bool,\n)\n time, dt = gettime(solver), getdt(solver)\n final_step = false\n if adjustfinalstep && time + dt > timeend\n orig_dt = dt\n dt = timeend - time\n updatedt!(solver, dt)\n final_step = true\n end\n @assert dt > 0\n\n dostep!(Q, solver, p, time)\n\n if !final_step\n updatetime!(solver, time + dt)\n else\n updatedt!(solver, orig_dt)\n updatetime!(solver, timeend)\n end\nend\n\n\"\"\"\n updatetime!(solver::AbstractODESolver, time)\n\nChange the current time to `time` for the ODE solver `solver`.\n\"\"\"\nupdatetime!(solver::AbstractODESolver, time) = (solver.t = time)\n\n\"\"\"\n updatedt!(solver::AbstractODESolver, dt)\n\nChange the time step size to `dt` for the ODE solver `solver`.\n\"\"\"\nupdatedt!(solver::AbstractODESolver, dt) = (solver.dt = dt)\n\n\"\"\"\n updatesteps!(solver::AbstractODESolver, dt)\n\nSet the number of elapsed time steps for the ODE solver `solver`.\n\"\"\"\nupdatesteps!(solver::AbstractODESolver, steps) = (solver.steps = steps)\n\nisadjustable(solver::AbstractODESolver) = true\n\n# {{{ run!\n\"\"\"\n solve!(Q, solver::AbstractODESolver; timeend,\n stopaftertimeend=true, numberofsteps, callbacks)\n\nSolves an ODE using the `solver` starting from a state `Q`. The state `Q` is\nupdated inplace. The final time `timeend` or `numberofsteps` must be specified.\n\nA series of optional callback functions can be specified using the tuple\n`callbacks`; see the `GenericCallbacks` module.\n\"\"\"\nfunction solve!(\n Q,\n solver::AbstractODESolver,\n param = nothing;\n timeend::Real = Inf,\n adjustfinalstep = true,\n numberofsteps::Integer = 0,\n callbacks = (),\n)\n\n @assert isfinite(timeend) || numberofsteps > 0\n if adjustfinalstep && !isadjustable(solver)\n error(\"$solver does not support time step adjustments. Can only be used with `adjustfinalstep=false`.\")\n end\n t0 = gettime(solver)\n\n # Loop through an initialize callbacks (if they need it)\n GenericCallbacks.init!(callbacks, solver, Q, param, t0)\n\n step = 0\n time = t0\n while time < timeend\n step += 1\n updatesteps!(solver, step)\n\n time = general_dostep!(\n Q,\n solver,\n param,\n timeend;\n adjustfinalstep = adjustfinalstep,\n )\n\n val = GenericCallbacks.call!(callbacks, solver, Q, param, time)\n if val !== nothing && val > 0\n return gettime(solver)\n end\n\n # Figure out if we should stop\n if step == numberofsteps\n break\n end\n end\n\n # Loop through to fini callbacks\n GenericCallbacks.fini!(callbacks, solver, Q, param, time)\n\n return gettime(solver)\nend\n# }}}\n\ninclude(\"BackwardEulerSolvers.jl\")\ninclude(\"MultirateInfinitesimalGARKExplicit.jl\")\ninclude(\"MultirateInfinitesimalGARKDecoupledImplicit.jl\")\ninclude(\"MultirateInfinitesimalStepMethod.jl\")\ninclude(\"LowStorageRungeKuttaMethod.jl\")\ninclude(\"LowStorageRungeKutta3NMethod.jl\")\ninclude(\"StrongStabilityPreservingRungeKuttaMethod.jl\")\ninclude(\"AdditiveRungeKuttaMethod.jl\")\ninclude(\"MultirateRungeKuttaMethod.jl\")\ninclude(\"SplitExplicitMethod.jl\")\ninclude(\"DifferentialEquations.jl\")\n\nend # module\n" }, { "path": "src/Numerics/ODESolvers/SplitExplicitMethod.jl", "content": "export SplitExplicitSolver\n\nusing ..BalanceLaws:\n initialize_states!,\n tendency_from_slow_to_fast!,\n cummulate_fast_solution!,\n reconcile_from_fast_to_slow!\n\nLSRK2N = LowStorageRungeKutta2N\n\n@doc \"\"\"\n SplitExplicitSolver(slow_solver, fast_solver; dt, t0 = 0, coupled = true)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q_{fast}} = f_{fast}(Q_{fast}, Q_{slow}, t)\n \\\\dot{Q_{slow}} = f_{slow}(Q_{slow}, Q_{fast}, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis method performs an operator splitting to timestep the vertical average\nof the model at a faster rate than the full model. This results in a first-\norder time stepper.\n\n\"\"\" SplitExplicitSolver\nmutable struct SplitExplicitSolver{SS, FS, RT, MSA} <: AbstractODESolver\n \"slow solver\"\n slow_solver::SS\n \"fast solver\"\n fast_solver::FS\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"storage for transfer tendency\"\n dQ2fast::MSA\n\n function SplitExplicitSolver(\n slow_solver::LSRK2N,\n fast_solver,\n Q = nothing;\n dt = getdt(slow_solver),\n t0 = slow_solver.t,\n ) where {AT <: AbstractArray}\n SS = typeof(slow_solver)\n FS = typeof(fast_solver)\n RT = real(eltype(slow_solver.dQ))\n\n dQ2fast = similar(slow_solver.dQ)\n dQ2fast .= -0\n MSA = typeof(dQ2fast)\n\n return new{SS, FS, RT, MSA}(\n slow_solver,\n fast_solver,\n RT(dt),\n RT(t0),\n 0,\n dQ2fast,\n )\n end\nend\n\nfunction dostep!(\n Qslow,\n split::SplitExplicitSolver{SS},\n param,\n time::Real,\n) where {SS <: LSRK2N}\n slow = split.slow_solver\n fast = split.fast_solver\n\n Qfast = slow.rhs!.modeldata.Q_2D\n\n dQslow = slow.dQ\n dQ2fast = split.dQ2fast\n\n slow_bl = slow.rhs!.balance_law\n fast_bl = fast.rhs!.balance_law\n\n groupsize = 256\n\n slow_dt = getdt(slow)\n fast_dt_in = getdt(fast)\n\n for slow_s in 1:length(slow.RKA)\n # Current slow state time\n slow_stage_time = time + slow.RKC[slow_s] * slow_dt\n\n # Initialize fast model and tendency adjustment\n # before evalution of slow mode\n initialize_states!(slow_bl, fast_bl, slow.rhs!, fast.rhs!, Qslow, Qfast)\n\n # Evaluate the slow mode\n # --> save tendency for the fast\n slow.rhs!(dQ2fast, Qslow, param, slow_stage_time, increment = false)\n\n # vertically integrate slow tendency to advance fast equation\n # and use vertical mean for slow model (negative source)\n # ---> work with dQ2fast as input\n tendency_from_slow_to_fast!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n dQ2fast,\n )\n\n # Compute (and RK update) slow tendency\n slow.rhs!(dQslow, Qslow, param, slow_stage_time, increment = true)\n\n # Fractional time for slow stage\n if slow_s == length(slow.RKA)\n γ = 1 - slow.RKC[slow_s]\n else\n γ = slow.RKC[slow_s + 1] - slow.RKC[slow_s]\n end\n\n # RKB for the slow with fractional time factor remove (since full\n # integration of fast will result in scaling by γ)\n nsubsteps = fast_dt_in > 0 ? ceil(Int, γ * slow_dt / fast_dt_in) : 1\n fast_dt = γ * slow_dt / nsubsteps\n\n updatedt!(fast, fast_dt)\n\n for substep in 1:nsubsteps\n fast_time = slow_stage_time + (substep - 1) * fast_dt\n dostep!(Qfast, fast, param, fast_time)\n cummulate_fast_solution!(\n slow_bl,\n fast_bl,\n fast.rhs!,\n Qfast,\n fast_time,\n fast_dt,\n substep,\n )\n end\n\n # Update (RK-stage) slow state\n event = Event(array_device(Qslow))\n event = update!(array_device(Qslow), groupsize)(\n realview(dQslow),\n realview(Qslow),\n slow.RKA[slow_s % length(slow.RKA) + 1],\n slow.RKB[slow_s],\n slow_dt,\n nothing,\n nothing,\n nothing;\n ndrange = length(realview(Qslow)),\n dependencies = (event,),\n )\n wait(array_device(Qslow), event)\n\n # reconcile slow equation using fast equation\n reconcile_from_fast_to_slow!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n )\n end\n updatedt!(fast, fast_dt_in)\n\n return nothing\nend\n" }, { "path": "src/Numerics/ODESolvers/StrongStabilityPreservingRungeKuttaMethod.jl", "content": "export StrongStabilityPreservingRungeKutta\nexport SSPRK22Heuns, SSPRK22Ralstons, SSPRK33ShuOsher, SSPRK34SpiteriRuuth\n\n\"\"\"\n StrongStabilityPreservingRungeKutta(f, RKA, RKB, RKC, Q; dt, t0 = 0)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThe constructor builds a strong-stability-preserving Runge--Kutta scheme\nbased on the provided `RKA`, `RKB` and `RKC` coefficient arrays.\n\nThe available concrete implementations are:\n\n - [`SSPRK33ShuOsher`](@ref)\n - [`SSPRK34SpiteriRuuth`](@ref)\n\"\"\"\nmutable struct StrongStabilityPreservingRungeKutta{T, RT, AT, Nstages} <:\n AbstractODESolver\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"rhs function\"\n rhs!::Any\n \"Storage for RHS during the `StrongStabilityPreservingRungeKutta` update\"\n Rstage::AT\n \"Storage for the stage state during the `StrongStabilityPreservingRungeKutta` update\"\n Qstage::AT\n \"RK coefficient vector A (rhs scaling)\"\n RKA::Array{RT, 2}\n \"RK coefficient vector B (rhs add in scaling)\"\n RKB::Array{RT, 1}\n \"RK coefficient vector C (time scaling)\"\n RKC::Array{RT, 1}\n\n function StrongStabilityPreservingRungeKutta(\n rhs!,\n RKA,\n RKB,\n RKC,\n Q::AT;\n dt = 0,\n t0 = 0,\n ) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n new{T, RT, AT, length(RKB)}(\n RT(dt),\n RT(t0),\n 0,\n rhs!,\n similar(Q),\n similar(Q),\n RKA,\n RKB,\n RKC,\n )\n end\nend\n\n\"\"\"\n dostep!(Q, ssp::StrongStabilityPreservingRungeKutta, p, time::Real,\n nsteps::Int, iStage::Int, [slow_δ, slow_rv_dQ, slow_scaling])\n\nWrapper function to use the strong stability preserving Runge--Kutta method `ssp`\nas the fast solver for a Multirate Infinitesimal Step method by calling dostep!(Q,\nssp::StrongStabilityPreservingRungeKutta, p, time::Real, [slow_δ, slow_rv_dQ,\nslow_scaling]) nsubsteps times.\n\"\"\"\nfunction dostep!(\n Q,\n ssp::StrongStabilityPreservingRungeKutta,\n p,\n time::Real,\n nsteps::Int,\n iStage::Int,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n slow_scaling = nothing,\n)\n for i in 1:nsteps\n dostep!(Q, ssp, p, time, slow_δ, slow_rv_dQ, slow_scaling)\n time += ssp.dt\n end\nend\n\n\"\"\"\n ODESolvers.dostep!(Q, ssp::StrongStabilityPreservingRungeKutta, p,\n time::Real, [slow_δ, slow_rv_dQ, slow_scaling])\n\nUse the strong stability preserving Runge--Kutta method `ssp` to step `Q`\nforward in time from the current time `time` to final time `time + getdt(ssp)`.\n\nIf the optional parameter `slow_δ !== nothing` then `slow_rv_dQ * slow_δ` is\nadded as an additional ODE right-hand side source. If the optional parameter\n`slow_scaling !== nothing` then after the final stage update the scaling\n`slow_rv_dQ *= slow_scaling` is performed.\n\"\"\"\nfunction dostep!(\n Q,\n ssp::StrongStabilityPreservingRungeKutta,\n p,\n time,\n slow_δ = nothing,\n slow_rv_dQ = nothing,\n in_slow_scaling = nothing,\n)\n dt = ssp.dt\n\n RKA, RKB, RKC = ssp.RKA, ssp.RKB, ssp.RKC\n rhs! = ssp.rhs!\n Rstage, Qstage = ssp.Rstage, ssp.Qstage\n\n rv_Q = realview(Q)\n rv_Rstage = realview(Rstage)\n rv_Qstage = realview(Qstage)\n groupsize = 256\n\n rv_Qstage .= rv_Q\n for s in 1:length(RKB)\n rhs!(Rstage, Qstage, p, time + RKC[s] * dt, increment = false)\n\n slow_scaling = nothing\n if s == length(RKB)\n slow_scaling = in_slow_scaling\n end\n event = Event(array_device(Q))\n event = update!(array_device(Q), groupsize)(\n rv_Rstage,\n rv_Q,\n rv_Qstage,\n RKA[s, 1],\n RKA[s, 2],\n RKB[s],\n dt,\n slow_δ,\n slow_rv_dQ,\n slow_scaling;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n rv_Q .= rv_Qstage\nend\n\n@kernel function update!(\n dQ,\n Q,\n Qstage,\n rka1,\n rka2,\n rkb,\n dt,\n slow_δ,\n slow_dQ,\n slow_scaling,\n)\n i = @index(Global, Linear)\n @inbounds begin\n if slow_δ !== nothing\n dQ[i] += slow_δ * slow_dQ[i]\n end\n Qstage[i] = rka1 * Q[i] + rka2 * Qstage[i] + dt * rkb * dQ[i]\n if slow_scaling !== nothing\n slow_dQ[i] *= slow_scaling\n end\n end\nend\n\n\"\"\"\n SSPRK22Heuns(f, Q; dt, t0 = 0)\n\nThis function returns a [`StrongStabilityPreservingRungeKutta`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the second-order, 2-stage, strong-stability-preserving, Runge--Kutta scheme\nof Shu and Osher (1988) (also known as Heun's method.)\nExact choice of coefficients from wikipedia page for Heun's method :)\n\n### References\n - [Shu1988](@cite)\n - [Heun1900](@cite)\n\"\"\"\nfunction SSPRK22Heuns(F, Q::AT; dt = 0, t0 = 0) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n RKA = [RT(1) RT(0); RT(1 // 2) RT(1 // 2)]\n RKB = [RT(1), RT(1 // 2)]\n RKC = [RT(0), RT(1)]\n StrongStabilityPreservingRungeKutta(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n SSPRK22Ralstons(f, Q; dt, t0 = 0)\n\nThis function returns a [`StrongStabilityPreservingRungeKutta`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the second-order, 2-stage, strong-stability-preserving, Runge--Kutta scheme\nof Shu and Osher (1988) (also known as Ralstons's method.)\nExact choice of coefficients from wikipedia page for Heun's method :)\n\n### References\n - [Shu1988](@cite)\n - [Ralston1962](@cite)\n\"\"\"\nfunction SSPRK22Ralstons(F, Q::AT; dt = 0, t0 = 0) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n RKA = [RT(1) RT(0); RT(5 // 8) RT(3 // 8)]\n RKB = [RT(2 // 3), RT(3 // 4)]\n RKC = [RT(0), RT(2 // 3)]\n StrongStabilityPreservingRungeKutta(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n SSPRK33ShuOsher(f, Q; dt, t0 = 0)\n\nThis function returns a [`StrongStabilityPreservingRungeKutta`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the third-order, 3-stage, strong-stability-preserving, Runge--Kutta scheme\nof Shu and Osher (1988)\n\n### References\n - [Shu1988](@cite)\n\"\"\"\nfunction SSPRK33ShuOsher(F, Q::AT; dt = 0, t0 = 0) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n RKA = [RT(1) RT(0); RT(3 // 4) RT(1 // 4); RT(1 // 3) RT(2 // 3)]\n RKB = [RT(1), RT(1 // 4), RT(2 // 3)]\n RKC = [RT(0), RT(1), RT(1 // 2)]\n StrongStabilityPreservingRungeKutta(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n\n\"\"\"\n SSPRK34SpiteriRuuth(f, Q; dt, t0 = 0)\n\nThis function returns a [`StrongStabilityPreservingRungeKutta`](@ref) time stepping object\nfor explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q} = f(Q, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis uses the third-order, 4-stage, strong-stability-preserving, Runge--Kutta scheme\nof Spiteri and Ruuth (1988)\n\n### References\n - [Spiteri2002](@cite)\n\"\"\"\nfunction SSPRK34SpiteriRuuth(\n F,\n Q::AT;\n dt = 0,\n t0 = 0,\n) where {AT <: AbstractArray}\n T = eltype(Q)\n RT = real(T)\n RKA = [RT(1) RT(0); RT(0) RT(1); RT(2 // 3) RT(1 // 3); RT(0) RT(1)]\n RKB = [RT(1 // 2); RT(1 // 2); RT(1 // 6); RT(1 // 2)]\n RKC = [RT(0); RT(1 // 2); RT(1); RT(1 // 2)]\n StrongStabilityPreservingRungeKutta(F, RKA, RKB, RKC, Q; dt = dt, t0 = t0)\nend\n" }, { "path": "src/Numerics/SystemSolvers/SystemSolvers.jl", "content": "module SystemSolvers\n\nusing ..MPIStateArrays\nusing ..MPIStateArrays: array_device, realview\n\nusing ..Mesh.Grids\nimport ..Mesh.Grids: polynomialorders, dimensionality\nusing ..Mesh.Topologies\nusing ..DGMethods\nusing ..DGMethods: DGModel, DGFVModel, SpaceDiscretization\nimport ..DGMethods.FVReconstructions: width\nusing ..BalanceLaws\n\nusing Adapt\nusing CUDA\nusing LinearAlgebra\nusing LazyArrays\nusing StaticArrays\nusing KernelAbstractions\n\nconst weighted_norm = false\n\n# just for testing SystemSolvers\nLinearAlgebra.norm(A::MVector, p::Real, weighted::Bool) = norm(A, p)\nLinearAlgebra.norm(A::MVector, weighted::Bool) = norm(A, 2, weighted)\nLinearAlgebra.dot(A::MVector, B::MVector, weighted) = dot(A, B)\nLinearAlgebra.norm(A::AbstractVector, p::Real, weighted::Bool) = norm(A, p)\nLinearAlgebra.norm(A::AbstractVector, weighted::Bool) = norm(A, 2, weighted)\nLinearAlgebra.dot(A::AbstractVector, B::AbstractVector, weighted) = dot(A, B)\n\nexport linearsolve!,\n settolerance!, prefactorize, construct_preconditioner, preconditioner_solve!\nexport AbstractSystemSolver,\n AbstractIterativeSystemSolver, AbstractNonlinearSolver\nexport nonlinearsolve!\n\n\"\"\"\n AbstractSystemSolver\n\nThis is an abstract type representing a generic linear solver.\n\"\"\"\nabstract type AbstractSystemSolver end\n\n\"\"\"\n AbstractNonlinearSolver\n\nThis is an abstract type representing a generic nonlinear solver.\n\"\"\"\nabstract type AbstractNonlinearSolver <: AbstractSystemSolver end\n\n\"\"\"\n LSOnly\n\nOnly applies the linear solver (no Newton solver)\n\"\"\"\nstruct LSOnly <: AbstractNonlinearSolver\n linearsolver::Any\nend\n\nfunction donewtoniteration!(\n rhs!,\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::LSOnly,\n args...,\n)\n @info \"donewtoniteration! linearsolve!\", args...\n linearsolve!(\n linearoperator!,\n preconditioner,\n solver.linearsolver,\n Q,\n Qrhs,\n args...;\n max_iters = getmaxiterations(solver.linearsolver),\n )\nend\n\n\n\"\"\"\n\nSolving rhs!(Q) = Qrhs via Newton,\n\nwhere `F = rhs!(Q) - Qrhs`\n\ndF/dQ(Q^n) ΔQ ≈ jvp!(ΔQ; Q^n, F(Q^n))\n\npreconditioner ≈ dF/dQ(Q)\n\n\"\"\"\nfunction nonlinearsolve!(\n rhs!,\n jvp!,\n preconditioner,\n solver::AbstractNonlinearSolver,\n Q::AT,\n Qrhs,\n args...;\n max_newton_iters = 10,\n cvg = Ref{Bool}(),\n) where {AT}\n\n FT = eltype(Q)\n tol = solver.tol\n converged = false\n iters = 0\n\n if preconditioner === nothing\n preconditioner = NoPreconditioner()\n end\n\n # Initialize NLSolver, compute initial residual\n initial_residual_norm = initialize!(rhs!, Q, Qrhs, solver, args...)\n if initial_residual_norm < tol\n converged = true\n end\n converged && return iters\n\n\n while !converged && iters < max_newton_iters\n\n # dF/dQ(Q^n) ΔQ ≈ jvp!(ΔQ; Q^n, F(Q^n)), update Q^n in jvp!\n update_Q!(jvp!, Q, args...)\n\n # update preconditioner based on finite difference, with jvp!\n preconditioner_update!(jvp!, rhs!.f!, preconditioner, args...)\n\n # do newton iteration with Q^{n+1} = Q^{n} - dF/dQ(Q^n)⁻¹ (rhs!(Q) - Qrhs)\n residual_norm, linear_iterations = donewtoniteration!(\n rhs!,\n jvp!,\n preconditioner,\n Q,\n Qrhs,\n solver,\n args...,\n )\n # @info \"Linear solver converged in $linear_iterations iterations\"\n iters += 1\n\n preconditioner_counter_update!(preconditioner)\n\n\n if !isfinite(residual_norm)\n error(\"norm of residual is not finite after $iters iterations of `donewtoniteration!`\")\n end\n\n # Check residual_norm / norm(R0)\n # Comment: Should we check \"correction\" magitude?\n # ||Delta Q|| / ||Q|| ?\n relresidual = residual_norm / initial_residual_norm\n if relresidual < tol || residual_norm < tol\n # @info \"Newton converged in $iters iterations!\"\n converged = true\n end\n end\n\n converged ||\n @warn \"Nonlinear solver did not converge after $iters iterations\"\n cvg[] = converged\n\n iters\nend\n\n\"\"\"\n AbstractIterativeSystemSolver\n\nThis is an abstract type representing a generic iterative\nlinear solver.\n\nThe available concrete implementations are:\n\n - [`GeneralizedConjugateResidual`](@ref)\n - [`GeneralizedMinimalResidual`](@ref)\n\"\"\"\nabstract type AbstractIterativeSystemSolver <: AbstractSystemSolver end\n\n\"\"\"\n settolerance!(solver::AbstractIterativeSystemSolver, tolerance, relative)\n\nSets the relative or absolute tolerance of the iterative linear solver\n`solver` to `tolerance`.\n\"\"\"\nsettolerance!(\n solver::AbstractIterativeSystemSolver,\n tolerance,\n relative = true,\n) = (relative ? (solver.rtol = tolerance) : (solver.atol = tolerance))\n\ndoiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::AbstractIterativeSystemSolver,\n threshold,\n args...,\n) = throw(MethodError(\n doiteration!,\n (linearoperator!, preconditioner, Q, Qrhs, solver, tolerance, args...),\n))\n\ninitialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::AbstractIterativeSystemSolver,\n args...,\n) = throw(MethodError(initialize!, (linearoperator!, Q, Qrhs, solver, args...)))\n\n\"\"\"\n prefactorize(linop!, linearsolver, args...)\n\nPrefactorize the in-place linear operator `linop!` for use with `linearsolver`.\n\"\"\"\nfunction prefactorize(\n linop!,\n linearsolver::AbstractIterativeSystemSolver,\n args...,\n)\n return nothing\nend\n\n\"\"\"\n linearsolve!(linearoperator!, solver::AbstractIterativeSystemSolver, Q, Qrhs, args...)\n\nSolves a linear problem defined by the `linearoperator!` function and the state\n`Qrhs`, i.e,\n\n```math\nL(Q) = Q_{rhs}\n```\n\nusing the `solver` and the initial guess `Q`. After the call `Q` contains the\nsolution. The arguments `args` is passed to `linearoperator!` when it is\ncalled.\n\"\"\"\nfunction linearsolve!(\n linearoperator!,\n preconditioner,\n solver::AbstractIterativeSystemSolver,\n Q,\n Qrhs,\n args...;\n max_iters = length(Q),\n cvg = Ref{Bool}(),\n)\n converged = false\n iters = 0\n\n if preconditioner === nothing\n preconditioner = NoPreconditioner()\n end\n\n converged, threshold =\n initialize!(linearoperator!, Q, Qrhs, solver, args...)\n converged && return iters\n\n while !converged && iters < max_iters\n converged, inner_iters, residual_norm = doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver,\n threshold,\n args...,\n )\n\n iters += inner_iters\n\n if !isfinite(residual_norm)\n error(\"norm of residual is not finite after $iters iterations of `doiteration!`\")\n end\n\n achieved_tolerance = residual_norm / threshold * solver.rtol\n end\n\n converged ||\n @warn \"Solver did not attain convergence after $iters iterations\"\n cvg[] = converged\n\n iters\nend\n\n@kernel function linearcombination!(Q, cs, Xs, increment::Bool)\n i = @index(Global, Linear)\n if !increment\n @inbounds Q[i] = -zero(eltype(Q))\n end\n @inbounds for j in 1:length(cs)\n Q[i] += cs[j] * Xs[j][i]\n end\nend\n\ninclude(\"generalized_minimal_residual_solver.jl\")\ninclude(\"generalized_conjugate_residual_solver.jl\")\ninclude(\"conjugate_gradient_solver.jl\")\ninclude(\"columnwise_lu_solver.jl\")\ninclude(\"preconditioners.jl\")\ninclude(\"batched_generalized_minimal_residual_solver.jl\")\ninclude(\"jacobian_free_newton_krylov_solver.jl\")\n\nend\n" }, { "path": "src/Numerics/SystemSolvers/batched_generalized_minimal_residual_solver.jl", "content": "\nusing CUDA\n\nexport BatchedGeneralizedMinimalResidual\n\n\"\"\"\n BatchedGeneralizedMinimalResidual(\n Q,\n dofperbatch,\n Nbatch;\n M = min(20, length(Q)),\n rtol = √eps(eltype(AT)),\n atol = eps(eltype(AT)),\n forward_reshape = size(Q),\n forward_permute = Tuple(1:length(size(Q))),\n )\n\n# BGMRES\nThis is an object for solving batched linear systems using the GMRES algorithm.\nThe constructor parameter `M` is the number of steps after which the algorithm\nis restarted (if it has not converged), `Q` is a reference state used only\nto allocate the solver internal state, `dofperbatch` is the size of each batched\nsystem (assumed to be the same throughout), `Nbatch` is the total number\nof independent linear systems, and `rtol` specifies the convergence\ncriterion based on the relative residual norm (max across all batched systems).\nThe argument `forward_reshape` is a tuple of integers denoting the reshaping\n(if required) of the solution vectors for batching the Arnoldi routines.\nThe argument `forward_permute` describes precisely which indices of the\narray `Q` to permute. This object is intended to be passed to\nthe [`linearsolve!`](@ref) command.\n\nThis uses a batched-version of the restarted Generalized Minimal Residual method\nof Saad and Schultz (1986).\n\n# Note\nEventually, we'll want to do something like this:\n\n i = @index(Global)\n linearoperator!(Q[:, :, :, i], args...)\n\nThis will help stop the need for constantly\nreshaping the work arrays. It would also potentially\nsave us some memory.\n\"\"\"\nmutable struct BatchedGeneralizedMinimalResidual{\n I,\n T,\n AT,\n BKT,\n OmT,\n HT,\n gT,\n sT,\n resT,\n res0T,\n FRS,\n FPR,\n BRS,\n BPR,\n} <: AbstractIterativeSystemSolver\n\n \"global Krylov basis at present step\"\n krylov_basis::AT\n \"global Krylov basis at previous step\"\n krylov_basis_prev::AT\n \"global batched Krylov basis\"\n batched_krylov_basis::BKT\n \"Storage for the Givens rotation matrices\"\n Ω::OmT\n \"Hessenberg matrix in each column\"\n H::HT\n \"rhs of the least squares problem in each column\"\n g0::gT\n \"The GMRES iterate in each batched column\"\n sol::sT\n \"Relative tolerance\"\n rtol::T\n \"Absolute tolerance\"\n atol::T\n \"Maximum number of GMRES iterations (global across all columns)\"\n max_iter::I\n \"total number of batched columns\"\n batch_size::I\n \"total number of dofs per batched column\"\n dofperbatch::I\n \"residual norm in each column\"\n resnorms::resT\n \"initial residual norm in each column\"\n initial_resnorms::res0T\n forward_reshape::FRS\n forward_permute::FPR\n backward_reshape::BRS\n backward_permute::BPR\n\n function BatchedGeneralizedMinimalResidual(\n Q::AT,\n dofperbatch,\n Nbatch;\n M = min(20, length(Q)),\n rtol = √eps(eltype(AT)),\n atol = eps(eltype(AT)),\n forward_reshape = size(Q),\n forward_permute = Tuple(1:length(size(Q))),\n ) where {AT}\n # Get ArrayType information\n if isa(array_device(Q), CPU)\n ArrayType = Array\n else\n # Sanity check since we don't support anything else\n @assert isa(array_device(Q), CUDADevice)\n ArrayType = CuArray\n end\n\n # FIXME: If we can batch the application of linearoperator!, then we dont\n # need these two temporary work vectors (unpermuted/reshaped)\n krylov_basis = similar(Q)\n krylov_basis_prev = similar(Q)\n\n FT = eltype(AT)\n # Create storage for holding the batched Krylov basis\n batched_krylov_basis =\n fill!(ArrayType{FT}(undef, M + 1, dofperbatch, Nbatch), 0)\n\n # Create storage for doing the batched Arnoldi process\n Ω = fill!(ArrayType{FT}(undef, Nbatch, 2 * M), 0)\n H = fill!(ArrayType{FT}(undef, Nbatch, M + 1, M), 0)\n g0 = fill!(ArrayType{FT}(undef, Nbatch, M + 1), 0)\n\n # Create storage for constructing the global gmres iterate\n # and recording column-norms\n sol = fill!(ArrayType{FT}(undef, dofperbatch, Nbatch), 0)\n resnorms = fill!(ArrayType{FT}(undef, Nbatch), 0)\n initial_resnorms = fill!(ArrayType{FT}(undef, Nbatch), 0)\n\n @assert dofperbatch * Nbatch == length(Q)\n\n # Define the back permutation and reshape\n backward_permute = Tuple(sortperm([forward_permute...]))\n tmp_reshape_tuple_b = [forward_reshape...]\n permute!(tmp_reshape_tuple_b, [forward_permute...])\n backward_reshape = Tuple(tmp_reshape_tuple_b)\n\n # FIXME: Is there a better way of doing this?\n BKT = typeof(batched_krylov_basis)\n OmT = typeof(Ω)\n HT = typeof(H)\n gT = typeof(g0)\n sT = typeof(sol)\n resT = typeof(resnorms)\n res0T = typeof(initial_resnorms)\n FRS = typeof(forward_reshape)\n FPR = typeof(forward_permute)\n BRS = typeof(backward_reshape)\n BPR = typeof(backward_permute)\n\n return new{\n typeof(Nbatch),\n eltype(Q),\n AT,\n BKT,\n OmT,\n HT,\n gT,\n sT,\n resT,\n res0T,\n FRS,\n FPR,\n BRS,\n BPR,\n }(\n krylov_basis,\n krylov_basis_prev,\n batched_krylov_basis,\n Ω,\n H,\n g0,\n sol,\n rtol,\n atol,\n M,\n Nbatch,\n dofperbatch,\n resnorms,\n initial_resnorms,\n forward_reshape,\n forward_permute,\n backward_reshape,\n backward_permute,\n )\n end\nend\n\n\"\"\"\n BatchedGeneralizedMinimalResidual(\n dg::SpaceDiscretization,\n Q::MPIStateArray;\n atol = sqrt(eps(eltype(Q))),\n rtol = sqrt(eps(eltype(Q))),\n max_subspace_size = nothing,\n independent_states = false,\n )\n\n# Description\nSpecialized constructor for `BatchedGeneralizedMinimalResidual` struct, using\na `SpaceDiscretization` to infer state-information and determine appropriate reshaping\nand permutations.\n\n# Arguments\n- `dg`: (SpaceDiscretization) A `SpaceDiscretization` containing all relevant grid and topology\n information.\n- `Q` : (MPIStateArray) An `MPIStateArray` containing field information.\n\n# Keyword Arguments\n- `atol`: (float) absolute tolerance. `DEFAULT = sqrt(eps(eltype(Q)))`\n- `rtol`: (float) relative tolerance. `DEFAULT = sqrt(eps(eltype(Q)))`\n- `max_subspace_size` : (Int). Maximal dimension of each (batched)\n Krylov subspace. DEFAULT = nothing\n- `independent_states`: (boolean) An optional flag indicating whether\n or not degrees of freedom are coupled\n internally (within a column).\n `DEFAULT = false`\n# Return\ninstance of `BatchedGeneralizedMinimalResidual` struct\n\"\"\"\nfunction BatchedGeneralizedMinimalResidual(\n dg::SpaceDiscretization,\n Q::MPIStateArray;\n atol = sqrt(eps(eltype(Q))),\n rtol = sqrt(eps(eltype(Q))),\n max_subspace_size = nothing,\n independent_states = false,\n)\n grid = dg.grid\n topology = grid.topology\n dim = dimensionality(grid)\n N = polynomialorders(grid)\n\n # Number of Gauss-Lobatto quadrature points in each direction\n Nq = N .+ 1\n\n # Number of states and elements (in vertical and horizontal directions)\n num_states = size(Q)[2]\n nelem = length(topology.realelems)\n nvertelem = topology.stacksize\n nhorzelem = div(nelem, nvertelem)\n\n # Definition of a \"column\" here is a vertical stack of degrees\n # of freedom. For example, consider a mesh consisting of a single\n # linear element:\n # o----------o\n # |\\ d1 d2 |\\\n # | \\ | \\\n # | \\ d3 d4 \\\n # | o----------o\n # o--d5---d6-o |\n # \\ | \\ |\n # \\ | \\ |\n # \\|d7 d8 \\|\n # o----------o\n # There are 4 total 1-D columns, each containing two\n # degrees of freedom. In general, a mesh of stacked elements will\n # have `Nq[1] * Nq[2] * nhorzelem` total 1-D columns.\n # A single 1-D column has `Nq[3] * nvertelem * num_states`\n # degrees of freedom.\n #\n # nql = length(Nq)\n # indices: (1...nql, nql + 1 , nql + 2, nql + 3)\n\n # for 3d case, this is [ni, nj, nk, num_states, nvertelem, nhorzelem]\n # here ni, nj, nk are number of Gauss quadrature points in each element in x-y-z directions\n # Q = reshape(Q, reshaping_tup), leads to the column-wise fashion Q\n reshaping_tup = (Nq..., num_states, nvertelem, nhorzelem)\n\n @inbounds Nqv = Nq[dim]\n @inbounds Nqh = dim == 2 ? Nq[1] : Nq[1] * Nq[2]\n if independent_states\n m = Nqv * nvertelem\n n = Nqh * nhorzelem * num_states\n else\n m = Nqv * nvertelem * num_states\n n = Nqh * nhorzelem\n end\n\n if max_subspace_size === nothing\n max_subspace_size = m\n end\n\n # permute [ni, nj, nk, num_states, nvertelem, nhorzelem]\n # to [nvertelem, nk, num_states, ni, nj, nhorzelem]\n permute_size = length(reshaping_tup)\n permute_tuple_f = (dim + 1, dim, dim + 2, (1:(dim - 1))..., permute_size)\n\n return BatchedGeneralizedMinimalResidual(\n Q,\n m,\n n;\n M = max_subspace_size,\n atol = atol,\n rtol = rtol,\n forward_reshape = reshaping_tup,\n forward_permute = permute_tuple_f,\n )\nend\n\nfunction initialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::BatchedGeneralizedMinimalResidual,\n args...;\n restart = false,\n)\n g0 = solver.g0\n krylov_basis = solver.krylov_basis\n rtol, atol = solver.rtol, solver.atol\n\n batched_krylov_basis = solver.batched_krylov_basis\n Ndof = solver.dofperbatch\n forward_reshape = solver.forward_reshape\n forward_permute = solver.forward_permute\n resnorms = solver.resnorms\n initial_resnorms = solver.initial_resnorms\n max_iter = solver.max_iter\n\n # Device and groupsize information\n device = array_device(Q)\n groupsize = 256\n\n @assert size(Q) == size(krylov_basis)\n\n # PRECONDITIONER: PQ0 -> P*Q0,\n # the first basis is (J Pinv)PQ0 = b, kry1 = b - J Q0\n linearoperator!(krylov_basis, Q, args...)\n krylov_basis .= Qrhs .- krylov_basis\n\n # Convert into a batched Krylov basis vector\n # REMARK: Ugly hack on the GPU. Can we fix this?\n tmp_array = similar(batched_krylov_basis, size(batched_krylov_basis)[2:3])\n convert_structure!(\n tmp_array,\n krylov_basis,\n forward_reshape,\n forward_permute,\n )\n batched_krylov_basis[1, :, :] .= tmp_array\n\n # Now we initialize across all columns (solver.batch_size).\n # This function also computes the residual norm in each column\n event = Event(device)\n event = batched_initialize!(device, groupsize)(\n resnorms,\n g0,\n batched_krylov_basis,\n Ndof,\n max_iter;\n ndrange = solver.batch_size,\n dependencies = (event,),\n )\n wait(device, event)\n\n # When restarting, we do not want to overwrite the initial threshold,\n # otherwise we may not get an accurate indication that we have sufficiently\n # reduced the GMRES residual.\n if !restart\n initial_resnorms .= resnorms\n end\n residual_norm = maximum(resnorms)\n initial_residual_norm = maximum(initial_resnorms)\n converged =\n check_convergence(residual_norm, initial_residual_norm, atol, rtol)\n\n converged, residual_norm\nend\n\nfunction doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::BatchedGeneralizedMinimalResidual,\n threshold,\n args...,\n)\n FT = eltype(Q)\n krylov_basis = solver.krylov_basis\n krylov_basis_prev = solver.krylov_basis_prev\n Hs = solver.H\n g0s = solver.g0\n Ωs = solver.Ω\n sols = solver.sol\n batched_krylov_basis = solver.batched_krylov_basis\n Ndof = solver.dofperbatch\n rtol, atol = solver.rtol, solver.atol\n max_iter = solver.max_iter\n\n forward_reshape = solver.forward_reshape\n forward_permute = solver.forward_permute\n backward_reshape = solver.backward_reshape\n backward_permute = solver.backward_permute\n resnorms = solver.resnorms\n initial_resnorms = solver.initial_resnorms\n initial_residual_norm = maximum(initial_resnorms)\n\n # Device and groupsize information\n device = array_device(Q)\n groupsize = 256\n\n # Main batched-GMRES iteration cycle\n converged = false\n residual_norm = typemax(FT)\n j = 1\n for outer j in 1:max_iter\n # FIXME: Remove this back-and-forth reshaping by exploiting the\n # data layout in a similar way that the ColumnwiseLU solver does\n\n convert_structure!(\n krylov_basis_prev,\n view(batched_krylov_basis, j, :, :),\n backward_reshape,\n backward_permute,\n )\n\n # PRECONDITIONER: batched_krylov_basis[j+1] = J P^{-1}batched_krylov_basis[j]\n # set krylov_basis_prev = P^{-1}batched_krylov_basis[j]\n preconditioner_solve!(preconditioner, krylov_basis_prev)\n\n # Global operator application to get new Krylov basis vector\n linearoperator!(krylov_basis, krylov_basis_prev, args...)\n\n # Now that we have a global Krylov vector, we reshape and batch\n # the Arnoldi iterations across all columns\n convert_structure!(\n view(batched_krylov_basis, j + 1, :, :),\n krylov_basis,\n forward_reshape,\n forward_permute,\n )\n\n event = Event(device)\n event = batched_arnoldi_process!(device, groupsize)(\n resnorms,\n g0s,\n Hs,\n Ωs,\n batched_krylov_basis,\n j,\n Ndof;\n ndrange = solver.batch_size,\n dependencies = (event,),\n )\n wait(device, event)\n\n # Current stopping criteria is based on the maximal column norm\n # TODO: Once we are able to batch the operator application, we\n # should revisit the termination criteria.\n residual_norm = maximum(resnorms)\n converged =\n check_convergence(residual_norm, initial_residual_norm, atol, rtol)\n if converged\n break\n end\n end\n\n # Reshape the solution vector to construct the new GMRES iterate\n # PRECONDITIONER Q = Q0 + Pinv PΔQ = Q0 + Pinv (Kry * y)\n # sol = PΔQ = Kry * y\n sols .= 0\n\n # Solve the triangular system (minimization problem for optimal linear coefficients\n # in the GMRES iterate) and construct the current iterate in each column\n event = Event(device)\n event = construct_batched_gmres_iterate!(device, groupsize)(\n batched_krylov_basis,\n Hs,\n g0s,\n sols,\n j,\n Ndof;\n ndrange = solver.batch_size,\n dependencies = (event,),\n )\n wait(device, event)\n\n # Use krylov_basis_prev as container for ΔQ\n ΔQ = krylov_basis_prev\n # Unwind reshaping and return solution in standard format\n convert_structure!(ΔQ, sols, backward_reshape, backward_permute)\n # PRECONDITIONER: Q -> Pinv Q\n preconditioner_solve!(preconditioner, ΔQ)\n Q .+= ΔQ\n\n # if not converged, then restart\n converged ||\n initialize!(linearoperator!, Q, Qrhs, solver, args...; restart = true)\n\n (converged, j, residual_norm)\nend\n\n@kernel function batched_initialize!(\n resnorms,\n g0,\n batched_krylov_basis,\n Ndof,\n M,\n)\n cidx = @index(Global)\n FT = eltype(batched_krylov_basis)\n\n # Initialize entire RHS storage\n @inbounds for j in 1:(M + 1)\n g0[cidx, j] = FT(0.0)\n end\n\n # Now we compute the first element of g0[cidx, :],\n # which is determined by the column norm of the initial residual ∥r0∥_2:\n # g0 = ∥r0∥_2 e1\n @inbounds for j in 1:Ndof\n g0[cidx, 1] +=\n batched_krylov_basis[1, j, cidx] * batched_krylov_basis[1, j, cidx]\n end\n @inbounds g0[cidx, 1] = sqrt(g0[cidx, 1])\n\n # Normalize the batched_krylov_basis by the (local) residual norm\n @inbounds for j in 1:Ndof\n batched_krylov_basis[1, j, cidx] /= g0[cidx, 1]\n end\n\n # Record initialize residual norm in the column\n @inbounds resnorms[cidx] = g0[cidx, 1]\n\n nothing\nend\n\n@kernel function batched_arnoldi_process!(\n resnorms,\n g0,\n H,\n Ω,\n batched_krylov_basis,\n j,\n Ndof,\n)\n cidx = @index(Global)\n FT = eltype(batched_krylov_basis)\n\n # Arnoldi process in the local column `cidx`\n @inbounds for i in 1:j\n H[cidx, i, j] = FT(0.0)\n # Modified Gram-Schmidt procedure to generate the Hessenberg matrix\n for k in 1:Ndof\n H[cidx, i, j] +=\n batched_krylov_basis[j + 1, k, cidx] *\n batched_krylov_basis[i, k, cidx]\n end\n # Orthogonalize new Krylov vector against previous one\n for k in 1:Ndof\n batched_krylov_basis[j + 1, k, cidx] -=\n H[cidx, i, j] * batched_krylov_basis[i, k, cidx]\n end\n end\n\n # And finally, normalize latest Krylov basis vector\n local_norm = FT(0.0)\n @inbounds for i in 1:Ndof\n local_norm +=\n batched_krylov_basis[j + 1, i, cidx] *\n batched_krylov_basis[j + 1, i, cidx]\n end\n @inbounds H[cidx, j + 1, j] = sqrt(local_norm)\n @inbounds for i in 1:Ndof\n batched_krylov_basis[j + 1, i, cidx] /= H[cidx, j + 1, j]\n end\n\n # Loop over previously computed Krylov basis vectors\n # and apply the Givens rotations\n @inbounds for i in 1:(j - 1)\n cos_tmp = Ω[cidx, 2 * i - 1]\n sin_tmp = Ω[cidx, 2 * i]\n\n # Apply the Givens rotations\n # | cos -sin | | hi |\n # | sin cos | | hi+1 |\n tmp1 = cos_tmp * H[cidx, i, j] + sin_tmp * H[cidx, i + 1, j]\n H[cidx, i + 1, j] =\n -sin_tmp * H[cidx, i, j] + cos_tmp * H[cidx, i + 1, j]\n H[cidx, i, j] = tmp1\n end\n\n # Eliminate the last element hj+1 and update the rotation matrix\n # | cos -sin | | hj | = | hj'|\n # | sin cos | | hj+1 | = | 0 |\n # where cos, sin = hj+1/sqrt(hj^2 + hj+1^2), hj/sqrt(hj^2 + hj+1^2),\n # and update for next iteration\n @inbounds begin\n Ω[cidx, 2 * j - 1] = H[cidx, j, j]\n Ω[cidx, 2 * j] = H[cidx, j + 1, j]\n H[cidx, j, j] = sqrt(Ω[cidx, 2 * j - 1]^2 + Ω[cidx, 2 * j]^2)\n H[cidx, j + 1, j] = FT(0.0)\n Ω[cidx, 2 * j - 1] /= H[cidx, j, j]\n Ω[cidx, 2 * j] /= H[cidx, j, j]\n\n # And now to the rhs g0\n cos_tmp = Ω[cidx, 2 * j - 1]\n sin_tmp = Ω[cidx, 2 * j]\n tmp1 = cos_tmp * g0[cidx, j] + sin_tmp * g0[cidx, j + 1]\n g0[cidx, j + 1] = -sin_tmp * g0[cidx, j] + cos_tmp * g0[cidx, j + 1]\n g0[cidx, j] = tmp1\n\n # Record estimate for the gmres residual\n resnorms[cidx] = abs(g0[cidx, j + 1])\n end\n\n nothing\nend\n\n@kernel function construct_batched_gmres_iterate!(\n batched_krylov_basis,\n Hs,\n g0s,\n sols,\n j,\n Ndof,\n)\n # Solve for the GMRES coefficients (yⱼ) at the `j`-th\n # iteration that minimizes ∥ b - A xⱼ ∥_2, where\n # xⱼ = ∑ᵢ yᵢ Ψᵢ, with Ψᵢ denoting the Krylov basis vectors\n cidx = @index(Global)\n\n # Do the upper-triangular backsolve\n @inbounds for i in j:-1:1\n g0s[cidx, i] /= Hs[cidx, i, i]\n for k in 1:(i - 1)\n g0s[cidx, k] -= Hs[cidx, k, i] * g0s[cidx, i]\n end\n end\n\n # Having determined yᵢ, we now construct the GMRES solution\n # in each column: xⱼ = ∑ᵢ yᵢ Ψᵢ\n @inbounds for i in 1:j\n for k in 1:Ndof\n sols[k, cidx] += g0s[cidx, i] * batched_krylov_basis[i, k, cidx]\n end\n end\n\n nothing\nend\n\n\"\"\"\n convert_structure!(\n x,\n y,\n reshape_tuple,\n permute_tuple,\n )\n\nComputes a tensor transpose and stores result in x\n\n# Arguments\n- `x`: (array) [OVERWRITTEN]. target destination for storing the y data\n- `y`: (array). data that we want to copy\n- `reshape_tuple`: (tuple) reshapes y to be like that of x, up to a permutation\n- `permute_tuple`: (tuple) permutes the reshaped array into the correct structure\n\"\"\"\n@inline function convert_structure!(x, y, reshape_tuple, permute_tuple)\n alias_y = reshape(y, reshape_tuple)\n permute_y = permutedims(alias_y, permute_tuple)\n copyto!(x, reshape(permute_y, size(x)))\n nothing\nend\n@inline convert_structure!(x, y::MPIStateArray, reshape_tuple, permute_tuple) =\n convert_structure!(x, y.realdata, reshape_tuple, permute_tuple)\n@inline convert_structure!(x::MPIStateArray, y, reshape_tuple, permute_tuple) =\n convert_structure!(x.realdata, y, reshape_tuple, permute_tuple)\n\nfunction check_convergence(residual_norm, initial_residual_norm, atol, rtol)\n relative_residual = residual_norm / initial_residual_norm\n converged = false\n if (residual_norm ≤ atol || relative_residual ≤ rtol)\n converged = true\n end\n return converged\nend\n" }, { "path": "src/Numerics/SystemSolvers/columnwise_lu_solver.jl", "content": "#### Columnwise LU Solver\n\nexport ManyColumnLU, SingleColumnLU\n\nabstract type AbstractColumnLUSolver <: AbstractSystemSolver end\n\n\"\"\"\n ManyColumnLU()\n\nThis solver is used for systems that are block diagonal where each block is\nassociated with a column of the mesh. The systems are solved using a\nnon-pivoted LU factorization.\n\"\"\"\nstruct ManyColumnLU <: AbstractColumnLUSolver end\n\n\"\"\"\n SingleColumnLU()\n\nThis solver is used for systems that are block diagonal where each block is\nassociated with a column of the mesh. Moreover, each block is assumed to be\nthe same. The systems are solved using a non-pivoted LU factorization.\n\"\"\"\nstruct SingleColumnLU <: AbstractColumnLUSolver end\n\nstruct ColumnwiseLU{AT}\n A::AT\nend\n\nstruct DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC, AT}\n data::AT\nend\nDGColumnBandedMatrix(\n A::DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC},\n data,\n) where {D, P, NS, EH, EV, EB, SC} =\n DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC, typeof(data)}(data)\nBase.eltype(A::DGColumnBandedMatrix) = eltype(A.data)\nBase.size(A::DGColumnBandedMatrix) = size(A.data)\ndimensionality(::DGColumnBandedMatrix{D}) where {D} = D\npolynomialorders(::DGColumnBandedMatrix{D, P}) where {D, P} = P\n# polynomialorders is polynomial orders P, which is a tuple,\n# vertical_polynomialorder is the vertical polynomial order\nvertical_polynomialorder(::DGColumnBandedMatrix{D, P}) where {D, P} = P[end]\nnum_state(::DGColumnBandedMatrix{D, P, NS}) where {D, P, NS} = NS\nnum_horz_elem(::DGColumnBandedMatrix{D, P, NS, EH}) where {D, P, NS, EH} = EH\nnum_vert_elem(\n ::DGColumnBandedMatrix{D, P, NS, EH, EV},\n) where {D, P, NS, EH, EV} = EV\nelem_band(\n ::DGColumnBandedMatrix{D, P, NS, EH, EV, EB},\n) where {D, P, NS, EH, EV, EB} = EB\nsingle_column(\n ::DGColumnBandedMatrix{D, P, NS, EH, EV, EB, SC},\n) where {D, P, NS, EH, EV, EB, SC} = SC\n\n# DG: lower_bandwidth is Nq_v*nstate * eband - 1, (does not include itself) \n# eband = 1 for inviscid, since nodal point at the face communicate to the overlaping point to its neighbor\n# and other points only communicate to points in the same element\n# eband = 2 for visous\n#\n# FV: lower_bandwidth is nstate * (stencil_width + 1 + 1) - 1, \n# since the reconstruction states depend on stencil_width points on each side, \n# and the flux depends on stencil_width + 1 points on each side\n# eband = (stencil_width + 2) for inviscid\n# eband = max{ (stencil_width + 2), 3} for viscous, \n# since the viscous flux is computed by applying first order FD twice\n# \n# The lower_bandwidth ls formulated as Nq_v*nstate * eband - 1\nlower_bandwidth(N, nstate, eband) = (N + 1) * nstate * eband - 1\nlower_bandwidth(A::DGColumnBandedMatrix) =\n lower_bandwidth(vertical_polynomialorder(A), num_state(A), elem_band(A))\nupper_bandwidth(N, nstate, eband) = lower_bandwidth(N, nstate, eband)\nupper_bandwidth(A::DGColumnBandedMatrix) =\n upper_bandwidth(vertical_polynomialorder(A), num_state(A), elem_band(A))\nBase.reshape(A::DGColumnBandedMatrix, args...) =\n DGColumnBandedMatrix(A, reshape(A.data, args...))\nAdapt.adapt_structure(to, A::DGColumnBandedMatrix) =\n DGColumnBandedMatrix(A, adapt(to, A.data))\n\n\nBase.@propagate_inbounds function Base.getindex(A::DGColumnBandedMatrix, I...)\n return A.data[I...]\nend\nBase.@propagate_inbounds function Base.setindex!(\n A::DGColumnBandedMatrix,\n val,\n I...,\n)\n A.data[I...] = val\nend\n\nfunction prefactorize(op, solver::AbstractColumnLUSolver, Q, args...)\n dg = op.f!\n\n # TODO: can we get away with just passing the grid?\n A = banded_matrix(\n op,\n dg,\n similar(Q),\n similar(Q),\n args...;\n single_column = typeof(solver) <: SingleColumnLU,\n )\n\n band_lu!(A)\n\n ColumnwiseLU(A)\nend\n\nfunction linearsolve!(\n linop,\n clu::ColumnwiseLU,\n ::AbstractColumnLUSolver,\n Q,\n Qrhs,\n args...,\n)\n A = clu.A\n Q .= Qrhs\n\n band_forward!(Q, A)\n band_back!(Q, A)\nend\n\n\"\"\"\n band_lu!(A)\n\n\"\"\"\nfunction band_lu!(A)\n device = array_device(A.data)\n\n nstate = num_state(A)\n Nq = polynomialorders(A) .+ 1\n @inbounds Nq_h = Nq[1]\n @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n nhorzelem = num_horz_elem(A)\n\n groupsize = (Nq_h, Nqj)\n ndrange = (nhorzelem * Nq_h, Nqj)\n\n if single_column(A)\n # single column case\n #\n # TODO Would it be faster to copy the matrix to the host and factorize it\n # there?\n groupsize = (1, 1)\n ndrange = groupsize\n A = reshape(A, 1, 1, size(A)..., 1)\n end\n\n event = Event(device)\n event = band_lu_kernel!(device, groupsize)(\n A,\n ndrange = ndrange,\n dependencies = (event,),\n )\n wait(device, event)\nend\n\nfunction band_forward!(Q, A)\n device = array_device(Q)\n\n Nq = polynomialorders(A) .+ 1\n @inbounds Nq_h = Nq[1]\n @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n nhorzelem = num_horz_elem(A)\n\n event = Event(device)\n event = band_forward_kernel!(device, (Nq_h, Nqj))(\n Q.data,\n A,\n ndrange = (nhorzelem * Nq_h, Nqj),\n dependencies = (event,),\n )\n wait(device, event)\nend\n\nfunction band_back!(Q, A)\n device = array_device(Q)\n\n Nq = polynomialorders(A) .+ 1\n @inbounds Nq_h = Nq[1]\n @inbounds Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n nhorzelem = num_horz_elem(A)\n\n event = Event(device)\n event = band_back_kernel!(device, (Nq_h, Nqj))(\n Q.data,\n A,\n ndrange = (nhorzelem * Nq_h, Nqj),\n dependencies = (event,),\n )\n wait(device, event)\nend\n\n\n\"\"\"\n banded_matrix(\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg);\n single_column = false,\n )\n\nForms the banded matrices for each the column operator defined by the `SpaceDiscretization`\ndg. If `single_column=false` then a banded matrix is stored for each column and\nif `single_column=true` only the banded matrix associated with the first column\nof the first element is stored. The bandwidth of the DG column banded matrix is\n`p = q = (vertical_polynomial + 1) * nstate * eband - 1` with `p` and `q` being\nthe upper and lower bandwidths.\n\nThe banded matrices are stored in the LAPACK band storage format\n.\n\nThe banded matrices are returned as an arrays where the array type matches that\nof `Q`. If `single_column=false` then the returned array has 5 dimensions, which\nare:\n- first horizontal column index\n- second horizontal column index\n- band index (-q:p)\n- vertical DOF index with state `s`, vertical DOF index `k`, and vertical\n element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`\n- horizontal element index\n\nIf the `single_column=true` then the returned array has 2 dimensions which are\nthe band index and the vertical DOF index.\n\"\"\"\nfunction banded_matrix(\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg);\n single_column = false,\n)\n banded_matrix(\n (dQ, Q) -> dg(dQ, Q, nothing, 0; increment = false),\n dg,\n Q,\n dQ;\n single_column = single_column,\n )\nend\n\n\"\"\"\n banded_matrix(\n f!,\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg),\n args...;\n single_column = false,\n )\n\nForms the banded matrices for each the column operator defined by the linear\noperator `f!` which is assumed to have the same banded structure as the\n`SpaceDiscretization` dg. If `single_column=false` then a banded matrix is stored for each\ncolumn and if `single_column=true` only the banded matrix associated with the\nfirst column of the first element is stored. The bandwidth of the DG column\nbanded matrix is `p = q = (vertical_polynomial + 1) * nstate * eband - 1` with\n`p` and `q` being the upper and lower bandwidths.\n\nThe banded matrices are stored in the LAPACK band storage format\n.\n\nThe banded matrices are returned as an arrays where the array type matches that\nof `Q`. If `single_column=false` then the returned array has 5 dimensions, which\nare:\n- first horizontal column index\n- second horizontal column index\n- band index (-q:p)\n- vertical DOF index with state `s`, vertical DOF index `k`, and vertical\n element `ev` mapping to `s + nstate * (k - 1) + nstate * nvertelem * (ev - 1)`\n- horizontal element index\n\nIf the `single_column=true` then the returned array has 2 dimensions which are\nthe band index and the vertical DOF index.\n\nHere `args` are passed to `f!`.\n\"\"\"\nfunction banded_matrix(\n f!,\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg),\n args...;\n single_column = false,\n)\n # Initialize banded matrix data structure\n A = empty_banded_matrix(dg, Q; single_column = single_column)\n\n # Populate matrix with data\n update_banded_matrix!(\n A,\n f!,\n dg,\n Q,\n dQ,\n args...;\n single_column = single_column,\n )\n\n A\nend\n\n\"\"\"\n empty_banded_matrix(\n dg::SpaceDiscretization,\n Q::MPIStateArray;\n single_column = false,\n )\n\nInitializes an empty banded matrix stored in the LAPACK band storage format\n.\n\"\"\"\nfunction empty_banded_matrix(\n dg::SpaceDiscretization,\n Q::MPIStateArray;\n single_column = false,\n)\n bl = dg.balance_law\n grid = dg.grid\n topology = grid.topology\n @assert isstacked(topology)\n @assert typeof(dg.direction) <: VerticalDirection\n\n FT = eltype(Q.data)\n device = array_device(Q)\n\n nstate = number_states(bl, Prognostic())\n N = polynomialorders(grid)\n dim = dimensionality(grid)\n Nq = N .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dim == 2 ? 1 : Nq[2]\n Nq_v = Nq[dim]\n end\n\n\n eband =\n (typeof(dg) <: DGModel) ?\n (number_states(bl, GradientFlux()) == 0 ? 1 : 2) :\n (\n number_states(bl, GradientFlux()) == 0 ?\n width(dg.fv_reconstruction) + 2 :\n max(width(dg.fv_reconstruction) + 2, 3)\n ) # else: DGFVModel\n\n p = lower_bandwidth(N[dim], nstate, eband)\n q = upper_bandwidth(N[dim], nstate, eband)\n\n nrealelem = length(topology.realelems)\n nvertelem = topology.stacksize\n nhorzelem = div(nrealelem, nvertelem)\n\n # first horizontal DOF index\n # second horizontal DOF index\n # band index -q:p\n # vertical DOF index\n # horizontal element index\n A = if single_column\n similar(Q.data, p + q + 1, Nq_v * nstate * nvertelem)\n else\n similar(Q.data, Nq_h, Nqj, p + q + 1, Nq_v * nstate * nvertelem, nhorzelem)\n end\n fill!(A, zero(FT))\n\n A = DGColumnBandedMatrix{\n dim,\n N,\n nstate,\n nhorzelem,\n nvertelem,\n eband,\n single_column,\n typeof(A),\n }(\n A,\n )\n\n A\nend\n\n\"\"\"\n update_banded_matrix!(\n A::DGColumnBandedMatrix,\n f!,\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg),\n args...;\n single_column = false,\n )\n\nUpdates the banded matrices for each the column operator defined by the linear\noperator `f!` which is assumed to have the same banded structure as the\n`SpaceDiscretization` dg. If `single_column=false` then a banded matrix is stored for each\ncolumn and if `single_column=true` only the banded matrix associated with the\nfirst column of the first element is stored. The bandwidth of the DG column\nbanded matrix is `p = q = (vertical_polynomial + 1) * nstate * eband - 1` with\n`p` and `q` being the upper and lower bandwidths.\n\nHere `args` are passed to `f!`.\n\"\"\"\nfunction update_banded_matrix!(\n A::DGColumnBandedMatrix,\n f!,\n dg::SpaceDiscretization,\n Q::MPIStateArray = MPIStateArray(dg),\n dQ::MPIStateArray = MPIStateArray(dg),\n args...;\n single_column = false,\n)\n bl = dg.balance_law\n grid = dg.grid\n topology = grid.topology\n @assert isstacked(topology)\n @assert typeof(dg.direction) <: VerticalDirection\n\n FT = eltype(Q.data)\n device = array_device(Q)\n\n nstate = number_states(bl, Prognostic())\n N = polynomialorders(grid)\n dim = dimensionality(grid)\n Nq = N .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dim == 2 ? 1 : Nq[2]\n Nq_v = Nq[dim]\n end\n\n eband = elem_band(A)\n\n nrealelem = length(topology.realelems)\n nvertelem = topology.stacksize\n nhorzelem = div(nrealelem, nvertelem)\n\n # loop through all DOFs in a column and compute the matrix column\n # loop only the first min(nvertelem, 2eband+1) elements\n # in each element loop, updating these columns correspond\n # to elements (ev :2eband+1 : nvertelem)\n for ev in 1:min(nvertelem, 2eband + 1)\n for s in 1:nstate\n for k in 1:Nq_v\n # Set a single 1 per column and rest 0\n event = Event(device)\n event = kernel_set_banded_data!(device, (Nq_h, Nqj, Nq_v))(\n Q.data,\n A,\n k,\n s,\n ev,\n 1:nhorzelem,\n 1:nvertelem;\n ndrange = (nvertelem * Nq_h, nhorzelem * Nqj, Nq_v),\n dependencies = (event,),\n )\n wait(device, event)\n\n # Get the matrix column\n f!(dQ, Q, args...)\n\n # Store the banded matrix\n event = Event(device)\n event = kernel_set_banded_matrix!(device, (Nq_h, Nqj, Nq_v))(\n A,\n dQ.data,\n k,\n s,\n ev,\n 1:nhorzelem,\n (-eband):eband;\n ndrange = ((2eband + 1) * Nq_h, nhorzelem * Nqj, Nq_v),\n dependencies = (event,),\n )\n wait(device, event)\n end\n end\n end\nend\n\n\"\"\"\n banded_matrix_vector_product!(\n A,\n dQ::MPIStateArray,\n Q::MPIStateArray\n )\n\nCompute a matrix vector product `dQ = A * Q` where `A` is assumed to be a matrix\ncreated using the `banded_matrix` function.\n\nThis function is primarily for testing purposes.\n\"\"\"\nfunction banded_matrix_vector_product!(A, dQ::MPIStateArray, Q::MPIStateArray)\n device = array_device(Q)\n\n Nq = polynomialorders(A) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n Nq_v = Nq[end]\n end\n nvertelem = num_vert_elem(A)\n nhorzelem = num_horz_elem(A)\n\n event = Event(device)\n event = kernel_banded_matrix_vector_product!(device, (Nq_h, Nqj, Nq_v))(\n dQ.data,\n A,\n Q.data,\n 1:nhorzelem,\n 1:nvertelem;\n ndrange = (nvertelem * Nq_h, nhorzelem * Nqj, Nq_v),\n dependencies = (event,),\n )\n wait(device, event)\nend\n\nusing StaticArrays\nusing KernelAbstractions.Extras: @unroll\n\n@doc \"\"\"\n band_lu_kernel!(A)\n\nThis performs Band Gaussian Elimination (Algorithm 4.3.1 of Golub and Van\nLoan). The array `A` contains a band matrix for each vertical column. For\nexample, `A[i, j, :, :, h]`, is the band matrix associated with the `(i, j)`th\ndegree of freedom in the horizontal element `h`.\n\nEach `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower\nbandwidth `p` where `n = nstate * Nq * nvertelem` and\n`p = q = (vertical_polynomial + 1) * nstate * eband - 1` \n\nEach band matrix is stored in the [LAPACK band storage](https://www.netlib.org/lapack/lug/node124.html).\nFor example the band matrix\n\n B = [b₁₁ b₁₂ 0 0 0\n b₂₁ b₂₂ b₂₃ 0 0\n b₃₁ b₃₂ b₃₃ b₃₄ 0\n 0 b₄₂ b₄₃ b₄₄ b₄₅\n 0 0 b₅₃ b₅₄ b₅₅]\n\nis stored as\n\n B = [0 b₁₂ b₂₃ b₃₄ b₄₅\n b₁₁ b₂₂ b₃₃ b₄₄ b₅₅\n b₂₁ b₃₂ b₄₃ b₅₄ 0\n b₃₁ b₄₂ b₅₃ 0 0]\n\n### Reference\n\n - [GolubVanLoan2013](@cite)\n\n\"\"\" band_lu_kernel!\n@kernel function band_lu_kernel!(A)\n @uniform begin\n Nq = polynomialorders(A) .+ 1\n @inbounds Nq_h = Nq[1]\n @inbounds Nq_v = Nq[end]\n nstate = num_state(A)\n nvertelem = num_vert_elem(A)\n n = nstate * Nq_v * nvertelem\n p, q = lower_bandwidth(A), upper_bandwidth(A)\n end\n\n h = @index(Group, Linear)\n i, j = @index(Local, NTuple)\n\n @inbounds begin\n for v in 1:nvertelem\n for k in 1:Nq_v\n for s in 1:nstate\n kk = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v\n\n Aq = A[i, j, q + 1, kk, h]\n for ii in 1:p\n A[i, j, q + ii + 1, kk, h] /= Aq\n end\n\n for jj in 1:q\n if jj + kk ≤ n\n Ajj = A[i, j, q - jj + 1, jj + kk, h]\n for ii in 1:p\n A[i, j, q + ii - jj + 1, jj + kk, h] -=\n A[i, j, q + ii + 1, kk, h] * Ajj\n end\n end\n end\n end\n end\n end\n end\nend\n\n@doc \"\"\"\n band_forward_kernel!(b, LU)\n\nThis performs Band Forward Substitution (Algorithm 4.3.2 of Golub and Van\nLoan), i.e., the right-hand side `b` is replaced with the solution of `L*x=b`.\n\nThe array `b` is of the size `(Nq * Nqj * Nq, nstate, nvertelem * nhorzelem)`.\n\nThe LU-factorization array `LU` contains a single band matrix or one\nfor each vertical column, see [`band_lu!`](@ref).\n\nEach `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower\nbandwidth `p` where `n = nstate * Nq * nvertelem` and\n`p = q = (vertical_polynomial + 1) * nstate * eband - 1` \n\n### Reference\n\n - [GolubVanLoan2013](@cite)\n\n\"\"\" band_forward_kernel!\n@kernel function band_forward_kernel!(b, LU)\n @uniform begin\n FT = eltype(b)\n nstate = num_state(LU)\n Nq = polynomialorders(LU) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dimensionality(LU) == 2 ? 1 : Nq[2]\n Nq_v = Nq[end]\n end\n nvertelem = num_vert_elem(LU)\n n = nstate * Nq_v * nvertelem\n eband = elem_band(LU)\n p, q = lower_bandwidth(LU), upper_bandwidth(LU)\n\n l_b = MArray{Tuple{p + 1}, FT}(undef)\n end\n\n h = @index(Group, Linear)\n i, j = @index(Local, NTuple)\n\n @inbounds begin\n @unroll for v in 1:eband\n @unroll for k in 1:Nq_v\n @unroll for s in 1:nstate\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n ee = v + nvertelem * (h - 1)\n ii = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v\n l_b[ii] = nvertelem ≥ v ? b[ijk, s, ee] : zero(FT)\n end\n end\n end\n\n for v in 1:nvertelem\n @unroll for k in 1:Nq_v\n @unroll for s in 1:nstate\n jj = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v\n\n @unroll for ii in 2:(p + 1)\n Lii =\n single_column(LU) ? LU[ii + q, jj] :\n LU[i, j, ii + q, jj, h]\n l_b[ii] -= Lii * l_b[1]\n end\n\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n ee = v + nvertelem * (h - 1)\n\n b[ijk, s, ee] = l_b[1]\n\n @unroll for ii in 1:p\n l_b[ii] = l_b[ii + 1]\n end\n\n if jj + p < n\n (idx, si) = fldmod1(jj + p + 1, nstate)\n (vi, ki) = fldmod1(idx, Nq_v)\n\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (ki - 1)\n ee = vi + nvertelem * (h - 1)\n\n l_b[p + 1] = b[ijk, si, ee]\n end\n end\n end\n end\n end\nend\n\n@doc \"\"\"\n band_back_kernel!(b, LU)\n\nThis performs Band Back Substitution (Algorithm 4.3.3 of Golub and Van\nLoan), i.e., the right-hand side `b` is replaced with the solution of `U*x=b`.\n\nThe array `b` is of the size `(Nq * Nqj * Nq, nstate, nvertelem * nhorzelem)`.\n\nThe LU-factorization array `LU` contains a single band matrix or one\nfor each vertical column, see [`band_lu!`](@ref).\n\nEach `n` by `n` band matrix is assumed to have upper bandwidth `q` and lower\nbandwidth `p` where `n = nstate * Nq * nvertelem` and\n`p = q = (vertical_polynomial + 1) * nstate * eband - 1` \n\n### Reference\n\n - [GolubVanLoan2013](@cite)\n\n\"\"\" band_back_kernel!\n@kernel function band_back_kernel!(b, LU)\n @uniform begin\n FT = eltype(b)\n nstate = num_state(LU)\n Nq = polynomialorders(LU) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dimensionality(LU) == 2 ? 1 : Nq[2]\n Nq_v = Nq[end]\n end\n nvertelem = num_vert_elem(LU)\n n = nstate * Nq_h * nvertelem\n q = upper_bandwidth(LU)\n eband = elem_band(LU)\n\n l_b = MArray{Tuple{q + 1}, FT}(undef)\n end\n\n h = @index(Group, Linear)\n i, j = @index(Local, NTuple)\n\n @inbounds begin\n @unroll for v in nvertelem:-1:(nvertelem - eband + 1)\n @unroll for k in Nq_v:-1:1\n @unroll for s in nstate:-1:1\n vi = eband - nvertelem + v\n ii = s + (k - 1) * nstate + (vi - 1) * nstate * Nq_v\n\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n ee = v + nvertelem * (h - 1)\n\n l_b[ii] = b[ijk, s, ee]\n end\n end\n end\n\n for v in nvertelem:-1:1\n @unroll for k in Nq_v:-1:1\n @unroll for s in nstate:-1:1\n jj = s + (k - 1) * nstate + (v - 1) * nstate * Nq_v\n\n l_b[q + 1] /=\n single_column(LU) ? LU[q + 1, jj] :\n LU[i, j, q + 1, jj, h]\n\n @unroll for ii in 1:q\n Uii =\n single_column(LU) ? LU[ii, jj] : LU[i, j, ii, jj, h]\n l_b[ii] -= Uii * l_b[q + 1]\n end\n\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n ee = v + nvertelem * (h - 1)\n\n b[ijk, s, ee] = l_b[q + 1]\n\n @unroll for ii in q:-1:1\n l_b[ii + 1] = l_b[ii]\n end\n\n if jj - q > 1\n (idx, si) = fldmod1(jj - q - 1, nstate)\n (vi, ki) = fldmod1(idx, Nq_v)\n\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (ki - 1)\n ee = vi + nvertelem * (h - 1)\n\n l_b[1] = b[ijk, si, ee]\n end\n end\n end\n end\n end\nend\n\n\n\n### TODO: Document this\n@kernel function kernel_set_banded_data!(\n Q,\n A::DGColumnBandedMatrix,\n kin,\n sin,\n evin0,\n helems,\n velems,\n)\n @uniform begin\n FT = eltype(Q)\n nstate = num_state(A)\n Nq = polynomialorders(A) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nq_v = Nq[end]\n Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n end\n nvertelem = num_vert_elem(A)\n eband = elem_band(A)\n end\n\n ev, eh = @index(Group, NTuple)\n i, j, k = @index(Local, NTuple)\n\n @inbounds begin\n e = ev + (eh - 1) * nvertelem\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n @unroll for s in 1:nstate\n if k == kin && s == sin && ((ev - evin0) % (2eband + 1) == 0)\n Q[ijk, s, e] = 1\n else\n Q[ijk, s, e] = 0\n end\n end\n end\nend\n\n\n@kernel function kernel_set_banded_matrix!(\n A,\n dQ,\n kin,\n sin,\n evin0,\n helems,\n vpelems,\n)\n @uniform begin\n FT = eltype(A)\n nstate = num_state(A)\n Nq = polynomialorders(A) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n Nq_v = Nq[end]\n end\n nvertelem = num_vert_elem(A)\n p = lower_bandwidth(A)\n q = upper_bandwidth(A)\n\n eband = elem_band(A)\n eshift = elem_band(A) + 1\n end\n\n ep, eh = @index(Group, NTuple)\n ep = ep - eshift\n i, j, k = @index(Local, NTuple)\n\n for evin in evin0:(2eband + 1):nvertelem\n # sin, kin, evin are the state, vertical dof, and vert element we are\n # handling\n # column index of matrix\n jj = sin + (kin - 1) * nstate + (evin - 1) * nstate * Nq_v\n\n # one thread is launch for dof that might contribute to column jj's band\n @inbounds begin\n # ep is the shift we need to add to evin to get the element we need to\n # consider\n ev = ep + evin\n if 1 ≤ ev ≤ nvertelem\n e = ev + (eh - 1) * nvertelem\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n @unroll for s in 1:nstate\n # row index of matrix\n ii = s + (k - 1) * nstate + (ev - 1) * nstate * Nq_v\n # row band index\n bb = ii - jj\n # make sure we're in the bandwidth\n if -q ≤ bb ≤ p\n if !single_column(A)\n A[i, j, bb + q + 1, jj, eh] = dQ[ijk, s, e]\n else\n if (i, j, eh) == (1, 1, 1)\n A[bb + q + 1, jj] = dQ[ijk, s, e]\n end\n end\n end\n end\n end\n end\n end\nend\n\n@kernel function kernel_banded_matrix_vector_product!(dQ, A, Q, helems, velems)\n @uniform begin\n FT = eltype(A)\n nstate = num_state(A)\n\n Nq = polynomialorders(A) .+ 1\n @inbounds begin\n Nq_h = Nq[1]\n Nq_v = Nq[end]\n Nqj = dimensionality(A) == 2 ? 1 : Nq[2]\n end\n eband = elem_band(A)\n nvertelem = num_vert_elem(A)\n p = lower_bandwidth(A)\n q = upper_bandwidth(A)\n\n elo = eband - 1\n eup = eband - 1\n end\n\n ev, eh = @index(Group, NTuple)\n i, j, k = @index(Local, NTuple)\n\n # matrix row loops\n @inbounds begin\n e = ev + nvertelem * (eh - 1)\n @unroll for s in 1:nstate\n Ax = -zero(FT)\n ii = s + (k - 1) * nstate + (ev - 1) * nstate * Nq_v\n\n # banded matrix column loops\n @unroll for evv in max(1, ev - elo):min(nvertelem, ev + eup)\n ee = evv + nvertelem * (eh - 1)\n @unroll for kk in 1:Nq_v\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (kk - 1)\n @unroll for ss in 1:nstate\n jj = ss + (kk - 1) * nstate + (evv - 1) * nstate * Nq_v\n bb = ii - jj\n if -q ≤ bb ≤ p\n if !single_column(A)\n Ax +=\n A[i, j, bb + q + 1, jj, eh] * Q[ijk, ss, ee]\n else\n Ax += A[bb + q + 1, jj] * Q[ijk, ss, ee]\n end\n end\n end\n end\n end\n ijk = i + Nqj * (j - 1) + Nq_h * Nqj * (k - 1)\n dQ[ijk, s, e] = Ax\n end\n end\nend\n" }, { "path": "src/Numerics/SystemSolvers/conjugate_gradient_solver.jl", "content": "#### Conjugate Gradient solver\n\nexport ConjugateGradient\n\nstruct ConjugateGradient{AT1, AT2, FT, RD, RT, IT} <:\n AbstractIterativeSystemSolver\n # tolerances (2)\n rtol::FT\n atol::FT\n # arrays of size reshape_tuple (7)\n r0::AT1\n z0::AT1\n p0::AT1\n r1::AT1\n z1::AT1\n p1::AT1\n Lp::AT1\n # arrays of size(MPIStateArray) which are aliased to two of the previous dimensions (2)\n alias_p0::AT2\n alias_Lp::AT2\n # reduction dimension (1)\n dims::RD\n # reshape dimension (1)\n reshape_tuple::RT\n # maximum number of iterations (1)\n max_iter::IT\n\nend\n\n# Define the outer constructor for the ConjugateGradient struct\n\"\"\"\n ConjugateGradient(\n Q::AT;\n rtol = eps(eltype(Q)),\n atol = eps(eltype(Q)),\n max_iter = length(Q),\n dims = :,\n reshape_tuple = size(Q),\n ) where {AT}\n\n# ConjugateGradient\n\n# Description\n- Outer constructor for the ConjugateGradient struct\n\n# Arguments\n- `Q`:(array). The kind of object that linearoperator! acts on.\n\n# Keyword Arguments\n- `rtol`: (float). relative tolerance\n- `atol`: (float). absolute tolerance\n- `dims`: (tuple or : ). the dimensions to reduce over\n- `reshape_tuple`: (tuple). the dimensions that the conjugate gradient solver operators over\n\n# Comment\n- The reshape tuple is necessary in case the linearoperator! is defined over vectors of a different size as compared to what plays nicely with the dimension reduction in the ConjugateGradient. It also allows the user to define preconditioners over arrays that are more convenienently shaped.\n\n# Return\n- ConjugateGradient struct\n\"\"\"\nfunction ConjugateGradient(\n Q::AT;\n rtol = eps(eltype(Q)),\n atol = eps(eltype(Q)),\n max_iter = length(Q),\n dims = :,\n reshape_tuple = size(Q),\n) where {AT}\n\n # allocate arrays (5)\n r0 = reshape(similar(Q), reshape_tuple)\n z0 = reshape(similar(Q), reshape_tuple)\n r1 = reshape(similar(Q), reshape_tuple)\n z1 = reshape(similar(Q), reshape_tuple)\n p1 = reshape(similar(Q), reshape_tuple)\n # allocate array of different shape (2)\n alias_p0 = similar(Q)\n alias_Lp = similar(Q)\n # allocate create aliased arrays (2)\n p0 = reshape(alias_p0, reshape_tuple)\n Lp = reshape(alias_Lp, reshape_tuple)\n\n container = [\n rtol,\n atol,\n r0,\n z0,\n p0,\n r1,\n z1,\n p1,\n Lp,\n alias_p0,\n alias_Lp,\n dims,\n reshape_tuple,\n max_iter,\n ]\n # create struct instance by splatting the container into the default constructor\n return ConjugateGradient{\n typeof(z0),\n typeof(Q),\n eltype(Q),\n typeof(dims),\n typeof(reshape_tuple),\n typeof(max_iter),\n }(container...)\nend\n\n# Define the outer constructor for the ConjugateGradient struct\n\"\"\"\n ConjugateGradient(\n Q::MPIStateArray;\n rtol = eps(eltype(Q)),\n atol = eps(eltype(Q)),\n max_iter = length(Q),\n dims = :,\n reshape_tuple = size(Q),\n )\n\n# Description\nOuter constructor for the ConjugateGradient struct with MPIStateArrays. THIS IS A HACK DUE TO RESHAPE FUNCTIONALITY ON MPISTATEARRAYS.\n\n# Arguments\n- `Q`:(array). The kind of object that linearoperator! acts on.\n\n# Keyword Arguments\n- `rtol`: (float). relative tolerance\n- `atol`: (float). absolute tolerance\n- `dims`: (tuple or : ). the dimensions to reduce over\n- `reshape_tuple`: (tuple). the dimensions that the conjugate gradient solver operators over\n\n# Comment\n- The reshape tuple is necessary in case the linearoperator! is defined over vectors of a different size as compared to what plays nicely with the dimension reduction in the ConjugateGradient. It also allows the user to define preconditioners over arrays that are more convenienently shaped.\n\n# Return\n- ConjugateGradient struct\n\"\"\"\nfunction ConjugateGradient(\n Q::MPIStateArray;\n rtol = eps(eltype(Q)),\n atol = eps(eltype(Q)),\n max_iter = length(Q),\n dims = :,\n reshape_tuple = size(Q),\n)\n\n # create empty container for pushing struct objects\n # allocate arrays (5)\n r0 = reshape(similar(Q.data), reshape_tuple)\n z0 = reshape(similar(Q.data), reshape_tuple)\n r1 = reshape(similar(Q.data), reshape_tuple)\n z1 = reshape(similar(Q.data), reshape_tuple)\n p1 = reshape(similar(Q.data), reshape_tuple)\n # allocate array of different shape (2)\n alias_p0 = similar(Q)\n alias_Lp = similar(Q)\n # allocate create aliased arrays (2)\n p0 = reshape(alias_p0.data, reshape_tuple)\n Lp = reshape(alias_Lp.data, reshape_tuple)\n container = [\n rtol,\n atol,\n r0,\n z0,\n p0,\n r1,\n z1,\n p1,\n Lp,\n alias_p0,\n alias_Lp,\n dims,\n reshape_tuple,\n max_iter,\n ]\n # create struct instance by splatting the container into the default constructor\n return ConjugateGradient{\n typeof(z0),\n typeof(Q),\n eltype(Q),\n typeof(dims),\n typeof(reshape_tuple),\n typeof(max_iter),\n }(container...)\nend\n\n\n\"\"\"\n initialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::ConjugateGradient,\n args...,\n )\n\n# Description\n\n- This function initializes the iterative solver. It is called as part of the AbstractIterativeSystemSolver routine. SEE CODEREF for documentation on AbstractIterativeSystemSolver\n\n# Arguments\n\n- `linearoperator!`: (function). This applies the predefined linear operator on an array. Applies a linear operator to object \"y\" and overwrites object \"z\". The function argument i s linearoperator!(z,y, args...) and it returns nothing.\n- `Q`: (array). This is an object that linearoperator! outputs\n- `Qrhs`: (array). This is an object that linearoperator! acts on\n- `solver`: (struct). This is a scruct for dispatch, in this case for ColumnwisePreconditionedConjugateGradient\n- `args...`: (arbitrary). This is optional arguments that can be passed into linearoperator! function.\n\n# Keyword Arguments\n\n- There are no keyword arguments\n\n# Return\n- `converged`: (bool). A boolean to say whether or not the iterative solver has converged.\n- `threshold`: (float). The value of the residual for the first timestep\n\n# Comment\n- This function does nothing for conjugate gradient\n\n\"\"\"\nfunction initialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::ConjugateGradient,\n args...,\n)\n\n return false, Inf\nend\n\n\n\"\"\"\n doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::ConjugateGradient,\n threshold,\n args...;\n applyPC! = (x, y) -> x .= y,\n )\n\n# Description\n\n- This function enacts the iterative solver. It is called as part of the AbstractIterativeSystemSolver routine. SEE CODEREF for documentation on AbstractIterativeSystemSolver\n\n# Arguments\n\n- `linearoperator!`: (function). This applies the predefined linear operator on an array. Applies a linear operator to object \"y\" and overwrites object \"z\". It is a function with arguments linearoperator!(z,y, args...), where \"z\" gets overwritten by \"y\" and \"args...\" are additional arguments passed to the linear operator. The linear operator is assumed to return nothing.\n- `Q`: (array). This is an object that linearoperator! overwrites\n- `Qrhs`: (array). This is an object that linearoperator! acts on. This is the rhs to the linear system\n- `solver`: (struct). This is a scruct for dispatch, in this case for ConjugateGradient\n- `threshold`: (float). Either an absolute or relative tolerance\n- `applyPC!`: (function). Applies a preconditioner to objecy \"y\" and overwrites object \"z\". applyPC!(z,y)\n- `args...`: (arbitrary). This is necessary for the linearoperator! function which has a signature linearoperator!(b, x, args....)\n\n# Keyword Arguments\n\n- There are no keyword arguments\n\n# Return\n- `converged`: (bool). A boolean to say whether or not the iterative solver has converged.\n- `iteration`: (int). Iteration number for the iterative solver\n- `threshold`: (float). The value of the residual for the first timestep\n\n# Comment\n- This function does conjugate gradient\n\n\"\"\"\nfunction doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::ConjugateGradient,\n threshold,\n args...;\n applyPC! = (x, y) -> x .= y,\n)\n\n\n # unroll names for convenience\n\n rtol = solver.rtol\n atol = solver.atol\n residual_norm = typemax(eltype(Q))\n dims = solver.dims\n converged = false\n\n max_iter = solver.max_iter\n r0 = solver.r0\n z0 = solver.z0\n p0 = solver.p0\n r1 = solver.r1\n z1 = solver.z1\n p1 = solver.p1\n Lp = solver.Lp\n alias_p0 = solver.alias_p0\n alias_Lp = solver.alias_Lp\n alias_Q = reshape(Q, solver.reshape_tuple)\n\n # Smack residual by linear operator\n linearoperator!(alias_Lp, Q, args...)\n # make sure that arrays are of the appropriate size\n alias_p0 .= Qrhs\n r0 .= p0 - Lp\n # apply the preconditioner\n applyPC!(z0, r0)\n # update p0\n p0 .= z0\n\n # TODO: FIX THIS\n absolute_residual = maximum(sqrt.(sum(r0 .* r0, dims = dims)))\n relative_residual =\n absolute_residual / maximum(sqrt.(sum(Qrhs .* Qrhs, dims = :)))\n # TODO: FIX THIS\n if (absolute_residual <= atol) || (relative_residual <= rtol)\n # wow! what a great guess\n converged = true\n return converged, 1, absolute_residual\n end\n\n for j in 1:max_iter\n linearoperator!(alias_Lp, alias_p0, args...)\n\n α = sum(r0 .* z0, dims = dims) ./ sum(p0 .* Lp, dims = dims)\n\n # Update along preconditioned direction, (note that broadcast will indeed work as expected)\n @. alias_Q += α * p0\n\n @. r1 = r0 - α * Lp\n\n # TODO: FIX THIS\n absolute_residual = maximum(sqrt.(sum(r1 .* r1, dims = dims)))\n relative_residual =\n absolute_residual / maximum(sqrt.(sum(Qrhs .* Qrhs, dims = :)))\n # TODO: FIX THIS\n converged = false\n if (absolute_residual <= atol) || (relative_residual <= rtol)\n converged = true\n return converged, j, absolute_residual\n end\n\n applyPC!(z1, r1)\n\n β = sum(z1 .* r1, dims = dims) ./ sum(z0 .* r0, dims = dims)\n\n # Update\n @. p0 = z1 + β * p0\n @. z0 = z1\n @. r0 = r1\n\n end\n\n # TODO: FIX THIS\n converged = true\n return converged, max_iter, absolute_residual\nend\n\n\"\"\"\n doiteration!(\n linearoperator!,\n preconditioner,\n Q::MPIStateArray,\n Qrhs::MPIStateArray,\n solver::ConjugateGradient,\n threshold,\n args...;\n applyPC! = (x, y) -> x .= y,\n )\n\n# Description\n\nThis function enacts the iterative solver. It is called as part of the AbstractIterativeSystemSolver routine. SEE CODEREF for documentation on AbstractIterativeSystemSolver. THIS IS A HACK TO WORK WITH MPISTATEARRAYS. THE ISSUE IS WITH RESHAPE.\n\n# Arguments\n\n- `linearoperator!`: (function). This applies the predefined linear operator on an array. Applies a linear operator to object \"y\" and overwrites object \"z\". It is a function with arguments linearoperator!(z,y, args...), where \"z\" gets overwritten by \"y\" and \"args...\" are additional arguments passed to the linear operator. The linear operator is assumed to return nothing.\n- `Q`: (array). This is an object that linearoperator! overwrites\n- `Qrhs`: (array). This is an object that linearoperator! acts on. This is the rhs to the linear system\n- `solver`: (struct). This is a scruct for dispatch, in this case for ConjugateGradient\n- `threshold`: (float). Either an absolute or relative tolerance\n- `applyPC!`: (function). Applies a preconditioner to objecy \"y\" and overwrites object \"z\". applyPC!(z,y)\n- `args...`: (arbitrary). This is necessary for the linearoperator! function which has a signature linearoperator!(b, x, args....)\n\n# Keyword Arguments\n\n- There are no keyword arguments\n\n# Return\n- `converged`: (bool). A boolean to say whether or not the iterative solver has converged.\n- `iteration`: (int). Iteration number for the iterative solver\n- `threshold`: (float). The value of the residual for the first timestep\n\n# Comment\n- This function does conjugate gradient\n\n\"\"\"\nfunction doiteration!(\n linearoperator!,\n preconditioner,\n Q::MPIStateArray,\n Qrhs::MPIStateArray,\n solver::ConjugateGradient,\n threshold,\n args...;\n applyPC! = (x, y) -> x .= y,\n)\n\n\n # unroll names for convenience\n\n rtol = solver.rtol\n atol = solver.atol\n residual_norm = typemax(eltype(Q))\n dims = solver.dims\n converged = false\n\n max_iter = solver.max_iter\n r0 = solver.r0\n z0 = solver.z0\n p0 = solver.p0\n r1 = solver.r1\n z1 = solver.z1\n p1 = solver.p1\n Lp = solver.Lp\n alias_p0 = solver.alias_p0\n alias_Lp = solver.alias_Lp\n alias_Q = reshape(Q.data, solver.reshape_tuple)\n\n # Smack residual by linear operator\n linearoperator!(alias_Lp, Q, args...)\n # make sure that arrays are of the appropriate size\n alias_p0 .= Qrhs.data\n r0 .= p0 - Lp\n # apply the preconditioner\n applyPC!(z0, r0)\n # update p0\n p0 .= z0\n\n # TODO: FIX THIS\n absolute_residual = maximum(sqrt.(sum(r0 .* r0, dims = dims)))\n relative_residual =\n absolute_residual / maximum(sqrt.(sum(Qrhs .* Qrhs, dims = :)))\n # TODO: FIX THIS\n if (absolute_residual <= atol) || (relative_residual <= rtol)\n # wow! what a great guess\n converged = true\n return converged, 1, absolute_residual\n end\n\n for j in 1:max_iter\n linearoperator!(alias_Lp, alias_p0, args...)\n\n α = sum(r0 .* z0, dims = dims) ./ sum(p0 .* Lp, dims = dims)\n\n # Update along preconditioned direction, (note that broadcast will indeed work as expected)\n @. alias_Q += α * p0\n\n @. r1 = r0 - α * Lp\n\n # TODO: FIX THIS\n absolute_residual = maximum(sqrt.(sum(r1 .* r1, dims = dims)))\n relative_residual =\n absolute_residual / maximum(sqrt.(sum(Qrhs .* Qrhs, dims = :)))\n # TODO: FIX THIS\n converged = false\n if (absolute_residual <= atol) || (relative_residual <= rtol)\n converged = true\n return converged, j, absolute_residual\n end\n\n applyPC!(z1, r1)\n\n β = sum(z1 .* r1, dims = dims) ./ sum(z0 .* r0, dims = dims)\n\n # Update\n @. p0 = z1 + β * p0\n @. z0 = z1\n @. r0 = r1\n\n end\n\n # TODO: FIX THIS\n converged = true\n return converged, max_iter, absolute_residual\nend\n" }, { "path": "src/Numerics/SystemSolvers/generalized_conjugate_residual_solver.jl", "content": "#### Generalized Conjugate Residual Solver\n\nexport GeneralizedConjugateResidual\n\n\"\"\"\n GeneralizedConjugateResidual(K, Q; rtol, atol)\n\n# Conjugate Residual\n\nThis is an object for solving linear systems using an iterative Krylov method.\nThe constructor parameter `K` is the number of steps after which the algorithm\nis restarted (if it has not converged), `Q` is a reference state used only\nto allocate the solver internal state, and `tolerance` specifies the convergence\ncriterion based on the relative residual norm. The amount of memory\nrequired by the solver state is roughly `(2K + 2) * size(Q)`.\nThis object is intended to be passed to the [`linearsolve!`](@ref) command.\n\nThis uses the restarted Generalized Conjugate Residual method of Eisenstat (1983).\n\n## References\n\n - [Eisenstat1983](@cite)\n\"\"\"\nmutable struct GeneralizedConjugateResidual{K, T, AT} <:\n AbstractIterativeSystemSolver\n residual::AT\n L_residual::AT\n p::NTuple{K, AT}\n L_p::NTuple{K, AT}\n alpha::MArray{Tuple{K}, T, 1, K}\n normsq::MArray{Tuple{K}, T, 1, K}\n rtol::T\n atol::T\n\n function GeneralizedConjugateResidual(\n K,\n Q::AT;\n rtol = √eps(eltype(AT)),\n atol = eps(eltype(AT)),\n ) where {AT}\n T = eltype(Q)\n\n residual = similar(Q)\n L_residual = similar(Q)\n p = ntuple(i -> similar(Q), K)\n L_p = ntuple(i -> similar(Q), K)\n alpha = @MArray zeros(K)\n normsq = @MArray zeros(K)\n\n new{K, T, AT}(residual, L_residual, p, L_p, alpha, normsq, rtol, atol)\n end\nend\n\nfunction initialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::GeneralizedConjugateResidual,\n args...,\n)\n residual = solver.residual\n p = solver.p\n L_p = solver.L_p\n\n @assert size(Q) == size(residual)\n rtol, atol = solver.rtol, solver.atol\n\n threshold = rtol * norm(Qrhs, weighted_norm)\n linearoperator!(residual, Q, args...)\n residual .-= Qrhs\n\n converged = false\n residual_norm = norm(residual, weighted_norm)\n if residual_norm < threshold\n converged = true\n return converged, threshold\n end\n\n p[1] .= residual\n linearoperator!(L_p[1], p[1], args...)\n\n threshold = max(atol, threshold)\n\n converged, threshold\nend\n\nfunction doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::GeneralizedConjugateResidual{K},\n threshold,\n args...,\n) where {K}\n\n residual = solver.residual\n p = solver.p\n L_residual = solver.L_residual\n L_p = solver.L_p\n normsq = solver.normsq\n alpha = solver.alpha\n\n residual_norm = typemax(eltype(Q))\n for k in 1:K\n normsq[k] = norm(L_p[k], weighted_norm)^2\n beta = -dot(residual, L_p[k], weighted_norm) / normsq[k]\n\n Q .+= beta * p[k]\n residual .+= beta * L_p[k]\n\n residual_norm = norm(residual, weighted_norm)\n\n if residual_norm <= threshold\n return (true, k, residual_norm)\n end\n\n linearoperator!(L_residual, residual, args...)\n\n for l in 1:k\n alpha[l] = -dot(L_residual, L_p[l], weighted_norm) / normsq[l]\n end\n\n if k < K\n rv_nextp = realview(p[k + 1])\n rv_L_nextp = realview(L_p[k + 1])\n else # restart\n rv_nextp = realview(p[1])\n rv_L_nextp = realview(L_p[1])\n end\n\n rv_residual = realview(residual)\n rv_p = realview.(p)\n rv_L_p = realview.(L_p)\n rv_L_residual = realview(L_residual)\n\n groupsize = 256\n T = eltype(alpha)\n\n event = Event(array_device(Q))\n event = linearcombination!(array_device(Q), groupsize)(\n rv_nextp,\n (one(T), alpha[1:k]...),\n (rv_residual, rv_p[1:k]...),\n false;\n ndrange = length(rv_nextp),\n dependencies = (event,),\n )\n\n event = linearcombination!(array_device(Q), groupsize)(\n rv_L_nextp,\n (one(T), alpha[1:k]...),\n (rv_L_residual, rv_L_p[1:k]...),\n false;\n ndrange = length(rv_nextp),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n (false, K, residual_norm)\nend\n" }, { "path": "src/Numerics/SystemSolvers/generalized_minimal_residual_solver.jl", "content": "#### Generalized Minimal Residual Solver\n\nexport GeneralizedMinimalResidual\n\n\"\"\"\n GeneralizedMinimalResidual(Q; M, rtol, atol)\n\n# GMRES\nThis is an object for solving linear systems using an iterative Krylov method.\nThe constructor parameter `M` is the number of steps after which the algorithm\nis restarted (if it has not converged), `Q` is a reference state used only\nto allocate the solver internal state, and `rtol` specifies the convergence\ncriterion based on the relative residual norm. The amount of memory\nrequired for the solver state is roughly `(M + 1) * size(Q)`.\nThis object is intended to be passed to the [`linearsolve!`](@ref) command.\n\nThis uses the restarted Generalized Minimal Residual method of Saad and Schultz (1986).\n\n## References\n\n - [Saad1986](@cite)\n\n\"\"\"\nmutable struct GeneralizedMinimalResidual{M, MP1, MMP1, T, AT} <:\n AbstractIterativeSystemSolver\n krylov_basis::NTuple{MP1, AT}\n \"Hessenberg matrix\"\n H::MArray{Tuple{MP1, M}, T, 2, MMP1}\n \"rhs of the least squares problem\"\n g0::MArray{Tuple{MP1, 1}, T, 2, MP1}\n rtol::T\n atol::T\n\n function GeneralizedMinimalResidual(\n Q::AT;\n M = min(20, eltype(Q)),\n rtol = √eps(eltype(AT)),\n atol = eps(eltype(AT)),\n ) where {AT}\n krylov_basis = ntuple(i -> similar(Q), M + 1)\n H = @MArray zeros(M + 1, M)\n g0 = @MArray zeros(M + 1)\n\n new{M, M + 1, M * (M + 1), eltype(Q), AT}(\n krylov_basis,\n H,\n g0,\n rtol,\n atol,\n )\n end\nend\n\nfunction initialize!(\n linearoperator!,\n Q,\n Qrhs,\n solver::GeneralizedMinimalResidual,\n args...,\n)\n g0 = solver.g0\n krylov_basis = solver.krylov_basis\n rtol, atol = solver.rtol, solver.atol\n\n @assert size(Q) == size(krylov_basis[1])\n\n # store the initial residual in krylov_basis[1]\n linearoperator!(krylov_basis[1], Q, args...)\n @. krylov_basis[1] = Qrhs - krylov_basis[1]\n\n threshold = rtol * norm(krylov_basis[1], weighted_norm)\n residual_norm = norm(krylov_basis[1], weighted_norm)\n\n converged = false\n # FIXME: Should only be true for threshold zero\n if threshold < atol\n converged = true\n return converged, threshold\n end\n\n fill!(g0, 0)\n g0[1] = residual_norm\n krylov_basis[1] ./= residual_norm\n\n converged, max(threshold, atol)\nend\n\nfunction doiteration!(\n linearoperator!,\n preconditioner,\n Q,\n Qrhs,\n solver::GeneralizedMinimalResidual{M},\n threshold,\n args...,\n) where {M}\n\n krylov_basis = solver.krylov_basis\n H = solver.H\n g0 = solver.g0\n\n converged = false\n residual_norm = typemax(eltype(Q))\n\n Ω = LinearAlgebra.Rotation{eltype(Q)}([])\n j = 1\n for outer j in 1:M\n\n # Arnoldi using the Modified Gram Schmidt orthonormalization\n linearoperator!(krylov_basis[j + 1], krylov_basis[j], args...)\n for i in 1:j\n H[i, j] = dot(krylov_basis[j + 1], krylov_basis[i], weighted_norm)\n @. krylov_basis[j + 1] -= H[i, j] * krylov_basis[i]\n end\n H[j + 1, j] = norm(krylov_basis[j + 1], weighted_norm)\n krylov_basis[j + 1] ./= H[j + 1, j]\n\n # apply the previous Givens rotations to the new column of H\n @views H[1:j, j:j] .= Ω * H[1:j, j:j]\n\n # compute a new Givens rotation to zero out H[j + 1, j]\n G, _ = givens(H, j, j + 1, j)\n\n # apply the new rotation to H and the rhs\n H .= G * H\n g0 .= G * g0\n\n # compose the new rotation with the others\n Ω = lmul!(G, Ω)\n\n residual_norm = abs(g0[j + 1])\n\n if residual_norm < threshold\n converged = true\n break\n end\n end\n\n # solve the triangular system\n y = SVector{j}(@views UpperTriangular(H[1:j, 1:j]) \\ g0[1:j])\n\n ## compose the solution\n rv_Q = realview(Q)\n rv_krylov_basis = realview.(krylov_basis)\n groupsize = 256\n event = Event(array_device(Q))\n event = linearcombination!(array_device(Q), groupsize)(\n rv_Q,\n y,\n rv_krylov_basis,\n true;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n\n # if not converged restart\n converged || initialize!(linearoperator!, Q, Qrhs, solver, args...)\n\n (converged, j, residual_norm)\nend\n" }, { "path": "src/Numerics/SystemSolvers/jacobian_free_newton_krylov_solver.jl", "content": "\nexport JacobianFreeNewtonKrylovSolver, JacobianAction\n\n\"\"\"\nmutable struct JacobianAction{FT, AT}\n rhs!\n ϵ::FT\n Q::AT\n Qdq::AT\n Fq::AT\n Fqdq::AT\nend\n\nSolve for Frhs = F(q), the Jacobian action is computed\n\n ∂F(Q) F(Q + eΔQ) - F(Q)\n ---- ΔQ ≈ -------------------\n ∂Q e\n\n\n...\n# Arguments \n- `rhs!` : nonlinear operator F(Q)\n- `ϵ::FT` : ϵ used for finite difference, e = e(Q, ϵ)\n- `Q::AT` : cache for Q\n- `Qdq::AT` : container for Q + ϵΔQ\n- `Fq::AT` : cache for F(Q)\n- `Fqdq::AT` : container for F(Q + ϵΔQ)\n...\n\"\"\"\nmutable struct JacobianAction{FT, AT}\n rhs!::Any\n ϵ::FT\n Q::AT\n Qdq::AT\n Fq::AT\n Fqdq::AT\nend\n\nfunction JacobianAction(rhs!, Q, ϵ)\n return JacobianAction(\n rhs!,\n ϵ,\n similar(Q),\n similar(Q),\n similar(Q),\n similar(Q),\n )\nend\n\n\"\"\"\nApproximates the action of the Jacobian of a nonlinear\nform on a vector `ΔQ` using the difference quotient:\n\n ∂F(Q) F(Q + e ΔQ) - F(Q)\nJΔQ = ---- ΔQ ≈ -------------------\n ∂Q e\n\n\nCompute JΔQ with cached Q and F(Q), and the direction dQ\n\"\"\"\nfunction (op::JacobianAction)(JΔQ, dQ, args...)\n rhs! = op.rhs!\n Q = op.Q\n Qdq = op.Qdq\n ϵ = op.ϵ\n Fq = op.Fq\n Fqdq = op.Fqdq\n\n FT = eltype(dQ)\n n = length(dQ)\n normdQ = norm(dQ, weighted_norm)\n\n β = √ϵ\n\n if normdQ > ϵ\n # for preconditioner reconstruction, it goes into this part\n # e depends only on the active freedoms in the preconditioner reconstruction\n factor = FT(1 / (n * normdQ))\n Qdq .= Q .* (abs.(dQ) .> 0)\n e = factor * β * norm(Qdq, 1, false) + β\n else\n factor = FT(1 / n)\n e = factor * β * norm(Q, 1, false) + β\n end\n\n\n\n Qdq .= Q .+ e .* dQ\n\n rhs!(Fqdq, Qdq, args...)\n\n JΔQ .= (Fqdq .- Fq) ./ e\n\nend\n\n\"\"\"\nupdate cached Q and F(Q) before each Newton iteration\n\"\"\"\nfunction update_Q!(op::JacobianAction, Q, args...)\n op.Q .= Q\n Fq = op.Fq\n\n op.rhs!(Fq, Q, args...)\nend\n\n\"\"\"\nSolve for Frhs = F(Q), by finite difference\n\n ∂F(Q) F(Q + eΔQ) - F(Q)\n ---- ΔQ ≈ -------------------\n ∂Q e\n\n Q^n+1 = Q^n - dF/dQ(Q^{n})⁻¹ (F(Q^n) - Frhs)\n\n set ΔQ = F(Q^n) - Frhs\n\"\"\"\nmutable struct JacobianFreeNewtonKrylovSolver{FT, AT} <: AbstractNonlinearSolver\n # small number used for finite difference\n ϵ::FT\n # tolerances for convergence\n tol::FT\n # Max number of Newton iterations\n M::Int\n # Linear solver for the Jacobian system\n linearsolver::Any\n # container for unknows ΔQ, which is updated for the linear solver\n ΔQ::AT\n # contrainer for F(Q)\n residual::AT\nend\n\n\"\"\"\nJacobianFreeNewtonKrylovSolver constructor\n\"\"\"\nfunction JacobianFreeNewtonKrylovSolver(\n Q,\n linearsolver;\n ϵ = 1.e-8,\n tol = 1.e-6,\n M = 30,\n)\n FT = eltype(Q)\n residual = similar(Q)\n ΔQ = similar(Q)\n return JacobianFreeNewtonKrylovSolver(\n FT(ϵ),\n FT(tol),\n M,\n linearsolver,\n ΔQ,\n residual,\n )\nend\n\n\"\"\"\nJacobianFreeNewtonKrylovSolver initialize the residual\n\"\"\"\nfunction initialize!(\n rhs!,\n Q,\n Qrhs,\n solver::JacobianFreeNewtonKrylovSolver,\n args...,\n)\n # where R = Qrhs - F(Q)\n R = solver.residual\n # Computes F(Q) and stores in R\n rhs!(R, Q, args...)\n # Computes R = R - Qrhs\n R .-= Qrhs\n return norm(R, weighted_norm)\nend\n\n\"\"\"\nSolve for Frhs = F(Q), by finite difference\n\nQ^n+1 = Q^n - dF/dQ(Q^{n})⁻¹ (F(Q^n) - Frhs)\n\nset ΔQ = F(Q^n) - Frhs\n\n...\n# Arguments \n- `rhs!`: functor rhs!(Q) = F(Q)\n- `jvp!`: Jacobian action jvp!(ΔQ) = dF/dQ(Q) ⋅ ΔQ\n- `preconditioner`: approximation of dF/dQ(Q)\n- `Q` : Q^n\n- `Qrhs` : Frhs\n- `solver`: linear solver\n...\n\"\"\"\nfunction donewtoniteration!(\n rhs!,\n jvp!,\n preconditioner::AbstractPreconditioner,\n Q,\n Qrhs,\n solver::JacobianFreeNewtonKrylovSolver,\n args...,\n)\n\n FT = eltype(Q)\n ΔQ = solver.ΔQ\n ΔQ .= FT(0.0)\n\n # R(Q) == 0, R = F(Q) - Qrhs, where F = rhs!\n # Compute right-hand side for Jacobian system:\n # J(Q)ΔQ = -R\n # where R = Qrhs - F(Q), which is computed at the end of last step or in the initialize function\n R = solver.residual\n\n # R = F(Q^n) - Frhs\n # ΔQ = dF/dQ(Q^{n})⁻¹ (Frhs - F(Q^n)) = -dF/dQ(Q^{n})⁻¹ R\n iters =\n linearsolve!(jvp!, preconditioner, solver.linearsolver, ΔQ, -R, args...)\n\n # Newton correction Q^{n+1} = Q^n + dF/dQ(Q^{n})⁻¹ (Frhs - F(Q^n))\n Q .+= ΔQ\n\n # Compute residual norm and residual for next step\n rhs!(R, Q, args...)\n R .-= Qrhs\n resnorm = norm(R, weighted_norm)\n\n return resnorm, iters\nend\n" }, { "path": "src/Numerics/SystemSolvers/preconditioners.jl", "content": "export AbstractPreconditioner,\n ColumnwiseLUPreconditioner,\n NoPreconditioner,\n preconditioner_update!,\n preconditioner_solve!,\n preconditioner_counter_update!\n\n\"\"\"\n Abstract base type for all preconditioners.\n\"\"\"\nabstract type AbstractPreconditioner end\n\n\"\"\"\nmutable struct NoPreconditioner\nend\n\nDo nothing\n\"\"\"\nmutable struct NoPreconditioner <: AbstractPreconditioner end\n\n\"\"\"\nDo nothing, when there is no preconditioner, preconditioner = Nothing\n\"\"\"\nfunction preconditioner_update!(\n op,\n dg,\n preconditioner::NoPreconditioner,\n args...,\n) end\n\n\"\"\"\nDo nothing, when there is no preconditioner, preconditioner = Nothing\n\"\"\"\nfunction preconditioner_solve!(preconditioner::NoPreconditioner, Q) end\n\n\"\"\"\nDo nothing, when there is no preconditioner, preconditioner = Nothing\n\"\"\"\nfunction preconditioner_counter_update!(preconditioner::NoPreconditioner) end\n\n\"\"\"\nmutable struct ColumnwiseLUPreconditioner{AT}\n A::DGColumnBandedMatrix\n Q::AT\n PQ::AT\n counter::Int\n update_freq::Int\nend\n\n...\n# Arguments\n- `A`: the lu factor of the precondition (approximated Jacobian), in the DGColumnBandedMatrix format\n- `Q`: MPIArray container, used to update A\n- `PQ`: MPIArray container, used to update A\n- `counter`: count the number of Newton, when counter > update_freq or counter < 0, update precondition\n- `update_freq`: preconditioner update frequency\n...\n\"\"\"\nmutable struct ColumnwiseLUPreconditioner{AT} <: AbstractPreconditioner\n A::DGColumnBandedMatrix\n Q::AT\n PQ::AT\n counter::Int\n update_freq::Int\nend\n\n\"\"\"\nColumnwiseLUPreconditioner constructor\nbuild an empty ColumnwiseLUPreconditioner\n\n...\n# Arguments\n- `dg`: DG model, use only the grid information\n- `Q0`: MPIArray, use only its size\n- `counter`: = -1, which indicates the preconditioner is empty\n- `update_freq`: preconditioner update frequency\n...\n\"\"\"\nfunction ColumnwiseLUPreconditioner(dg, Q0, update_freq = 100)\n single_column = false\n Q = similar(Q0)\n PQ = similar(Q0)\n\n A = empty_banded_matrix(dg, Q; single_column = single_column)\n\n # counter = -1, which indicates the preconditioner is empty\n ColumnwiseLUPreconditioner(A, Q, PQ, -1, update_freq)\nend\n\n\"\"\"\nupdate the DGColumnBandedMatrix by the finite difference approximation\n...\n# Arguments\n- `op`: operator used to compute the finte difference information\n- `dg`: the DG model, use only the grid information\n...\n\"\"\"\nfunction preconditioner_update!(\n op,\n dg,\n preconditioner::ColumnwiseLUPreconditioner,\n args...,\n)\n\n # preconditioner.counter < 0, means newly constructed empty preconditioner\n if preconditioner.counter >= 0 &&\n (preconditioner.counter < preconditioner.update_freq)\n return\n end\n\n A = preconditioner.A\n Q = preconditioner.Q\n PQ = preconditioner.PQ\n\n update_banded_matrix!(A, op, dg, Q, PQ, args...)\n band_lu!(A)\n\n preconditioner.counter = 0\nend\n\n\"\"\"\nInplace applying the preconditioner\n\nQ = P⁻¹ * Q\n\"\"\"\nfunction preconditioner_solve!(preconditioner::ColumnwiseLUPreconditioner, Q)\n A = preconditioner.A\n band_forward!(Q, A)\n band_back!(Q, A)\n\nend\n\n\"\"\"\nUpdate the preconditioner counter, after each Newton iteration\n\"\"\"\nfunction preconditioner_counter_update!(\n preconditioner::ColumnwiseLUPreconditioner,\n)\n preconditioner.counter += 1\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/Courant.jl", "content": "using Logging, Printf\nusing LinearAlgebra: norm\n\nusing ...Mesh.Grids:\n VerticalDirection, HorizontalDirection, EveryDirection, min_node_distance\nusing ...DGMethods: courant\n\nimport ...Courant:\n advective_courant, nondiffusive_courant, diffusive_courant, viscous_courant\n\nimport ...DGMethods: calculate_dt\n\n\"\"\"\n advective_courant(::HBModel)\n\ncalculates the CFL condition due to advection\n\n\"\"\"\n@inline function advective_courant(\n m::HBModel,\n Q::Vars,\n A::Vars,\n D::Vars,\n Δx,\n Δt,\n t,\n direction = VerticalDirection(),\n)\n if direction isa VerticalDirection\n ū = norm(A.w)\n elseif direction isa HorizontalDirection\n ū = norm(Q.u)\n else\n v = @SVector [Q.u[1], Q.u[2], A.w]\n ū = norm(v)\n end\n\n return Δt * ū / Δx\nend\n\n\"\"\"\n nondiffusive_courant(::HBModel)\n\ncalculates the CFL condition due to gravity waves\n\n\"\"\"\n@inline function nondiffusive_courant(\n m::HBModel,\n Q::Vars,\n A::Vars,\n D::Vars,\n Δx,\n Δt,\n t,\n direction = HorizontalDirection(),\n)\n return Δt * m.cʰ / Δx\nend\n\"\"\"\n viscous_courant(::HBModel)\n\ncalculates the CFL condition due to viscosity\n\n\"\"\"\n@inline function viscous_courant(\n m::HBModel,\n Q::Vars,\n A::Vars,\n D::Vars,\n Δx,\n Δt,\n t,\n direction = VerticalDirection(),\n)\n ν̄ = norm_viscosity(m, direction)\n\n return Δt * ν̄ / Δx^2\nend\n\n@inline norm_viscosity(m::HBModel, ::VerticalDirection) = m.νᶻ\n@inline norm_viscosity(m::HBModel, ::HorizontalDirection) = sqrt(2) * m.νʰ\n@inline norm_viscosity(m::HBModel, ::EveryDirection) = sqrt(2 * m.νʰ^2 + m.νᶻ^2)\n\n\n\"\"\"\n diffusive_courant(::HBModel)\n\ncalculates the CFL condition due to temperature diffusivity\nfactor of 1000 is for convective adjustment\n\n\"\"\"\n@inline function diffusive_courant(\n m::HBModel,\n Q::Vars,\n A::Vars,\n D::Vars,\n Δx,\n Δt,\n t,\n direction = VerticalDirection(),\n)\n κ̄ = norm_diffusivity(m, direction)\n\n return Δt * κ̄ / Δx^2\nend\n\n@inline norm_diffusivity(m::HBModel, ::VerticalDirection) = 1000 * m.κᶻ\n@inline norm_diffusivity(m::HBModel, ::HorizontalDirection) = sqrt(2) * m.κʰ\n@inline norm_diffusivity(m::HBModel, ::EveryDirection) =\n sqrt(2 * m.κʰ^2 + (1000 * m.κᶻ)^2)\n\n\"\"\"\n calculate_dt(dg, model::HBModel, Q, Courant_number, direction::EveryDirection, t)\n\ncalculates the time step based on grid spacing and model parameters\ntakes minimum of advective, gravity wave, diffusive, and viscous CFL\n\n\"\"\"\n@inline function calculate_dt(\n dg,\n model::HBModel,\n Q,\n Courant_number,\n t,\n ::EveryDirection,\n)\n Δt = one(eltype(Q))\n\n CFL_advective =\n courant(advective_courant, dg, model, Q, Δt, t, VerticalDirection())\n CFL_gravity = courant(\n nondiffusive_courant,\n dg,\n model,\n Q,\n Δt,\n t,\n HorizontalDirection(),\n )\n CFL_viscous =\n courant(viscous_courant, dg, model, Q, Δt, t, VerticalDirection())\n CFL_diffusive =\n courant(diffusive_courant, dg, model, Q, Δt, t, VerticalDirection())\n\n CFLs = [CFL_advective, CFL_gravity, CFL_viscous, CFL_diffusive]\n dts = [Courant_number / CFL for CFL in CFLs]\n dt = min(dts...)\n\n @info @sprintf(\n \"\"\"Calculating timestep\n Advective Constraint = %.1f seconds\n Nondiffusive Constraint = %.1f seconds\n Viscous Constraint = %.1f seconds\n Diffusive Constrait = %.1f seconds\n Timestep = %.1f seconds\"\"\",\n dts...,\n dt\n )\n\n return dt\nend\n\n\n\"\"\"\n calculate_dt(dg, bl::LinearHBModel, Q, Courant_number,\n direction::EveryDirection)\n\ncalculates the time step based on grid spacing and model parameters\ntakes minimum of gravity wave, diffusive, and viscous CFL\n\n\"\"\"\n@inline function calculate_dt(\n dg,\n model::LinearHBModel,\n Q,\n Courant_number,\n t,\n ::EveryDirection,\n)\n Δt = one(eltype(Q))\n ocean = model.ocean\n\n CFL_advective =\n courant(advective_courant, dg, ocean, Q, Δt, t, VerticalDirection())\n CFL_gravity = courant(\n nondiffusive_courant,\n dg,\n ocean,\n Q,\n Δt,\n t,\n HorizontalDirection(),\n )\n CFL_viscous =\n courant(viscous_courant, dg, ocean, Q, Δt, t, HorizontalDirection())\n CFL_diffusive =\n courant(diffusive_courant, dg, ocean, Q, Δt, t, HorizontalDirection())\n\n CFLs = [CFL_advective, CFL_gravity, CFL_viscous, CFL_diffusive]\n dts = [Courant_number / CFL for CFL in CFLs]\n dt = min(dts...)\n\n @info @sprintf(\n \"\"\"Calculating timestep\n Advective Constraint = %.1f seconds\n Nondiffusive Constraint = %.1f seconds\n Viscous Constraint = %.1f seconds\n Diffusive Constrait = %.1f seconds\n Timestep = %.1f seconds\"\"\",\n dts...,\n dt\n )\n\n\n return dt\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/HydrostaticBoussinesq.jl", "content": "module HydrostaticBoussinesq\n\nexport HydrostaticBoussinesqModel, Forcing\n\nusing StaticArrays\nusing LinearAlgebra: dot, Diagonal\nusing CLIMAParameters.Planet: grav\n\nusing ..Ocean\nusing ...VariableTemplates\nusing ...MPIStateArrays\nusing ...Mesh.Filters: apply!\nusing ...Mesh.Grids: VerticalDirection\nusing ...Mesh.Geometry\nusing ...DGMethods\nusing ...DGMethods: init_state_auxiliary!\nusing ...DGMethods.NumericalFluxes\nusing ...DGMethods.NumericalFluxes: RusanovNumericalFlux\nusing ...BalanceLaws\n\nimport ..Ocean: coriolis_parameter\nimport ...DGMethods.NumericalFluxes: update_penalty!\n\nimport ...BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_state!,\n boundary_conditions,\n update_auxiliary_state!,\n update_auxiliary_state_gradient!,\n integral_load_auxiliary_state!,\n integral_set_auxiliary_state!,\n indefinite_stack_integral!,\n reverse_indefinite_stack_integral!,\n reverse_integral_load_auxiliary_state!,\n reverse_integral_set_auxiliary_state!\n\nimport ..Ocean:\n ocean_init_state!,\n ocean_init_aux!,\n ocean_boundary_state!,\n _ocean_boundary_state!\n\n×(a::SVector, b::SVector) = StaticArrays.cross(a, b)\n⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)\n⊗(a::SVector, b::SVector) = a * b'\n\ninclude(\"hydrostatic_boussinesq_model.jl\")\ninclude(\"bc_velocity.jl\")\ninclude(\"bc_temperature.jl\")\ninclude(\"LinearHBModel.jl\")\ninclude(\"Courant.jl\")\n\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/LinearHBModel.jl", "content": "export LinearHBModel\n\n# Linear model for 1D IMEX\n\"\"\"\n LinearHBModel <: BalanceLaw\n\nA `BalanceLaw` for modeling vertical diffusion implicitly.\n\nwrite out the equations here\n\n# Usage\n\n model = HydrostaticBoussinesqModel(problem)\n linear = LinearHBModel(model)\n\n\"\"\"\nstruct LinearHBModel{M} <: BalanceLaw\n ocean::M\n function LinearHBModel(ocean::M) where {M}\n return new{M}(ocean)\n end\nend\n\n\"\"\"\n Copy over state, aux, and diff variables from HBModel\n\"\"\"\nvars_state(lm::LinearHBModel, ::Prognostic, FT) =\n vars_state(lm.ocean, Prognostic(), FT)\nvars_state(lm::LinearHBModel, st::Gradient, FT) = vars_state(lm.ocean, st, FT)\nvars_state(lm::LinearHBModel, ::GradientFlux, FT) =\n vars_state(lm.ocean, GradientFlux(), FT)\nvars_state(lm::LinearHBModel, st::Auxiliary, FT) = vars_state(lm.ocean, st, FT)\nvars_state(lm::LinearHBModel, ::UpwardIntegrals, FT) = @vars()\n\n\"\"\"\n No integration, hyperbolic flux, or source terms\n\"\"\"\n@inline integrate_aux!(::LinearHBModel, _...) = nothing\n@inline flux_first_order!(::LinearHBModel, _...) = nothing\n@inline source!(::LinearHBModel, _...) = nothing\n\n\"\"\"\n No need to init, initialize by full model\n\"\"\"\ninit_state_auxiliary!(\n lm::LinearHBModel,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n) = nothing\ninit_state_prognostic!(lm::LinearHBModel, Q::Vars, A::Vars, localgeo, t) =\n nothing\n\n\"\"\"\n compute_gradient_argument!(::LinearHBModel)\n\ncopy u and θ to var_gradient\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `G`: array of gradient variables\n- `Q`: array of state variables\n- `A`: array of aux variables\n- `t`: time, not used\n\"\"\"\n@inline function compute_gradient_argument!(\n m::LinearHBModel,\n G::Vars,\n Q::Vars,\n A,\n t,\n)\n G.∇u = Q.u\n G.∇θ = Q.θ\n\n return nothing\nend\n\n\"\"\"\n compute_gradient_flux!(::LinearHBModel)\n\ncopy ν∇u and κ∇θ to var_diffusive\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `D`: array of diffusive variables\n- `G`: array of gradient variables\n- `Q`: array of state variables\n- `A`: array of aux variables\n- `t`: time, not used\n\"\"\"\n@inline function compute_gradient_flux!(\n lm::LinearHBModel,\n D::Vars,\n G::Grad,\n Q::Vars,\n A::Vars,\n t,\n)\n ν = viscosity_tensor(lm.ocean)\n D.ν∇u = -ν * G.∇u\n\n κ = diffusivity_tensor(lm.ocean, G.∇θ[3])\n D.κ∇θ = -κ * G.∇θ\n\n return nothing\nend\n\n\"\"\"\n flux_second_order!(::HBModel)\n\ncalculates the parabolic flux contribution to state variables\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `F`: array of fluxes for each state variable\n- `Q`: array of state variables\n- `D`: array of diff variables\n- `A`: array of aux variables\n- `t`: time, not used\n\n# computations\n∂ᵗu = -∇⋅(ν∇u)\n∂ᵗθ = -∇⋅(κ∇θ)\n\"\"\"\n@inline function flux_second_order!(\n lm::LinearHBModel,\n F::Grad,\n Q::Vars,\n D::Vars,\n HD::Vars,\n A::Vars,\n t::Real,\n)\n F.u += D.ν∇u\n F.θ += D.κ∇θ\n\n return nothing\nend\n\n\"\"\"\n wavespeed(::LinaerHBModel)\n\ncalculates the wavespeed for rusanov flux\n\"\"\"\nfunction wavespeed(lm::LinearHBModel, n⁻, _...)\n C = abs(SVector(lm.ocean.cʰ, lm.ocean.cʰ, lm.ocean.cᶻ)' * n⁻)\n return C\nend\n\nboundary_conditions(linear::LinearHBModel) = boundary_conditions(linear.ocean)\n\n\"\"\"\n boundary_state!(nf, ::LinearHBModel, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@inline function boundary_state!(nf, bc, linear::LinearHBModel, args...)\n return _ocean_boundary_state!(nf, bc, linear.ocean, args...)\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/bc_temperature.jl", "content": "using ..Ocean: surface_flux\n\n\"\"\"\n ocean_temperature_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Insulating, ::HBModel)\n\napply insulating boundary condition for temperature\nsets transmissive ghost point\n\"\"\"\nfunction ocean_temperature_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc_temperature::Insulating,\n ocean,\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n t,\n)\n Q⁺.θ = Q⁻.θ\n\n return nothing\nend\n\n\"\"\"\n ocean_temperature_boundary_state!(::NumericalFluxSecondOrder, ::Insulating, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for κ∇θ\n\"\"\"\n@inline function ocean_temperature_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_temperature::Insulating,\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.θ = Q⁻.θ\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n ocean_temperature_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::TemperatureFlux, ::HBModel)\n\napply temperature flux boundary condition for velocity\napplies insulating conditions for first-order and gradient fluxes\n\"\"\"\nfunction ocean_temperature_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc_velocity::TemperatureFlux,\n ocean,\n args...,\n)\n return ocean_temperature_boundary_state!(nf, Insulating(), ocean, args...)\nend\n\n\"\"\"\n ocean_temperature_boundary_state!(::NumericalFluxSecondOrder, ::TemperatureFlux, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have specified flux on the boundary for κ∇θ\n\"\"\"\n@inline function ocean_temperature_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_temperature::TemperatureFlux,\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.θ = Q⁻.θ\n D⁺.κ∇θ = n⁻ * surface_flux(ocean.problem, A⁻.y, Q⁻.θ)\n\n return nothing\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/bc_velocity.jl", "content": "using ..Ocean: kinematic_stress\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::HBModel)\n\napply no slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxFirstOrder,\n bc_velocity::Impenetrable{NoSlip},\n ocean,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n Q⁺.u = -Q⁻.u\n A⁺.w = -A⁻.w\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::HBModel)\n\napply no slip boundary condition for velocity\nset numerical flux to zero for u\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxGradient,\n bc_velocity::Impenetrable{NoSlip},\n ocean,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n FT = eltype(Q⁺)\n Q⁺.u = SVector(-zero(FT), -zero(FT))\n A⁺.w = -zero(FT)\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::HBModel)\n\napply no slip boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for u\n\"\"\"\n@inline function ocean_velocity_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_velocity::Impenetrable{NoSlip},\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = -Q⁻.u\n A⁺.w = -A⁻.w\n D⁺.ν∇u = D⁻.ν∇u\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::HBModel)\n\napply free slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxFirstOrder,\n bc_velocity::Impenetrable{FreeSlip},\n ocean,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n v⁻ = @SVector [Q⁻.u[1], Q⁻.u[2], A⁻.w]\n v⁺ = v⁻ - 2 * n⁻ ⋅ v⁻ .* SVector(n⁻)\n Q⁺.u = @SVector [v⁺[1], v⁺[2]]\n A⁺.w = v⁺[3]\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::HBModel)\n\napply free slip boundary condition for velocity\nsets non-reflective ghost point\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxGradient,\n bc_velocity::Impenetrable{FreeSlip},\n ocean,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n v⁻ = @SVector [Q⁻.u[1], Q⁻.u[2], A⁻.w]\n v⁺ = v⁻ - n⁻ ⋅ v⁻ .* SVector(n⁻)\n Q⁺.u = @SVector [v⁺[1], v⁺[2]]\n A⁺.w = v⁺[3]\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::HBModel)\n\napply free slip boundary condition for velocity\napply zero numerical flux in the normal direction\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_velocity::Impenetrable{FreeSlip},\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = Q⁻.u\n A⁺.w = A⁻.w\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::HBModel)\n\napply free slip boundary condition for velocity\nsets non-reflective ghost point\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc_velocity::Penetrable{FreeSlip},\n ocean,\n args...,\n)\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::HBModel)\n\napply free slip boundary condition for velocity\nsets non-reflective ghost point\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_velocity::Penetrable{FreeSlip},\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = Q⁻.u\n A⁺.w = A⁻.w\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc_velocity::Impenetrable{<:KinematicStress},\n ocean,\n args...,\n)\n return ocean_velocity_boundary_state!(\n nf,\n Impenetrable(FreeSlip()),\n ocean,\n args...,\n )\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function ocean_velocity_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_velocity::Impenetrable{<:KinematicStress},\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = Q⁻.u\n D⁺.ν∇u =\n n⁻ * kinematic_stress(ocean.problem, A⁻.y, ocean.ρₒ, bc_velocity.drag)'\n\n return nothing\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction ocean_velocity_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc_velocity::Penetrable{<:KinematicStress},\n ocean,\n args...,\n)\n return ocean_velocity_boundary_state!(\n nf,\n Penetrable(FreeSlip()),\n ocean,\n args...,\n )\nend\n\n\"\"\"\n ocean_velocity_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function ocean_velocity_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc_velocity::Penetrable{<:KinematicStress},\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = Q⁻.u\n D⁺.ν∇u =\n n⁻ *\n kinematic_stress(\n ocean.problem,\n A⁻.y,\n ocean.ρₒ,\n bc_velocity.drag.stress,\n )'\n\n return nothing\nend\n" }, { "path": "src/Ocean/HydrostaticBoussinesq/hydrostatic_boussinesq_model.jl", "content": "\"\"\"\n HydrostaticBoussinesqModel <: BalanceLaw\n\nA `BalanceLaw` for ocean modeling.\n\nwrite out the equations here\n\nρₒ = reference density of sea water\ncʰ = maximum horizontal wave speed\ncᶻ = maximum vertical wave speed\nαᵀ = thermal expansitivity coefficient\nνʰ = horizontal viscosity\nνᶻ = vertical viscosity\nκʰ = horizontal diffusivity\nκᶻ = vertical diffusivity\nfₒ = first coriolis parameter (constant term)\nβ = second coriolis parameter (linear term)\n\n# Usage\n\n HydrostaticBoussinesqModel(problem)\n\n\"\"\"\nstruct HydrostaticBoussinesqModel{C, PS, P, MA, TA, F, FT, I} <: BalanceLaw\n param_set::PS\n problem::P\n coupling::C\n momentum_advection::MA\n tracer_advection::TA\n forcing::F\n state_filter::I\n ρₒ::FT\n cʰ::FT\n cᶻ::FT\n αᵀ::FT\n νʰ::FT\n νᶻ::FT\n κʰ::FT\n κᶻ::FT\n κᶜ::FT\n fₒ::FT\n β::FT\n function HydrostaticBoussinesqModel{FT}(\n param_set::PS,\n problem::P;\n coupling::C = Uncoupled(),\n momentum_advection::MA = nothing,\n tracer_advection::TA = NonLinearAdvectionTerm(),\n forcing::F = Forcing(),\n state_filter::I = nothing,\n ρₒ = FT(1000), # kg / m^3\n cʰ = FT(0), # m/s\n cᶻ = FT(0), # m/s\n αᵀ = FT(2e-4), # (m/s)^2 / K\n νʰ = FT(5e3), # m^2 / s\n νᶻ = FT(5e-3), # m^2 / s\n κʰ = FT(1e3), # m^2 / s # horizontal diffusivity\n κᶻ = FT(1e-4), # m^2 / s # background vertical diffusivity\n κᶜ = FT(1e-1), # m^2 / s # diffusivity for convective adjustment\n fₒ = FT(1e-4), # Hz\n β = FT(1e-11), # Hz / m\n ) where {FT <: AbstractFloat, PS, P, C, MA, TA, F, I}\n return new{C, PS, P, MA, TA, F, FT, I}(\n param_set,\n problem,\n coupling,\n momentum_advection,\n tracer_advection,\n forcing,\n state_filter,\n ρₒ,\n cʰ,\n cᶻ,\n αᵀ,\n νʰ,\n νᶻ,\n κʰ,\n κᶻ,\n κᶜ,\n fₒ,\n β,\n )\n end\nend\n\nHBModel = HydrostaticBoussinesqModel\n\nboundary_conditions(ocean::HBModel) = ocean.problem.boundary_conditions\n\n@inline noforcing(args...) = 0\n\nfunction Forcing(; u = noforcing, v = noforcing, η = noforcing, θ = noforcing)\n return (u = u, v = v, η = η, θ = θ)\nend\n\n\"\"\"\n vars_state(::HBModel, ::Prognostic)\n\nprognostic variables evolved forward in time\n\nu = (u,v) = (zonal velocity, meridional velocity)\nη = sea surface height\nθ = temperature\n\"\"\"\nfunction vars_state(m::HBModel, ::Prognostic, T)\n @vars begin\n u::SVector{2, T}\n η::T # real a 2-D variable TODO: should be 2D\n θ::T\n end\nend\n\n\"\"\"\n init_state_prognostic!(::HBModel)\n\nsets the initial value for state variables\ndispatches to ocean_init_state! which is defined in a problem file such as SimpleBoxProblem.jl\n\"\"\"\nfunction init_state_prognostic!(m::HBModel, Q::Vars, A::Vars, local_geometry, t)\n return ocean_init_state!(m, m.problem, Q, A, local_geometry, t)\nend\n\n\"\"\"\n vars_state(::HBModel, ::Auxiliary)\nhelper variables for computation\n\nsecond half is because there is no dedicated integral kernels\nthese variables are used to compute vertical integrals\nw = vertical velocity\nwz0 = w at z = 0\npkin = bulk hydrostatic pressure contribution\n\n\nfirst half of these are fields that are used for computation\ny = north-south coordinate\n\n\"\"\"\nfunction vars_state(m::HBModel, ::Auxiliary, T)\n @vars begin\n y::T # y-coordinate of the box\n w::T # ∫(-∇⋅u)\n pkin::T # ∫(-αᵀθ)\n wz0::T # w at z=0\n uᵈ::SVector{2, T} # velocity deviation from vertical mean\n ΔGᵘ::SVector{2, T} # vertically averaged tendency\n end\nend\n\nfunction ocean_init_aux! end\n\n\"\"\"\n init_state_auxiliary!(::HBModel)\n\nsets the initial value for auxiliary variables (those that aren't related to vertical integrals)\ndispatches to ocean_init_aux! which is defined in a problem file such as SimpleBoxProblem.jl\n\"\"\"\nfunction init_state_auxiliary!(\n m::HBModel,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n m,\n (m, A, tmp, geom) -> ocean_init_aux!(m, m.problem, A, geom),\n state_auxiliary,\n grid,\n direction,\n )\nend\n\n\"\"\"\n vars_state(::HBModel, ::Gradient)\n\nvariables that you want to take a gradient of\nthese are just copies in our model\n\"\"\"\nfunction vars_state(m::HBModel, ::Gradient, T)\n @vars begin\n ∇u::SVector{2, T}\n ∇uᵈ::SVector{2, T}\n ∇θ::T\n end\nend\n\n\"\"\"\n compute_gradient_argument!(::HBModel)\n\ncopy u and θ to var_gradient\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `G`: array of gradient variables\n- `Q`: array of state variables\n- `A`: array of aux variables\n- `t`: time, not used\n\"\"\"\n@inline function compute_gradient_argument!(m::HBModel, G::Vars, Q::Vars, A, t)\n G.∇θ = Q.θ\n\n velocity_gradient_argument!(m, m.coupling, G, Q, A, t)\n\n return nothing\nend\n\n@inline function velocity_gradient_argument!(\n m::HBModel,\n ::Uncoupled,\n G,\n Q,\n A,\n t,\n)\n G.∇u = Q.u\n\n return nothing\nend\n\n\"\"\"\n vars_state(::HBModel, ::GradientFlux, FT)\n\nthe output of the gradient computations\nmultiplies ∇u by viscosity tensor and ∇θ by the diffusivity tensor\n\"\"\"\nfunction vars_state(m::HBModel, ::GradientFlux, T)\n @vars begin\n ∇ʰu::T\n ν∇u::SMatrix{3, 2, T, 6}\n κ∇θ::SVector{3, T}\n end\nend\n\n\"\"\"\n compute_gradient_flux!(::HBModel)\n\ncopy ∇u and ∇θ to var_diffusive\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `D`: array of diffusive variables\n- `G`: array of gradient variables\n- `Q`: array of state variables\n- `A`: array of aux variables\n- `t`: time, not used\n\"\"\"\n@inline function compute_gradient_flux!(\n m::HBModel,\n D::Vars,\n G::Grad,\n Q::Vars,\n A::Vars,\n t,\n)\n # store ∇ʰu for continuity equation (convert gradient to divergence)\n D.∇ʰu = G.∇u[1, 1] + G.∇u[2, 2]\n\n velocity_gradient_flux!(m, m.coupling, D, G, Q, A, t)\n\n κ = diffusivity_tensor(m, G.∇θ[3])\n D.κ∇θ = -κ * G.∇θ\n\n return nothing\nend\n\n@inline function velocity_gradient_flux!(m::HBModel, ::Uncoupled, D, G, Q, A, t)\n ν = viscosity_tensor(m)\n D.ν∇u = -ν * G.∇u\n\n return nothing\nend\n\n\"\"\"\n viscosity_tensor(::HBModel)\n\nuniform viscosity with different values for horizontal and vertical directions\n\n# Arguments\n- `m`: model object to dispatch on and get viscosity parameters\n\"\"\"\n@inline viscosity_tensor(m::HBModel) = Diagonal(@SVector [m.νʰ, m.νʰ, m.νᶻ])\n\n\"\"\"\n diffusivity_tensor(::HBModel)\n\nuniform diffusivity in the horizontal direction\napplies convective adjustment in the vertical, bump by 1000 if ∂θ∂z < 0\n\n# Arguments\n- `m`: model object to dispatch on and get diffusivity parameters\n- `∂θ∂z`: value of the derivative of temperature in the z-direction\n\"\"\"\n@inline function diffusivity_tensor(m::HBModel, ∂θ∂z)\n ∂θ∂z < 0 ? κ = m.κᶜ : κ = m.κᶻ\n\n return Diagonal(@SVector [m.κʰ, m.κʰ, κ])\nend\n\n\"\"\"\n vars_integral(::HBModel)\n\nlocation to store integrands for bottom up integrals\n∇hu = the horizontal divegence of u, e.g. dw/dz\n\"\"\"\nfunction vars_state(m::HBModel, ::UpwardIntegrals, T)\n @vars begin\n ∇ʰu::T\n αᵀθ::T\n end\nend\n\n\"\"\"\n integral_load_auxiliary_state!(::HBModel)\n\ncopy w to var_integral\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case HBModel\nI -> array of integrand variables\nQ -> array of state variables\nA -> array of aux variables\n\"\"\"\n@inline function integral_load_auxiliary_state!(\n m::HBModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.∇ʰu = A.w # borrow the w value from A...\n I.αᵀθ = -m.αᵀ * Q.θ # integral will be reversed below\n\n return nothing\nend\n\n\"\"\"\n integral_set_auxiliary_state!(::HBModel)\n\ncopy integral results back out to aux\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case HBModel\nA -> array of aux variables\nI -> array of integrand variables\n\"\"\"\n@inline function integral_set_auxiliary_state!(m::HBModel, A::Vars, I::Vars)\n A.w = I.∇ʰu\n A.pkin = I.αᵀθ\n\n return nothing\nend\n\n\"\"\"\n vars_reverse_integral(::HBModel)\n\nlocation to store integrands for top down integrals\nαᵀθ = density perturbation\n\"\"\"\nfunction vars_state(m::HBModel, ::DownwardIntegrals, T)\n @vars begin\n αᵀθ::T\n end\nend\n\n\"\"\"\n reverse_integral_load_auxiliary_state!(::HBModel)\n\ncopy αᵀθ to var_reverse_integral\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case HBModel\nI -> array of integrand variables\nA -> array of aux variables\n\"\"\"\n@inline function reverse_integral_load_auxiliary_state!(\n m::HBModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.αᵀθ = A.pkin\n\n return nothing\nend\n\n\"\"\"\n reverse_integral_set_auxiliary_state!(::HBModel)\n\ncopy reverse integral results back out to aux\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case HBModel\nA -> array of aux variables\nI -> array of integrand variables\n\"\"\"\n@inline function reverse_integral_set_auxiliary_state!(\n m::HBModel,\n A::Vars,\n I::Vars,\n)\n A.pkin = I.αᵀθ\n\n return nothing\nend\n\n\"\"\"\n flux_first_order!(::HBModel)\n\ncalculates the hyperbolic flux contribution to state variables\nthis computation is done pointwise at each nodal point\n\n# arguments:\nm -> model in this case HBModel\nF -> array of fluxes for each state variable\nQ -> array of state variables\nA -> array of aux variables\nt -> time, not used\n\n# computations\n∂ᵗu = ∇⋅(g*η + g∫αᵀθdz + v⋅u)\n∂ᵗθ = ∇⋅(vθ) where v = (u,v,w)\n\"\"\"\n@inline function flux_first_order!(\n m::HBModel,\n F::Grad,\n Q::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n # ∇ʰ • (g η)\n hydrostatic_pressure!(m, m.coupling, F, Q, A, t)\n\n # ∇ʰ • (- ∫(αᵀ θ))\n kinematic_pressure!(m, F, Q, A, t)\n\n # ∇ʰ • (v ⊗ u)\n momentum_advection!(m, m.momentum_advection, F, Q, A, t)\n\n # ∇ • (u θ)\n tracer_advection!(m, m.tracer_advection, F, Q, A, t)\n\n return nothing\nend\n\n@inline function hydrostatic_pressure!(m::HBModel, ::Uncoupled, F, Q, A, t)\n η = Q.η\n Iʰ = @SMatrix [\n 1 -0\n -0 1\n -0 -0\n ]\n\n F.u += grav(parameter_set(m)) * η * Iʰ\n\n return nothing\nend\n\n@inline function kinematic_pressure!(m::HBModel, F, Q, A, t)\n pkin = A.pkin\n Iʰ = @SMatrix [\n 1 -0\n -0 1\n -0 -0\n ]\n F.u += grav(parameter_set(m)) * pkin * Iʰ\n\n return nothing\nend\n\nmomentum_advection!(::HBModel, ::Nothing, _...) = nothing\n@inline function momentum_advection!(\n ::HBModel,\n ::NonLinearAdvectionTerm,\n F,\n Q,\n A,\n t,\n)\n u = Q.u\n @inbounds v = @SVector [Q.u[1], Q.u[2], A.w]\n\n F.u += v * u'\n\n return nothing\nend\n\ntracer_advection!(::HBModel, ::Nothing, _...) = nothing\n@inline function tracer_advection!(\n ::HBModel,\n ::NonLinearAdvectionTerm,\n F,\n Q,\n A,\n t,\n)\n θ = Q.θ\n @inbounds v = @SVector [Q.u[1], Q.u[2], A.w]\n\n F.θ += v * θ\n\n return nothing\nend\n\n\"\"\"\n flux_second_order!(::HBModel)\n\ncalculates the parabolic flux contribution to state variables\nthis computation is done pointwise at each nodal point\n\n# arguments:\n- `m`: model in this case HBModel\n- `F`: array of fluxes for each state variable\n- `Q`: array of state variables\n- `D`: array of diff variables\n- `A`: array of aux variables\n- `t`: time, not used\n\n# computations\n∂ᵗu = -∇⋅(ν∇u)\n∂ᵗθ = -∇⋅(κ∇θ)\n\"\"\"\n@inline function flux_second_order!(\n m::HBModel,\n F::Grad,\n Q::Vars,\n D::Vars,\n HD::Vars,\n A::Vars,\n t::Real,\n)\n F.u += D.ν∇u\n F.θ += D.κ∇θ\n\n return nothing\nend\n\n\"\"\"\n source!(::HBModel)\n\nCalculates the source term contribution to state variables.\nThis computation is done pointwise at each nodal point.\n\nArguments:\n m -> model in this case HBModel\n F -> array of fluxes for each state variable\n Q -> array of state variables\n A -> array of aux variables\n t -> time, not used\n\nComputations:\n ∂ᵗu = -f × u\n ∂ᵗη = w|(z=0)\n\"\"\"\n@inline function source!(\n m::HBModel,\n S::Vars,\n Q::Vars,\n D::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n # explicit forcing for SSH\n wz0 = A.wz0\n S.η += wz0\n\n coriolis_force!(m, m.coupling, S, Q, A, t)\n\n # Arguments for forcing functions\n # args = y, t, u, v, w, η, θ\n args = tuple(A.y, t, Q.u[1], Q.u[2], A.w, Q.η, Q.θ)\n\n Su = m.forcing.u(args...)\n Sv = m.forcing.v(args...)\n\n S.u += @SVector [Su, Sv]\n S.η += m.forcing.η(args...)\n S.θ += m.forcing.θ(args...)\n\n return nothing\nend\n\n@inline function coriolis_force!(m::HBModel, ::Uncoupled, S, Q, A, t)\n # f × u\n f = coriolis_parameter(m, m.problem, A.y)\n u, v = Q.u # Horizontal components of velocity\n S.u -= @SVector [-f * v, f * u]\n\n return nothing\nend\n\n\"\"\"\n wavespeed(::HBModel)\n\ncalculates the wavespeed for rusanov flux\n\"\"\"\n@inline wavespeed(m::HBModel, n⁻, _...) = abs(SVector(m.cʰ, m.cʰ, m.cᶻ)' * n⁻)\n\n\"\"\"\n update_penalty(::HBModel)\n set Δη = 0 when computing numerical fluxes\n\"\"\"\n# We want not have jump penalties on η (since not a flux variable)\nfunction update_penalty!(\n ::RusanovNumericalFlux,\n ::HBModel,\n n⁻,\n λ,\n ΔQ::Vars,\n Q⁻,\n A⁻,\n Q⁺,\n A⁺,\n t,\n)\n ΔQ.η = -0\n\n return nothing\nend\n\nfilter_state!(Q, filter::Nothing, grid) = nothing\nfilter_state!(Q, filter, grid) = apply!(Q, UnitRange(1, size(Q, 2)), grid)\n\n\"\"\"\n update_auxiliary_state!(::HBModel)\n\nApplies the vertical filter to the zonal and meridional velocities to preserve numerical incompressibility\nApplies an exponential filter to θ to anti-alias the non-linear advective term\n\nDoesn't actually touch the aux variables any more, but we need a better filter interface than this anyways\n\"\"\"\nfunction update_auxiliary_state!(\n dg::DGModel,\n m::HBModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n FT = eltype(Q)\n MD = dg.modeldata\n\n # `update_aux!` gets called twice, once for the real elements and once for\n # the ghost elements. Only apply the filters to the real elems.\n if elems == dg.grid.topology.realelems\n # required to ensure that after integration velocity field is divergence free\n vert_filter = MD.vert_filter\n apply!(Q, (:u,), dg.grid, vert_filter, direction = VerticalDirection())\n\n exp_filter = MD.exp_filter\n apply!(Q, (:θ,), dg.grid, exp_filter, direction = VerticalDirection())\n\n filter_state!(Q, m.state_filter, dg.grid)\n end\n\n compute_flow_deviation!(dg, m, m.coupling, Q, t)\n\n return true\nend\n\n@inline compute_flow_deviation!(dg, ::HBModel, ::Uncoupled, _...) = nothing\n\n\n\"\"\"\n update_auxiliary_state_gradient!(::HBModel)\n\n ∇hu to w for integration\n performs integration for w and pkin (should be moved to its own integral kernels)\n copies down w and wz0 because we don't have 2D structures\n\"\"\"\nfunction update_auxiliary_state_gradient!(\n dg::DGModel,\n m::HBModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n FT = eltype(Q)\n A = dg.state_auxiliary\n D = dg.state_gradient_flux\n\n # load -∇ʰu as ∂ᶻw\n index_w = varsindex(vars_state(m, Auxiliary(), FT), :w)\n index_∇ʰu = varsindex(vars_state(m, GradientFlux(), FT), :∇ʰu)\n @views @. A.data[:, index_w, elems] = -D.data[:, index_∇ʰu, elems]\n\n # compute integrals for w and pkin\n indefinite_stack_integral!(dg, m, Q, A, t, elems) # bottom -> top\n reverse_indefinite_stack_integral!(dg, m, Q, A, t, elems) # top -> bottom\n\n # We are unable to use vars (ie A.w) for this because this operation will\n # return a SubArray, and adapt (used for broadcasting along reshaped arrays)\n # has a limited recursion depth for the types allowed.\n number_aux = number_states(m, Auxiliary())\n index_wz0 = varsindex(vars_state(m, Auxiliary(), FT), :wz0)\n info = basic_grid_info(dg)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n # project w(z=0) down the stack\n data = reshape(A.data, Nqh, Nqk, number_aux, nelemv, nelemh)\n flat_wz0 = @view data[:, end:end, index_w, end:end, 1:nrealelemh]\n boxy_wz0 = @view data[:, :, index_wz0, :, 1:nrealelemh]\n boxy_wz0 .= flat_wz0\n\n return true\nend\n\n\"\"\"\n boundary_state!(nf, bc, ::HBModel, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@inline boundary_state!(nf, bc, ocean::HBModel, args...) =\n _ocean_boundary_state!(nf, bc, ocean, args...)\n\n#=\n\"\"\"\n boundary_state!(nf, ::HBModel, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@inline function boundary_state!(nf, ocean::HBModel, args...)\n boundary_conditions = ocean.problem.boundary_conditions\n return ocean_boundary_state!(nf, boundary_conditions, ocean, args...)\nend\n=#\n\n\"\"\"\n ocean_boundary_state!(nf, bc::OceanBC, ::HBModel, args...)\n\nsplits boundary condition application into velocity and temperature conditions\n\"\"\"\n@inline function ocean_boundary_state!(nf, bc::OceanBC, ocean::HBModel, args...)\n ocean_velocity_boundary_state!(nf, bc.velocity, ocean, args...)\n ocean_temperature_boundary_state!(nf, bc.temperature, ocean, args...)\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(nf, boundaries::Tuple, ::HBModel,\n Q⁺, A⁺, n, Q⁻, A⁻, bctype)\n\napplies boundary conditions for the first-order and gradient fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@generated function ocean_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n boundaries::Tuple,\n ocean,\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n bctype,\n t,\n args...,\n)\n N = fieldcount(boundaries)\n return quote\n Base.Cartesian.@nif(\n $(N + 1),\n i -> bctype == i, # conditionexpr\n i -> ocean_boundary_state!(\n nf,\n boundaries[i],\n ocean,\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n t,\n ), # expr\n i -> error(\"Invalid boundary tag\")\n ) # elseexpr\n return nothing\n end\nend\n\n\"\"\"\n ocean_boundary_state!(nf, boundaries::Tuple, ::HBModel,\n Q⁺, A⁺, D⁺, n, Q⁻, A⁻, D⁻, bctype)\n\napplies boundary conditions for the second-order fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@generated function ocean_boundary_state!(\n nf::NumericalFluxSecondOrder,\n boundaries::Tuple,\n ocean,\n Q⁺,\n D⁺,\n HD⁺,\n A⁺,\n n,\n Q⁻,\n D⁻,\n HD⁻,\n A⁻,\n bctype,\n t,\n args...,\n)\n N = fieldcount(boundaries)\n return quote\n Base.Cartesian.@nif(\n $(N + 1),\n i -> bctype == i, # conditionexpr\n i -> ocean_boundary_state!(\n nf,\n boundaries[i],\n ocean,\n Q⁺,\n D⁺,\n A⁺,\n n,\n Q⁻,\n D⁻,\n A⁻,\n t,\n ), # expr\n i -> error(\"Invalid boundary tag\")\n ) # elseexpr\n return nothing\n end\nend\n" }, { "path": "src/Ocean/JLD2Writer.jl", "content": "module JLD2Writers\n\nusing JLD2\n\nusing ..CartesianDomains: DiscontinuousSpectralElementGrid\nusing ..Ocean: current_step, current_time\nusing ..CartesianFields: SpectralElementField\n\nstruct JLD2Writer{A, F, M, O}\n filepath::F\n model::M\n outputs::O\n array_type::A\nend\n\nfunction Base.show(io::IO, writer::JLD2Writer{A}) where {A}\n\n header = \"JLD2Writer{$A}\"\n filepath = \" ├── filepath: $(writer.filepath)\"\n outputs = \" └── $(length(writer.outputs)) outputs: $(keys(ow.outputs))\"\n\n print(io, header, '\\n', filepath, '\\n', outputs)\n\n return nothing\nend\n\n\"\"\"\n JLD2Writer(model, outputs=model.fields; filepath, array_type=Array, overwrite_existing=true)\n\nReturns a utility for writing field output to JLD2 files, `overwrite_existing` file at `filepath`.\n\n`write!(jld2_writer::JLD2Writer)` writes `outputs` to `filepath`, where `outputs` is either a `NamedTuple` or\n`Dict`ionary of `fields` or functions of the form `output(model)`.\n\nField data is converted to `array_type` before outputting.\n\"\"\"\nfunction JLD2Writer(\n model,\n outputs = model.fields;\n filepath,\n array_type = Array,\n overwrite_existing = true,\n)\n\n # Convert grid to CPU\n cpu_grid =\n DiscontinuousSpectralElementGrid(model.domain, array_type = array_type)\n\n # Initialize output\n overwrite_existing && isfile(filepath) && rm(filepath; force = true)\n\n file = jldopen(filepath, \"a+\")\n\n file[\"domain\"] = model.domain\n file[\"grid\"] = cpu_grid\n\n close(file)\n\n writer = JLD2Writer(filepath, model, outputs, array_type)\n\n write!(writer, first = true)\n\n return writer\nend\n\ninitialize_output!(file, args...) = nothing\n\ninitialize_output!(file, field::SpectralElementField, name) =\n file[\"$name/meta/realelems\"] = field.realelems\n\nfunction write!(writer::JLD2Writer; first = false)\n\n model = writer.model\n filepath = writer.filepath\n outputs = writer.outputs\n\n # Add new data to file\n file = jldopen(filepath, \"a+\")\n\n N_output = first ? 0 : length(keys(file[\"times\"]))\n\n step = current_step(model)\n time = current_time(model)\n\n file[\"times/$N_output\"] = time\n file[\"steps/$N_output\"] = step\n\n for (name, output) in zip(keys(outputs), values(outputs))\n first && initialize_output!(file, output, name)\n write_single_output!(file, output, name, N_output, writer)\n end\n\n close(file)\n\n return nothing\nend\n\nfunction write_field!(file, field, name, N_output, writer)\n data = convert(writer.array_type, field.data)\n file[\"$name/$N_output\"] = data\n return nothing\nend\n\nwrite_single_output!(file, field::SpectralElementField, args...) =\n write_field!(file, field, args...)\n\nfunction write_single_output!(file, output, name, N_output, writer)\n data = output(writer.model)\n file[\"$name/$N_output\"] = data\n return nothing\nend\n\nstruct OutputTimeSeries{F, N, D, G, T, S}\n filepath::F\n name::N\n domain::D\n grid::G\n times::T\n steps::S\nend\n\nfunction OutputTimeSeries(name, filepath)\n file = jldopen(filepath)\n\n domain, grid, times, steps = [nothing for i in 1:4]\n\n try\n domain = file[\"domain\"]\n grid = file[\"grid\"]\n\n output_indices = keys(file[\"times\"])\n\n times = [file[\"times/$i\"] for i in output_indices]\n steps = [file[\"steps/$i\"] for i in output_indices]\n\n catch err\n @warn \"Could not build time series of $name from $filepath because $(sprint(showerror, err))\"\n\n finally\n close(file)\n\n end\n\n return OutputTimeSeries(filepath, name, domain, grid, times, steps)\nend\n\nfunction Base.length(timeseries::OutputTimeSeries)\n file = jldopen(timeseries.filepath)\n\n timeseries_length = length(keys(file[\"times\"]))\n\n close(file)\n\n return timeseries_length\nend\n\n\nfunction Base.getindex(timeseries::OutputTimeSeries, i)\n\n name = timeseries.name\n domain = timeseries.domain\n grid = timeseries.grid\n\n file = jldopen(timeseries.filepath)\n\n data = file[\"$name/$(i-1)\"]\n realelems = file[\"$name/meta/realelems\"]\n\n close(file)\n\n realdata = view(data, :, realelems)\n\n return SpectralElementField(domain, grid, realdata, data, realelems)\nend\n\nend # module\n" }, { "path": "src/Ocean/Ocean.jl", "content": "module Ocean\n\nusing ..BalanceLaws\nusing ..CartesianDomains\nusing ..CartesianFields\nusing ..Problems\n\nexport AbstractOceanModel,\n AbstractOceanProblem,\n AbstractOceanCoupling,\n Uncoupled,\n Coupled,\n AdvectionTerm,\n NonLinearAdvectionTerm,\n InitialConditions\n\nabstract type AbstractOceanModel <: BalanceLaw end\nabstract type AbstractOceanProblem <: AbstractProblem end\n\nabstract type AbstractOceanCoupling end\nstruct Uncoupled <: AbstractOceanCoupling end\nstruct Coupled <: AbstractOceanCoupling end\n\nabstract type AdvectionTerm end\nstruct NonLinearAdvectionTerm <: AdvectionTerm end\n\nfunction ocean_init_state! end\nfunction ocean_init_aux! end\nfunction ocean_boundary_state! end\n\nfunction coriolis_parameter end\nfunction kinematic_stress end\nfunction surface_flux end\n\ninclude(\"OceanBC.jl\")\n\ninclude(\"HydrostaticBoussinesq/HydrostaticBoussinesq.jl\")\n\nusing .HydrostaticBoussinesq: HydrostaticBoussinesqModel, Forcing\n\ninclude(\"ShallowWater/ShallowWaterModel.jl\")\ninclude(\"SplitExplicit/SplitExplicitModel.jl\")\ninclude(\"SplitExplicit01/SplitExplicitModel.jl\")\ninclude(\"OceanProblems/OceanProblems.jl\")\n\ninclude(\"SuperModels.jl\")\n\nusing .OceanProblems: InitialConditions\nusing .SuperModels:\n HydrostaticBoussinesqSuperModel, current_time, current_step, Δt\n\ninclude(\"JLD2Writer.jl\")\n\nusing .JLD2Writers: JLD2Writer, OutputTimeSeries, write!\n\nend\n" }, { "path": "src/Ocean/OceanBC.jl", "content": "export OceanBC,\n VelocityBC,\n VelocityDragBC,\n TemperatureBC,\n Impenetrable,\n Penetrable,\n NoSlip,\n FreeSlip,\n KinematicStress,\n Insulating,\n TemperatureFlux\n\nusing StaticArrays\n\nusing ..BalanceLaws\nusing ..DGMethods.NumericalFluxes\n\n\"\"\"\n OceanBC(velocity = Impenetrable(NoSlip())\n temperature = Insulating())\n\nThe standard boundary condition for OceanModel. The default options imply a \"no flux\" boundary condition.\n\"\"\"\nBase.@kwdef struct OceanBC{M, T}\n velocity::M = Impenetrable(NoSlip())\n temperature::T = Insulating()\nend\n\nabstract type VelocityBC end\nabstract type VelocityDragBC end\nabstract type TemperatureBC end\n\n\"\"\"\n Impenetrable(drag::VelocityDragBC) :: VelocityBC\n\nDefines an impenetrable wall model for velocity. This implies:\n - no flow in the direction normal to the boundary, and\n - flow parallel to the boundary is subject to the `drag` condition.\n\"\"\"\nstruct Impenetrable{D <: VelocityDragBC} <: VelocityBC\n drag::D\nend\n\n\"\"\"\n Penetrable(drag::VelocityDragBC) :: VelocityBC\n\nDefines an penetrable wall model for velocity. This implies:\n - no constraint on flow in the direction normal to the boundary, and\n - flow parallel to the boundary is subject to the `drag` condition.\n\"\"\"\nstruct Penetrable{D <: VelocityDragBC} <: VelocityBC\n drag::D\nend\n\n\"\"\"\n NoSlip() :: VelocityDragBC\n\nZero velocity at the boundary.\n\"\"\"\nstruct NoSlip <: VelocityDragBC end\n\n\"\"\"\n FreeSlip() :: VelocityDragBC\n\nNo surface drag on velocity parallel to the boundary.\n\"\"\"\nstruct FreeSlip <: VelocityDragBC end\n\n\"\"\"\n KinematicStress(stress) :: VelocityDragBC\n\nApplies the specified kinematic stress on velocity normal to the boundary.\nPrescribe the net inward kinematic stress across the boundary by `stress`,\na function with signature `stress(problem, state, aux, t)`, returning the flux (in m²/s²).\n\"\"\"\nstruct KinematicStress{S} <: VelocityDragBC\n stress::S\n\n function KinematicStress(stress::S = nothing) where {S}\n new{S}(stress)\n end\nend\n\nkinematic_stress(problem, y, ρ₀) = @SVector [0, 0] # fallback for generic problems\nkinematic_stress(problem, y, ρ₀, ::Nothing) = kinematic_stress(problem, y, ρ₀)\nkinematic_stress(problem, y, ρ₀, stress) = stress(y)\n\n\"\"\"\n Insulating() :: TemperatureBC\n\nNo temperature flux across the boundary\n\"\"\"\nstruct Insulating <: TemperatureBC end\n\n\"\"\"\n TemperatureFlux(flux) :: TemperatureBC\n\nPrescribe the net inward temperature flux across the boundary by `flux`,\na function with signature `flux(problem, state, aux, t)`, returning the flux (in m⋅K/s).\n\"\"\"\nstruct TemperatureFlux{T} <: TemperatureBC\n flux::T\n\n function TemperatureFlux(flux::T = nothing) where {T}\n new{T}(flux)\n end\nend\n\n# these functions just trim off the extra arguments\nfunction _ocean_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc,\n ocean,\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n t,\n _...,\n)\n return ocean_boundary_state!(nf, bc, ocean, Q⁺, A⁺, n, Q⁻, A⁻, t)\nend\n\nfunction _ocean_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc,\n ocean,\n Q⁺,\n D⁺,\n HD⁺,\n A⁺,\n n,\n Q⁻,\n D⁻,\n HD⁻,\n A⁻,\n t,\n _...,\n)\n return ocean_boundary_state!(nf, bc, ocean, Q⁺, D⁺, A⁺, n, Q⁻, D⁻, A⁻, t)\nend\n" }, { "path": "src/Ocean/OceanProblems/OceanProblems.jl", "content": "module OceanProblems\n\nexport SimpleBox, Fixed, Rotating, HomogeneousBox, OceanGyre\n\nusing StaticArrays\nusing CLIMAParameters.Planet: grav\n\nusing ...Problems\n\nusing ..Ocean\nusing ..BalanceLaws: parameter_set\nusing ..HydrostaticBoussinesq\nusing ..ShallowWater\nusing ..SplitExplicit01\n\nimport ..Ocean:\n ocean_init_state!,\n ocean_init_aux!,\n kinematic_stress,\n surface_flux,\n coriolis_parameter\n\nHBModel = HydrostaticBoussinesqModel\nSWModel = ShallowWaterModel\n\ninclude(\"simple_box_problem.jl\")\ninclude(\"homogeneous_box.jl\")\n\ninclude(\"shallow_water_initial_states.jl\")\n\nfunction ocean_init_state!(\n m::SWModel,\n p::HomogeneousBox,\n Q,\n A,\n local_geometry,\n t,\n)\n if t == 0\n null_init_state!(p, m.turbulence, Q, A, local_geometry, 0)\n else\n gyre_init_state!(m, p, m.turbulence, Q, A, local_geometry, t)\n end\nend\n\n@inline coriolis_parameter(m::SWModel, p::HomogeneousBox, y) =\n m.fₒ + m.β * (y - p.Lʸ / 2)\n\ninclude(\"ocean_gyre.jl\")\n\ninclude(\"initial_value_problem.jl\")\n\nend # module\n" }, { "path": "src/Ocean/OceanProblems/homogeneous_box.jl", "content": "##########################\n# Homogenous wind stress #\n# Constant temperature #\n##########################\n\n\"\"\"\n HomogeneousBox <: AbstractSimpleBoxProblem\n\nContainer structure for a simple box problem with wind-stress.\nLˣ = zonal (east-west) length\nLʸ = meridional (north-south) length\nH = height of the ocean\nτₒ = maximum value of wind-stress (amplitude)\n\"\"\"\nstruct HomogeneousBox{T, BC} <: AbstractSimpleBoxProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n boundary_conditions::BC\n function HomogeneousBox{FT}(\n Lˣ, # m\n Lʸ, # m\n H; # m\n τₒ = FT(1e-1), # N/m²\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), Insulating()),\n ),\n ) where {FT <: AbstractFloat}\n return new{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, BC)\n end\nend\n\n\"\"\"\n ocean_init_state!(::HomogeneousBox)\n\ninitialize u,v with random values, η with 0, and θ with a constant (20)\n\n# Arguments\n- `p`: HomogeneousBox problem object, used to dispatch on\n- `Q`: state vector\n- `A`: auxiliary state vector, not used\n- `localgeo`: the local geometry, not used\n- `t`: time to evaluate at, not used\n\"\"\"\nfunction ocean_init_state!(m::HBModel, p::HomogeneousBox, Q, A, localgeo, t)\n Q.u = @SVector [0, 0]\n Q.η = 0\n Q.θ = 20\n\n return nothing\nend\n\n\"\"\"\n kinematic_stress(::HomogeneousBox)\n\njet stream like windstress\n\n# Arguments\n- `p`: problem object to dispatch on and get additional parameters\n- `y`: y-coordinate in the box\n\"\"\"\n@inline kinematic_stress(p::HomogeneousBox, y, ρ) =\n @SVector [(p.τₒ / ρ) * cos(y * π / p.Lʸ), -0]\n\n@inline kinematic_stress(p::HomogeneousBox, y) =\n @SVector [-p.τₒ * cos(π * y / p.Lʸ), -0]\n" }, { "path": "src/Ocean/OceanProblems/initial_value_problem.jl", "content": "#####\n##### Initial value problem\n#####\n\nstruct InitialValueProblem{FT, IC, BC} <: AbstractSimpleBoxProblem\n Lˣ::FT\n Lʸ::FT\n H::FT\n initial_conditions::IC\n boundary_conditions::BC\n\n \"\"\"\n InitialValueProblem{FT}(; dimensions, initial_conditions=InitialConditions(),\n boundary_conditions = (OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating())))\n\n Returns an `InitialValueProblem` with `dimensions = (Lˣ, Lʸ, H)`, `initial_conditions`,\n and `boundary_conditions`.\n\n The default `initial_conditions` are resting with no temperature perturbation;\n the default list of `boundary_conditions` provide implementations for an impenetrable, insulating\n boundary, and a penetrable, insulating boundary.\n \"\"\"\n function InitialValueProblem{FT}(;\n dimensions,\n initial_conditions = InitialConditions(),\n boundary_conditions = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n ),\n ) where {FT}\n\n return new{FT, typeof(initial_conditions), typeof(boundary_conditions)}(\n FT.(dimensions)...,\n initial_conditions,\n boundary_conditions,\n )\n end\n\nend\n\n\n#####\n##### Initial conditions\n#####\n\nresting(x, y, z) = 0\n\nstruct InitialConditions{U, V, T, E}\n u::U\n v::V\n θ::T\n η::E\nend\n\n\n\"\"\"\n InitialConditions(; u=resting, v=resting, θ=resting, η=resting)\n\nStores initial conditions for each prognostic variable provided as functions of `x, y, z`.\n\nExample\n=======\n\n# A Gaussian surface perturbation\na = 0.1 # m, amplitude\nL = 1e5 # m, horizontal scale of the perturbation\n\nηᵢ(x, y, z) = a * exp(-(x^2 + y^2) / 2L^2)\n\nics = InitialConditions(η=ηᵢ)\n\"\"\"\nInitialConditions(; u = resting, v = resting, θ = resting, η = resting) =\n InitialConditions(u, v, θ, η)\n\n\"\"\"\n ocean_init_state!(::HydrostaticBoussinesqModel, ic::InitialCondition, state, aux, local_geometry, time)\n\nInitialize the state variables `u = (u, v)` (a vector), `θ`, and `η`. Mutates `state`.\n\nThis function is called by `init_state_prognostic!(::HydrostaticBoussinesqModel, ...)`.\n\"\"\"\nfunction ocean_init_state!(\n ::HydrostaticBoussinesqModel,\n ivp::InitialValueProblem,\n state,\n aux,\n local_geometry,\n time,\n)\n\n ics = ivp.initial_conditions\n x, y, z = local_geometry.coord\n\n state.u = @SVector [ics.u(x, y, z), ics.v(x, y, z)]\n state.θ = ics.θ(x, y, z)\n state.η = ics.η(x, y, z)\n\n return nothing\nend\n" }, { "path": "src/Ocean/OceanProblems/ocean_gyre.jl", "content": "\"\"\"\n OceanGyre <: AbstractSimpleBoxProblem\n\nContainer structure for a simple box problem with wind-stress, coriolis force, and temperature forcing.\nLˣ = zonal (east-west) length\nLʸ = meridional (north-south) length\nH = height of the ocean\nτₒ = maximum value of wind-stress (amplitude)\nλʳ = temperature relaxation penetration constant (meters / second)\nθᴱ = maximum surface temperature\n\"\"\"\nstruct OceanGyre{T, BC} <: AbstractSimpleBoxProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n λʳ::T\n θᴱ::T\n boundary_conditions::BC\n function OceanGyre{FT}(\n Lˣ, # m\n Lʸ, # m\n H; # m\n τₒ = FT(1e-1), # N/m²\n λʳ = FT(4 // 86400), # m/s\n θᴱ = FT(10), # K\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), TemperatureFlux()),\n ),\n ) where {FT <: AbstractFloat}\n return new{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, λʳ, θᴱ, BC)\n end\nend\n\n\"\"\"\n ocean_init_state!(::OceanGyre)\n\ninitialize u,v,η with 0 and θ linearly distributed between 9 at z=0 and 1 at z=H\n\n# Arguments\n- `p`: OceanGyre problem object, used to dispatch on and obtain ocean height H\n- `Q`: state vector\n- `A`: auxiliary state vector, not used\n- `localgeo`: the local geometry information\n- `t`: time to evaluate at, not used\n\"\"\"\nfunction ocean_init_state!(\n ::Union{HBModel, OceanModel},\n p::OceanGyre,\n Q,\n A,\n localgeo,\n t,\n)\n coords = localgeo.coord\n @inbounds y = coords[2]\n @inbounds z = coords[3]\n @inbounds H = p.H\n\n Q.u = @SVector [-0, -0]\n Q.η = -0\n Q.θ = (5 + 4 * cos(y * π / p.Lʸ)) * (1 + z / H)\n\n return nothing\nend\n\nfunction ocean_init_state!(\n ::Union{SWModel, BarotropicModel},\n ::OceanGyre,\n Q,\n A,\n localgeo,\n t,\n)\n Q.U = @SVector [-0, -0]\n Q.η = -0\n\n return nothing\nend\n\n\"\"\"\n kinematic_stress(::OceanGyre)\n\njet stream like windstress\n\n# Arguments\n- `p`: problem object to dispatch on and get additional parameters\n- `y`: y-coordinate in the box\n\"\"\"\n@inline kinematic_stress(p::OceanGyre, y, ρ) =\n @SVector [(p.τₒ / ρ) * cos(y * π / p.Lʸ), -0]\n\n@inline kinematic_stress(p::OceanGyre, y) =\n @SVector [-p.τₒ * cos(π * y / p.Lʸ), -0]\n\n\"\"\"\n surface_flux(::OceanGyre)\n\ncool-warm north-south linear temperature gradient\n\n# Arguments\n- `p`: problem object to dispatch on and get additional parameters\n- `y`: y-coordinate in the box\n- `θ`: temperature within element on boundary\n\"\"\"\n@inline function surface_flux(p::OceanGyre, y, θ)\n Lʸ = p.Lʸ\n θᴱ = p.θᴱ\n λʳ = p.λʳ\n\n θʳ = θᴱ * (1 - y / Lʸ)\n return λʳ * (θ - θʳ)\nend\n" }, { "path": "src/Ocean/OceanProblems/shallow_water_initial_states.jl", "content": "using ..ShallowWater: TurbulenceClosure, LinearDrag, ConstantViscosity\n\nfunction null_init_state!(\n ::HomogeneousBox,\n ::TurbulenceClosure,\n state,\n aux,\n local_geometry,\n t,\n)\n T = eltype(state.U)\n state.U = @SVector zeros(T, 2)\n state.η = 0\n return nothing\nend\n\nη_lsw(x, y, t) = cos(π * x) * cos(π * y) * cos(√2 * π * t)\nu_lsw(x, y, t) = 2^(-0.5) * sin(π * x) * cos(π * y) * sin(√2 * π * t)\nv_lsw(x, y, t) = 2^(-0.5) * cos(π * x) * sin(π * y) * sin(√2 * π * t)\n\nfunction lsw_init_state!(\n m::ShallowWaterModel,\n p::HomogeneousBox,\n state,\n aux,\n local_geometry,\n t,\n)\n\n coords = local_geometry.coord\n\n state.U = @SVector [\n u_lsw(coords[1], coords[2], t),\n v_lsw(coords[1], coords[2], t),\n ]\n\n state.η = η_lsw(coords[1], coords[2], t)\n\n return nothing\nend\n\nv_lkw(x, y, t) = 0\nu_lkw(x, y, t) = exp(-0.5 * y^2) * exp(-0.5 * (x - t + 5)^2)\nη_lkw(x, y, t) = 1 + u_lkw(x, y, t)\n\nfunction lkw_init_state!(\n m::ShallowWaterModel,\n p::HomogeneousBox,\n state,\n aux,\n local_geometry,\n t,\n)\n\n coords = local_geometry.coord\n\n state.U = @SVector [\n u_lkw(coords[1], coords[2], t),\n v_lkw(coords[1], coords[2], t),\n ]\n\n state.η = η_lkw(coords[1], coords[2], t)\n\n return nothing\nend\n\nR₋(ϵ) = (-1 - sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)\nR₊(ϵ) = (-1 + sqrt(1 + (2 * π * ϵ)^2)) / (2ϵ)\nD(ϵ) =\n (R₊(ϵ) * (exp(R₋(ϵ)) - 1) + R₋(ϵ) * (1 - exp(R₊(ϵ)))) /\n (exp(R₊(ϵ)) - exp(R₋(ϵ)))\nR₂(x₁, ϵ) =\n (1 / D(ϵ)) * (\n (\n (R₊(ϵ) * (exp(R₋(ϵ)) - 1)) * exp(R₊(ϵ) * x₁) +\n (R₋(ϵ) * (1 - exp(R₊(ϵ)))) * exp(R₋(ϵ) * x₁)\n ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))\n )\nR₁(x₁, ϵ) =\n (π / D(ϵ)) * (\n 1 .+\n (\n (exp(R₋(ϵ)) - 1) * exp(R₊(ϵ) * x₁) .+\n (1 - exp(R₊(ϵ))) * exp(R₋(ϵ) * x₁)\n ) / (exp(R₊(ϵ)) - exp(R₋(ϵ)))\n )\n\n𝒱(x₁, y₁, ϵ) = R₂(x₁, ϵ) * sin.(π * y₁)\n𝒰(x₁, y₁, ϵ) = -R₁(x₁, ϵ) * cos.(π * y₁)\nℋ(x₁, y₁, ϵ, βᵖ, fₒ, γ) =\n (R₂(x₁, ϵ) / (π * fₒ)) * γ * cos(π * y₁) +\n (R₁(x₁, ϵ) / π) *\n (sin(π * y₁) * (1.0 + βᵖ * (y₁ - 0.5)) + (βᵖ / π) * cos(π * y₁))\n\nfunction gyre_init_state!(\n m::SWModel,\n p::HomogeneousBox,\n T::LinearDrag,\n state,\n aux,\n local_geometry,\n t,\n)\n\n coords = local_geometry.coord\n\n FT = eltype(state)\n τₒ = p.τₒ\n fₒ = m.fₒ\n β = m.β\n Lˣ = p.Lˣ\n Lʸ = p.Lʸ\n H = p.H\n\n γ = T.λ\n\n βᵖ = β * Lʸ / fₒ\n ϵ = γ / (Lˣ * β)\n\n _grav::FT = grav(parameter_set(m))\n\n uˢ(ϵ) = (τₒ * D(ϵ)) / (H * γ * π)\n hˢ(ϵ) = (fₒ * Lˣ * uˢ(ϵ)) / _grav\n\n u = uˢ(ϵ) * 𝒰(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)\n v = uˢ(ϵ) * 𝒱(coords[1] / Lˣ, coords[2] / Lʸ, ϵ)\n h = hˢ(ϵ) * ℋ(coords[1] / Lˣ, coords[2] / Lʸ, ϵ, βᵖ, fₒ, γ)\n\n state.U = @SVector [H * u, H * v]\n\n state.η = h\n\n return nothing\nend\n\nt1(x, δᵐ) = cos((√3 * x) / (2 * δᵐ)) + (√3^-1) * sin((√3 * x) / (2 * δᵐ))\nt2(x, δᵐ) = 1 - exp((-x) / (2 * δᵐ)) * t1(x, δᵐ)\nt3(y, Lʸ) = π * sin(π * y / Lʸ)\nt4(x, Lˣ, C) = C * (1 - x / Lˣ)\n\nη_munk(x, y, Lˣ, Lʸ, δᵐ, C) = t4(x, Lˣ, C) * t3(y, Lʸ) * t2(x, δᵐ)\n\nfunction gyre_init_state!(\n m::SWModel,\n p::HomogeneousBox,\n V::ConstantViscosity,\n state,\n aux,\n local_geometry,\n t,\n)\n\n coords = local_geometry.coord\n\n FT = eltype(state.U)\n _grav::FT = grav(parameter_set(m))\n\n τₒ = p.τₒ\n fₒ = m.fₒ\n β = m.β\n Lˣ = p.Lˣ\n Lʸ = p.Lʸ\n H = p.H\n\n ν = V.ν\n\n δᵐ = (ν / β)^(1 / 3)\n C = τₒ / (_grav * H) * (fₒ / β)\n\n state.η = η_munk(coords[1], coords[2], Lˣ, Lʸ, δᵐ, C)\n state.U = @SVector zeros(FT, 2)\n\n return nothing\nend\n" }, { "path": "src/Ocean/OceanProblems/simple_box_problem.jl", "content": "abstract type AbstractSimpleBoxProblem <: AbstractOceanProblem end\n\n\"\"\"\n ocean_init_aux!(::HBModel, ::AbstractSimpleBoxProblem)\n\nsave y coordinate for computing coriolis, wind stress, and sea surface temperature\n\n# Arguments\n- `m`: model object to dispatch on and get viscosities and diffusivities\n- `p`: problem object to dispatch on and get additional parameters\n- `A`: auxiliary state vector\n- `geom`: geometry stuff\n\"\"\"\nfunction ocean_init_aux!(::HBModel, ::AbstractSimpleBoxProblem, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # needed for proper CFL condition calculation\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n\n A.uᵈ = @SVector [-0, -0]\n A.ΔGᵘ = @SVector [-0, -0]\n\n return nothing\nend\n\nfunction ocean_init_aux!(::OceanModel, ::AbstractSimpleBoxProblem, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # needed for proper CFL condition calculation\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n\n A.u_d = @SVector [-0, -0]\n A.ΔGu = @SVector [-0, -0]\n\n return nothing\nend\n\nfunction ocean_init_aux!(::SWModel, ::AbstractSimpleBoxProblem, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.Δu = @SVector [-0, -0]\n\n return nothing\nend\n\nfunction ocean_init_aux!(::BarotropicModel, ::AbstractSimpleBoxProblem, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.U_c = @SVector [-0, -0]\n A.η_c = -0\n A.U_s = @SVector [-0, -0]\n A.η_s = -0\n A.Δu = @SVector [-0, -0]\n A.η_diag = -0\n A.Δη = -0\n\n return nothing\nend\n\n\"\"\"\n coriolis_parameter\n\nnorthern hemisphere coriolis\n\n# Arguments\n- `m`: model object to dispatch on and get coriolis parameters\n- `y`: y-coordinate in the box\n\"\"\"\n@inline coriolis_parameter(\n m::Union{HBModel, OceanModel},\n ::AbstractSimpleBoxProblem,\n y,\n) = m.fₒ + m.β * y\n@inline coriolis_parameter(\n m::Union{SWModel, BarotropicModel},\n ::AbstractSimpleBoxProblem,\n y,\n) = m.fₒ + m.β * y\n\n############################\n# Basic box problem #\n# Set up dimensions of box #\n############################\n\nabstract type AbstractRotation end\nstruct Rotating <: AbstractRotation end\nstruct Fixed <: AbstractRotation end\n\n\"\"\"\n SimpleBoxProblem <: AbstractSimpleBoxProblem\n\nStub structure with the dimensions of the box.\nLˣ = zonal (east-west) length\nLʸ = meridional (north-south) length\nH = height of the ocean\n\"\"\"\nstruct SimpleBox{R, T, BC} <: AbstractSimpleBoxProblem\n rotation::R\n Lˣ::T\n Lʸ::T\n H::T\n boundary_conditions::BC\n function SimpleBox{FT}(\n Lˣ, # m\n Lʸ, # m\n H; # m\n rotation = Fixed(),\n BC = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n ),\n ) where {FT <: AbstractFloat}\n return new{typeof(rotation), FT, typeof(BC)}(rotation, Lˣ, Lʸ, H, BC)\n end\nend\n\n@inline coriolis_parameter(\n m::Union{HBModel, OceanModel},\n ::SimpleBox{R},\n y,\n) where {R <: Fixed} = -0\n@inline coriolis_parameter(\n m::Union{SWModel, BarotropicModel},\n ::SimpleBox{R},\n y,\n) where {R <: Fixed} = -0\n\n@inline coriolis_parameter(\n m::Union{HBModel, OceanModel},\n ::SimpleBox{R},\n y,\n) where {R <: Rotating} = m.fₒ\n@inline coriolis_parameter(\n m::Union{SWModel, BarotropicModel},\n ::SimpleBox{R},\n y,\n) where {R <: Rotating} = m.fₒ\n\nfunction ocean_init_state!(\n m::Union{SWModel, BarotropicModel},\n p::SimpleBox,\n Q,\n A,\n local_geometry,\n t,\n)\n coords = local_geometry.coord\n\n k = (2π / p.Lˣ, 2π / p.Lʸ, 2π / p.H)\n ν = viscosity(m)\n\n gH = gravity_speed(m)\n @inbounds f = coriolis_parameter(m, p, coords[2])\n\n U, V, η = barotropic_state!(p.rotation, (coords..., t), ν, k, (gH, f))\n\n Q.U = @SVector [U, V]\n Q.η = η\n\n return nothing\nend\n\nviscosity(m::SWModel) = (m.turbulence.ν, m.turbulence.ν, -0)\nviscosity(m::BarotropicModel) = (m.baroclinic.νʰ, m.baroclinic.νʰ, -0)\n\ngravity_speed(m::SWModel) = grav(parameter_set(m)) * m.problem.H\ngravity_speed(m::BarotropicModel) =\n grav(parameter_set(m)) * m.baroclinic.problem.H\n\nfunction ocean_init_state!(\n m::Union{HBModel, OceanModel},\n p::SimpleBox,\n Q,\n A,\n local_geometry,\n t,\n)\n\n coords = local_geometry.coord\n\n k = (2π / p.Lˣ, 2π / p.Lʸ, 2π / p.H)\n ν = (m.νʰ, m.νʰ, m.νᶻ)\n\n gH = grav(parameter_set(m)) * p.H\n @inbounds f = coriolis_parameter(m, p, coords[2])\n\n U, V, η = barotropic_state!(p.rotation, (coords..., t), ν, k, (gH, f))\n u°, v° = baroclinic_deviation(p.rotation, (coords..., t), ν, k, f)\n\n u = u° + U / p.H\n v = v° + V / p.H\n\n Q.u = @SVector [u, v]\n Q.η = η\n Q.θ = -0\n\n return nothing\nend\n\nfunction barotropic_state!(\n ::Fixed,\n (x, y, z, t),\n (νˣ, νʸ, νᶻ),\n (kˣ, kʸ, kᶻ),\n params,\n)\n gH, _ = params\n\n M = @SMatrix [-νˣ*kˣ^2 gH*kˣ; -kˣ 0]\n A = exp(M * t) * @SVector [1, 1]\n\n U = A[1] * sin(kˣ * x)\n V = -0\n η = A[2] * cos(kˣ * x)\n\n return (U = U, V = V, η = η)\nend\n\nfunction baroclinic_deviation(\n ::Fixed,\n (x, y, z, t),\n (νˣ, νʸ, νᶻ),\n (kˣ, kʸ, kᶻ),\n f,\n)\n λ = νˣ * kˣ^2 + νᶻ * kᶻ^2\n\n u° = exp(-λ * t) * cos(kᶻ * z) * sin(kˣ * x)\n v° = -0\n\n return (u° = u°, v° = v°)\nend\n\nfunction barotropic_state!(\n ::Rotating,\n (x, y, z, t),\n (νˣ, νʸ, νᶻ),\n (kˣ, kʸ, kᶻ),\n params,\n)\n gH, f = params\n\n M = @SMatrix [-νˣ*kˣ^2 f gH*kˣ; -f -νˣ*kˣ^2 0; -kˣ 0 0]\n A = exp(M * t) * @SVector [1, 1, 1]\n\n U = A[1] * sin(kˣ * x)\n V = A[2] * sin(kˣ * x)\n η = A[3] * cos(kˣ * x)\n\n return (U = U, V = V, η = η)\nend\n\nfunction baroclinic_deviation(\n ::Rotating,\n (x, y, z, t),\n (νˣ, νʸ, νᶻ),\n (kˣ, kʸ, kᶻ),\n f,\n)\n λ = νˣ * kˣ^2 + νᶻ * kᶻ^2\n\n M = @SMatrix[-λ f; -f -λ]\n A = exp(M * t) * @SVector[1, 1]\n\n u° = A[1] * cos(kᶻ * z) * sin(kˣ * x)\n v° = A[2] * cos(kᶻ * z) * sin(kˣ * x)\n\n return (u° = u°, v° = v°)\nend\n\n@inline kinematic_stress(p::SimpleBox, y) = @SVector [-0, -0]\n" }, { "path": "src/Ocean/README.md", "content": "Code for shallow water problem configurations and for specializing ClimateMachine core tools to ocean scenarios.\n\n# Code structure (aspirational)\n\n# Equations\n - ExplicitHydrostaticBoussinesq\n - SplitExplicitHydrostaticBoussinesq\n - ShallowWater\n\n# Models\n - ExplicitHydrostaticBoussinesq\n - SplitExplicitHydrostaticBoussinesq\n - ShallowWater\n" }, { "path": "src/Ocean/ShallowWater/ShallowWaterModel.jl", "content": "module ShallowWater\n\nexport ShallowWaterModel\n\nusing StaticArrays\nusing ...MPIStateArrays: MPIStateArray\nusing LinearAlgebra: dot, Diagonal\nusing CLIMAParameters.Planet: grav\n\nusing ..Ocean\nusing ...VariableTemplates\nusing ...Mesh.Geometry\nusing ...DGMethods\nusing ...DGMethods.NumericalFluxes\nusing ...BalanceLaws\nusing ..Ocean: kinematic_stress, coriolis_parameter\n\nimport ...DGMethods.NumericalFluxes: update_penalty!\nimport ...BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_conditions,\n boundary_state!\nimport ..Ocean:\n ocean_init_state!,\n ocean_init_aux!,\n ocean_boundary_state!,\n _ocean_boundary_state!\n\nusing ...Mesh.Geometry: LocalGeometry\n\n×(a::SVector, b::SVector) = StaticArrays.cross(a, b)\n⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)\n⊗(a::SVector, b::SVector) = a * b'\n\nabstract type TurbulenceClosure end\nstruct LinearDrag{L} <: TurbulenceClosure\n λ::L\nend\nstruct ConstantViscosity{L} <: TurbulenceClosure\n ν::L\nend\n\n\"\"\"\n ShallowWaterModel <: BalanceLaw\n\nA `BalanceLaw` for shallow water modeling.\n\nwrite out the equations here\n\n# Usage\n\n ShallowWaterModel(problem)\n\n\"\"\"\nstruct ShallowWaterModel{C, PS, P, T, A, FT} <: BalanceLaw\n param_set::PS\n problem::P\n coupling::C\n turbulence::T\n advection::A\n c::FT\n fₒ::FT\n β::FT\n function ShallowWaterModel{FT}(\n param_set::PS,\n problem::P,\n turbulence::T,\n advection::A;\n coupling::C = Uncoupled(),\n c = FT(0), # m/s\n fₒ = FT(1e-4), # Hz\n β = FT(1e-11), # Hz / m\n ) where {FT <: AbstractFloat, PS, P, T, A, C}\n return new{C, PS, P, T, A, FT}(\n param_set,\n problem,\n coupling,\n turbulence,\n advection,\n c,\n fₒ,\n β,\n )\n end\nend\nSWModel = ShallowWaterModel\n\nfunction vars_state(m::SWModel, ::Prognostic, T)\n @vars begin\n η::T\n U::SVector{2, T}\n end\nend\n\nfunction init_state_prognostic!(m::SWModel, state::Vars, aux::Vars, localgeo, t)\n ocean_init_state!(m, m.problem, state, aux, localgeo, t)\nend\n\nfunction vars_state(m::SWModel, ::Auxiliary, T)\n @vars begin\n y::T\n Gᵁ::SVector{2, T} # integral of baroclinic tendency\n Δu::SVector{2, T} # reconciliation Δu = 1/H * (Ū - ∫u)\n end\nend\n\nfunction init_state_auxiliary!(\n m::SWModel,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n m,\n (m, A, tmp, geom) -> ocean_init_aux!(m, m.problem, A, geom),\n state_auxiliary,\n grid,\n direction,\n )\nend\n\nfunction vars_state(m::SWModel, ::Gradient, T)\n @vars begin\n ∇U::SVector{2, T}\n end\nend\n\nfunction compute_gradient_argument!(\n m::SWModel,\n f::Vars,\n q::Vars,\n α::Vars,\n t::Real,\n)\n compute_gradient_argument!(m.turbulence, f, q, α, t)\nend\n\ncompute_gradient_argument!(::LinearDrag, _...) = nothing\n\n@inline function compute_gradient_argument!(\n T::ConstantViscosity,\n f::Vars,\n q::Vars,\n α::Vars,\n t::Real,\n)\n f.∇U = q.U\n\n return nothing\nend\n\nfunction vars_state(m::SWModel, ::GradientFlux, T)\n @vars begin\n ν∇U::SMatrix{3, 2, T, 6}\n end\nend\n\nfunction compute_gradient_flux!(\n m::SWModel,\n σ::Vars,\n δ::Grad,\n q::Vars,\n α::Vars,\n t::Real,\n)\n compute_gradient_flux!(m, m.turbulence, σ, δ, q, α, t)\nend\n\ncompute_gradient_flux!(::SWModel, ::LinearDrag, _...) = nothing\n\n@inline function compute_gradient_flux!(\n ::SWModel,\n T::ConstantViscosity,\n σ::Vars,\n δ::Grad,\n q::Vars,\n α::Vars,\n t::Real,\n)\n ν = Diagonal(@SVector [T.ν, T.ν, -0])\n ∇U = δ.∇U\n\n σ.ν∇U = -ν * ∇U\n\n return nothing\nend\n\n@inline function flux_first_order!(\n m::SWModel,\n F::Grad,\n q::Vars,\n α::Vars,\n t::Real,\n direction,\n)\n U = @SVector [q.U[1], q.U[2], -0]\n η = q.η\n H = m.problem.H\n Iʰ = @SMatrix [\n 1 -0\n -0 1\n -0 -0\n ]\n\n F.η += U\n F.U += grav(parameter_set(m)) * H * η * Iʰ\n\n advective_flux!(m, m.advection, F, q, α, t)\n\n return nothing\nend\n\nadvective_flux!(::SWModel, ::Nothing, _...) = nothing\n\n@inline function advective_flux!(\n m::SWModel,\n ::NonLinearAdvectionTerm,\n F::Grad,\n q::Vars,\n α::Vars,\n t::Real,\n)\n U = q.U\n H = m.problem.H\n V = @SVector [U[1], U[2], -0]\n\n F.U += 1 / H * V ⊗ U\n\n return nothing\nend\n\nfunction flux_second_order!(\n m::SWModel,\n G::Grad,\n q::Vars,\n σ::Vars,\n ::Vars,\n α::Vars,\n t::Real,\n)\n flux_second_order!(m, m.turbulence, G, q, σ, α, t)\nend\n\nflux_second_order!(::SWModel, ::LinearDrag, _...) = nothing\n\n@inline function flux_second_order!(\n ::SWModel,\n ::ConstantViscosity,\n G::Grad,\n q::Vars,\n σ::Vars,\n α::Vars,\n t::Real,\n)\n G.U += σ.ν∇U\n\n return nothing\nend\n\n@inline wavespeed(m::SWModel, n⁻, q::Vars, α::Vars, t::Real, direction) = m.c\n\n@inline function source!(\n m::SWModel{P},\n S::Vars,\n q::Vars,\n d::Vars,\n α::Vars,\n t::Real,\n direction,\n) where {P}\n # f × u\n U, V = q.U\n f = coriolis_parameter(m, m.problem, α.y)\n S.U -= @SVector [-f * V, f * U]\n\n forcing_term!(m, m.coupling, S, q, α, t)\n linear_drag!(m.turbulence, S, q, α, t)\n\n return nothing\nend\n\n@inline function forcing_term!(m::SWModel, ::Uncoupled, S, Q, A, t)\n S.U += kinematic_stress(m.problem, A.y)\n\n return nothing\nend\n\nlinear_drag!(::ConstantViscosity, _...) = nothing\n\n@inline function linear_drag!(T::LinearDrag, S::Vars, q::Vars, α::Vars, t::Real)\n λ = T.λ\n U = q.U\n\n S.U -= λ * U\n\n return nothing\nend\n\nboundary_conditions(shallow::SWModel) = shallow.problem.boundary_conditions\n\n\"\"\"\n boundary_state!(nf, ::SWModel, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@inline function boundary_state!(nf, bc, shallow::SWModel, args...)\n return _ocean_boundary_state!(nf, bc, shallow, args...)\nend\n\n\"\"\"\n ocean_boundary_state!(nf, bc::OceanBC, ::SWModel)\n\nsplits boundary condition application into velocity\n\"\"\"\n@inline function ocean_boundary_state!(nf, bc::OceanBC, m::SWModel, args...)\n return ocean_boundary_state!(nf, bc.velocity, m, m.turbulence, args...)\nend\n\ninclude(\"bc_velocity.jl\")\n\n\nend\n" }, { "path": "src/Ocean/ShallowWater/bc_velocity.jl", "content": "\"\"\"\n ocean_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::SWModel)\n\napply free slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{FreeSlip},\n ::SWModel,\n ::TurbulenceClosure,\n q⁺,\n a⁺,\n n⁻,\n q⁻,\n a⁻,\n t,\n args...,\n)\n q⁺.η = q⁻.η\n\n V⁻ = @SVector [q⁻.U[1], q⁻.U[2], -0]\n V⁺ = V⁻ - 2 * n⁻ ⋅ V⁻ .* SVector(n⁻)\n q⁺.U = @SVector [V⁺[1], V⁺[2]]\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::Union{NumericalFluxGradient, NumericalFluxSecondOrder}, ::Impenetrable{FreeSlip}, ::SWModel)\n\nno second order flux computed for linear drag\n\"\"\"\nocean_boundary_state!(\n ::Union{NumericalFluxGradient, NumericalFluxSecondOrder},\n ::VelocityBC,\n ::SWModel,\n ::LinearDrag,\n _...,\n) = nothing\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::SWModel)\n\napply free slip boundary condition for velocity\nsets non-reflective ghost point\n\"\"\"\nfunction ocean_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{FreeSlip},\n ::SWModel,\n ::ConstantViscosity,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n args...,\n)\n V⁻ = @SVector [Q⁻.U[1], Q⁻.U[2], -0]\n V⁺ = V⁻ - n⁻ ⋅ V⁻ .* SVector(n⁻)\n Q⁺.U = @SVector [V⁺[1], V⁺[2]]\n\n return nothing\nend\n\n\"\"\"\n shallow_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::SWModel)\n\napply free slip boundary condition for velocity\napply zero numerical flux in the normal direction\n\"\"\"\nfunction ocean_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{FreeSlip},\n ::SWModel,\n ::ConstantViscosity,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n args...,\n)\n Q⁺.U = Q⁻.U\n D⁺.ν∇U = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::SWModel)\n\napply no slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{NoSlip},\n ::SWModel,\n ::TurbulenceClosure,\n q⁺,\n α⁺,\n n⁻,\n q⁻,\n α⁻,\n t,\n args...,\n)\n q⁺.η = q⁻.η\n q⁺.U = -q⁻.U\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::SWModel)\n\napply no slip boundary condition for velocity\nset numerical flux to zero for U\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{NoSlip},\n ::SWModel,\n ::ConstantViscosity,\n q⁺,\n α⁺,\n n⁻,\n q⁻,\n α⁻,\n t,\n args...,\n)\n FT = eltype(q⁺)\n q⁺.U = @SVector zeros(FT, 2)\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::SWModel)\n\napply no slip boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for U\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{NoSlip},\n ::SWModel,\n ::ConstantViscosity,\n q⁺,\n σ⁺,\n α⁺,\n n⁻,\n q⁻,\n σ⁻,\n α⁻,\n t,\n args...,\n)\n q⁺.U = -q⁻.U\n σ⁺.ν∇U = σ⁻.ν∇U\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::SWModel)\n\nno mass boundary condition for penetrable\n\"\"\"\nocean_boundary_state!(\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{FreeSlip},\n ::SWModel,\n ::ConstantViscosity,\n _...,\n) = nothing\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::SWModel)\n\napply free slip boundary condition for velocity\napply zero numerical flux in the normal direction\n\"\"\"\nfunction ocean_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Penetrable{FreeSlip},\n ::SWModel,\n ::ConstantViscosity,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n args...,\n)\n Q⁺.U = Q⁻.U\n D⁺.ν∇U = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction ocean_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Impenetrable{<:KinematicStress},\n shallow::SWModel,\n turb::TurbulenceClosure,\n args...,\n)\n return ocean_boundary_state!(\n nf,\n Impenetrable(FreeSlip()),\n shallow,\n turb,\n args...,\n )\nend\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{<:KinematicStress},\n shallow::SWModel,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.U = Q⁻.U\n D⁺.ν∇U = n⁻ * kinematic_stress(shallow.problem, A⁻.y, 1000)'\n # applies windstress for now, will be fixed in a later PR\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction ocean_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{<:KinematicStress},\n shallow::SWModel,\n turb::TurbulenceClosure,\n args...,\n)\n return ocean_boundary_state!(\n nf,\n Penetrable(FreeSlip()),\n shallow,\n turb,\n args...,\n )\nend\n\n\"\"\"\n ocean_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{KinematicStress}, ::HBModel)\n\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function ocean_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Penetrable{<:KinematicStress},\n shallow::SWModel,\n ::TurbulenceClosure,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n Q⁺.u = Q⁻.u\n D⁺.ν∇u = n⁻ * kinematic_stress(shallow.problem, A⁻.y, 1000)'\n # applies windstress for now, will be fixed in a later PR\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit/Communication.jl", "content": "@inline function initialize_states!(\n baroclinic::HBModel{C},\n barotropic::SWModel{C},\n model_bc,\n model_bt,\n state_bc,\n state_bt,\n) where {C <: Coupled}\n model_bc.state_auxiliary.ΔGᵘ .= -0\n\n return nothing\nend\n\n@inline function tendency_from_slow_to_fast!(\n baroclinic::HBModel{C},\n barotropic::SWModel{C},\n model_bc,\n model_bt,\n state_bc,\n state_bt,\n forcing_tendency,\n) where {C <: Coupled}\n FT = eltype(state_bc)\n info = basic_grid_info(model_bc)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n #### integrate the tendency\n model_int = model_bc.modeldata.integral_model\n integral = model_int.balance_law\n update_auxiliary_state!(model_int, integral, forcing_tendency, 0)\n\n ### properly shape MPIStateArrays\n num_aux_int = number_states(integral, Auxiliary())\n data_int = model_int.state_auxiliary.data\n data_int = reshape(data_int, Nqh, Nqk, num_aux_int, nelemv, nelemh)\n\n num_aux_bt = number_states(barotropic, Auxiliary())\n data_bt = model_bt.state_auxiliary.data\n data_bt = reshape(data_bt, Nqh, num_aux_bt, nelemh)\n\n num_aux_bc = number_states(baroclinic, Auxiliary())\n data_bc = model_bc.state_auxiliary.data\n data_bc = reshape(data_bc, Nqh, Nqk, num_aux_bc, nelemv, nelemh)\n\n ### get vars indices\n index_∫du = varsindex(vars_state(integral, Auxiliary(), FT), :(∫x))\n index_Gᵁ = varsindex(vars_state(barotropic, Auxiliary(), FT), :Gᵁ)\n index_ΔGᵘ = varsindex(vars_state(baroclinic, Auxiliary(), FT), :ΔGᵘ)\n\n ### get top value (=integral over full depth) of ∫du\n ∫du = @view data_int[:, end, index_∫du, end, 1:nrealelemh]\n\n ### copy into Gᵁ of barotropic model\n Gᵁ = @view data_bt[:, index_Gᵁ, 1:nrealelemh]\n Gᵁ .= ∫du\n\n ### get top value (=integral over full depth) of ∫du\n ∫du = @view data_int[:, end:end, index_∫du, end:end, 1:nrealelemh]\n\n ### save vertically averaged tendency to remove from 3D tendency\n ### need to reshape for the broadcast\n ΔGᵘ = @view data_bc[:, :, index_ΔGᵘ, :, 1:nrealelemh]\n ΔGᵘ .-= ∫du / baroclinic.problem.H\n\n return nothing\nend\n\n@inline function cummulate_fast_solution!(\n baroclinic::HBModel{C},\n barotropic::SWModel{C},\n model_bt,\n state_bt,\n fast_time,\n fast_dt,\n substep,\n) where {C <: Coupled}\n return nothing\nend\n\n@inline function reconcile_from_fast_to_slow!(\n baroclinic::HBModel{C},\n barotropic::SWModel{C},\n model_bc,\n model_bt,\n state_bc,\n state_bt,\n) where {C <: Coupled}\n FT = eltype(state_bc)\n info = basic_grid_info(model_bc)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n ### integrate the horizontal velocity\n model_int = model_bc.modeldata.integral_model\n integral = model_int.balance_law\n update_auxiliary_state!(model_int, integral, state_bc, 0)\n\n ### properly shape MPIStateArrays\n num_aux_int = number_states(integral, Auxiliary())\n data_int = model_int.state_auxiliary.data\n data_int = reshape(data_int, Nqh, Nqk, num_aux_int, nelemv, nelemh)\n\n num_aux_bt = number_states(barotropic, Auxiliary())\n data_bt_aux = model_bt.state_auxiliary.data\n data_bt_aux = reshape(data_bt_aux, Nqh, num_aux_bt, nelemh)\n\n num_state_bt = number_states(barotropic, Prognostic())\n data_bt_state = state_bt.data\n data_bt_state = reshape(data_bt_state, Nqh, num_state_bt, nelemh)\n\n num_state_bc = number_states(baroclinic, Prognostic())\n data_bc_state = state_bc.data\n data_bc_state =\n reshape(data_bc_state, Nqh, Nqk, num_state_bc, nelemv, nelemh)\n\n ### get vars indices\n index_∫u = varsindex(vars_state(integral, Auxiliary(), FT), :(∫x))\n index_Δu = varsindex(vars_state(barotropic, Auxiliary(), FT), :Δu)\n index_U = varsindex(vars_state(barotropic, Prognostic(), FT), :U)\n index_u = varsindex(vars_state(baroclinic, Prognostic(), FT), :u)\n index_η_3D = varsindex(vars_state(baroclinic, Prognostic(), FT), :η)\n index_η_2D = varsindex(vars_state(barotropic, Prognostic(), FT), :η)\n\n ### get top value (=integral over full depth)\n ∫u = @view data_int[:, end, index_∫u, end, 1:nrealelemh]\n\n ### Δu is a place holder for 1/H * (Ū - ∫u)\n Δu = @view data_bt_aux[:, index_Δu, 1:nrealelemh]\n U = @view data_bt_state[:, index_U, 1:nrealelemh]\n Δu .= 1 / baroclinic.problem.H * (U - ∫u)\n\n ### copy the 2D contribution down the 3D solution\n ### need to reshape for the broadcast\n data_bt_aux = reshape(data_bt_aux, Nqh, 1, num_aux_bt, 1, nelemh)\n Δu = @view data_bt_aux[:, :, index_Δu, :, 1:nrealelemh]\n u = @view data_bc_state[:, :, index_u, :, 1:nrealelemh]\n u .+= Δu\n\n ### copy η from barotropic mode to baroclinic mode\n ### need to reshape for the broadcast\n data_bt_state = reshape(data_bt_state, Nqh, 1, num_state_bt, 1, nelemh)\n η_2D = @view data_bt_state[:, :, index_η_2D, :, 1:nrealelemh]\n η_3D = @view data_bc_state[:, :, index_η_3D, :, 1:nrealelemh]\n η_3D .= η_2D\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit/HydrostaticBoussinesqCoupling.jl", "content": "using ..Ocean: coriolis_parameter\nusing ..HydrostaticBoussinesq\nusing ...DGMethods\n\nimport ...BalanceLaws\nimport ..HydrostaticBoussinesq:\n viscosity_tensor,\n coriolis_force!,\n velocity_gradient_argument!,\n velocity_gradient_flux!,\n hydrostatic_pressure!,\n compute_flow_deviation!\n\n@inline function velocity_gradient_argument!(m::HBModel, ::Coupled, G, Q, A, t)\n G.∇u = Q.u\n G.∇uᵈ = A.uᵈ\n\n return nothing\nend\n\n@inline function velocity_gradient_flux!(m::HBModel, ::Coupled, D, G, Q, A, t)\n ν = viscosity_tensor(m)\n ∇u = @SMatrix [\n G.∇uᵈ[1, 1] G.∇uᵈ[1, 2]\n G.∇uᵈ[2, 1] G.∇uᵈ[2, 2]\n G.∇u[3, 1] G.∇u[3, 2]\n ]\n D.ν∇u = -ν * ∇u\n\n return nothing\nend\n\n@inline hydrostatic_pressure!(::HBModel, ::Coupled, _...) = nothing\n\n@inline function coriolis_force!(m::HBModel, ::Coupled, S, Q, A, t)\n # f × u\n f = coriolis_parameter(m, m.problem, A.y)\n uᵈ, vᵈ = A.uᵈ # Horizontal components of velocity\n S.u -= @SVector [-f * vᵈ, f * uᵈ]\n\n return nothing\nend\n\n# Compute Horizontal Flow deviation from vertical mean\n@inline function compute_flow_deviation!(dg, m::HBModel, ::Coupled, Q, t)\n FT = eltype(Q)\n info = basic_grid_info(dg)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n #### integrate the tendency\n model_int = dg.modeldata.integral_model\n integral = model_int.balance_law\n update_auxiliary_state!(model_int, integral, Q, 0)\n\n ### properly shape MPIStateArrays\n num_int = number_states(integral, Auxiliary())\n data_int = model_int.state_auxiliary.data\n data_int = reshape(data_int, Nqh, Nqk, num_int, nelemv, nelemh)\n\n num_aux = number_states(m, Auxiliary())\n data_aux = dg.state_auxiliary.data\n data_aux = reshape(data_aux, Nqh, Nqk, num_aux, nelemv, nelemh)\n\n num_state = number_states(m, Prognostic())\n data_state = reshape(Q.data, Nqh, Nqk, num_state, nelemv, nelemh)\n\n ### get vars indices\n index_∫u = varsindex(vars_state(integral, Auxiliary(), FT), :(∫x))\n index_uᵈ = varsindex(vars_state(m, Auxiliary(), FT), :uᵈ)\n index_u = varsindex(vars_state(m, Prognostic(), FT), :u)\n\n ### get top value (=integral over full depth)\n ∫u = @view data_int[:, end:end, index_∫u, end:end, 1:nrealelemh]\n uᵈ = @view data_aux[:, :, index_uᵈ, :, 1:nrealelemh]\n u = @view data_state[:, :, index_u, :, 1:nrealelemh]\n\n ## make a copy of horizontal velocity\n ## and remove vertical mean velocity\n uᵈ .= u\n uᵈ .-= ∫u / m.problem.H\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit/ShallowWaterCoupling.jl", "content": "import ..ShallowWater: forcing_term!\n\n@inline function forcing_term!(::SWModel, ::Coupled, S, Q, A, t)\n S.U += A.Gᵁ\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit/SplitExplicitModel.jl", "content": "module SplitExplicit\n\nusing StaticArrays\n\nusing ..Ocean\nusing ..HydrostaticBoussinesq\nusing ..ShallowWater\n\nusing ...VariableTemplates\nusing ...MPIStateArrays\nusing ...Mesh.Geometry\nusing ...DGMethods\nusing ...BalanceLaws\n\nimport ...BalanceLaws:\n initialize_states!,\n tendency_from_slow_to_fast!,\n cummulate_fast_solution!,\n reconcile_from_fast_to_slow!,\n boundary_conditions\n\nHBModel = HydrostaticBoussinesqModel\nSWModel = ShallowWaterModel\n\nfunction initialize_states!(\n ::HBModel{C},\n ::SWModel{C},\n _...,\n) where {C <: Uncoupled}\n return nothing\nend\nfunction tendency_from_slow_to_fast!(\n ::HBModel{C},\n ::SWModel{C},\n _...,\n) where {C <: Uncoupled}\n return nothing\nend\nfunction cummulate_fast_solution!(\n ::HBModel{C},\n ::SWModel{C},\n _...,\n) where {C <: Uncoupled}\n return nothing\nend\nfunction reconcile_from_fast_to_slow!(\n ::HBModel{C},\n ::SWModel{C},\n _...,\n) where {C <: Uncoupled}\n return nothing\nend\n\ninclude(\"VerticalIntegralModel.jl\")\ninclude(\"Communication.jl\")\ninclude(\"ShallowWaterCoupling.jl\")\ninclude(\"HydrostaticBoussinesqCoupling.jl\")\n\nend\n" }, { "path": "src/Ocean/SplitExplicit/VerticalIntegralModel.jl", "content": "import ...BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n update_auxiliary_state!,\n integral_load_auxiliary_state!,\n integral_set_auxiliary_state!\n\nstruct VerticalIntegralModel{M} <: BalanceLaw\n ocean::M\n function VerticalIntegralModel(ocean::M) where {M}\n return new{M}(ocean)\n end\nend\n\nvars_state(tm::VerticalIntegralModel, st::Prognostic, FT) =\n vars_state(tm.ocean, st, FT)\n\nfunction vars_state(m::VerticalIntegralModel, ::Auxiliary, T)\n @vars begin\n ∫x::SVector{2, T}\n end\nend\n\ninit_state_auxiliary!(\n tm::VerticalIntegralModel,\n A::MPIStateArray,\n grid,\n direction,\n) = nothing\n\nfunction vars_state(m::VerticalIntegralModel, ::UpwardIntegrals, T)\n @vars begin\n ∫x::SVector{2, T}\n end\nend\n\n@inline function integral_load_auxiliary_state!(\n m::VerticalIntegralModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.∫x = A.∫x\n\n return nothing\nend\n\n@inline function integral_set_auxiliary_state!(\n m::VerticalIntegralModel,\n A::Vars,\n I::Vars,\n)\n A.∫x = I.∫x\n\n return nothing\nend\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n tm::VerticalIntegralModel,\n x::MPIStateArray,\n t::Real,\n)\n A = dg.state_auxiliary\n\n # copy tendency vector to aux state for integration\n function f!(::VerticalIntegralModel, x, A, t)\n @inbounds begin\n A.∫x = @SVector [x.u[1], x.u[2]]\n end\n\n return nothing\n end\n update_auxiliary_state!(f!, dg, tm, x, t)\n\n # compute integral for Gᵁ\n indefinite_stack_integral!(dg, tm, x, A, t) # bottom -> top\n\n return true\nend\n" }, { "path": "src/Ocean/SplitExplicit01/BarotropicModel.jl", "content": "struct BarotropicModel{M} <: AbstractOceanModel\n baroclinic::M\n function BarotropicModel(baroclinic::M) where {M}\n return new{M}(baroclinic)\n end\nend\n\nparameter_set(m::BarotropicModel) = parameter_set(m.baroclinic)\n\nfunction vars_state(m::BarotropicModel, ::Prognostic, T)\n @vars begin\n U::SVector{2, T}\n η::T\n end\nend\n\nfunction init_state_prognostic!(\n m::BarotropicModel,\n Q::Vars,\n A::Vars,\n localgeo,\n t,\n)\n Q.U = @SVector [-0, -0]\n Q.η = -0\n return nothing\nend\n\nfunction vars_state(m::BarotropicModel, ::Auxiliary, T)\n @vars begin\n Gᵁ::SVector{2, T} # integral of baroclinic tendency\n U_c::SVector{2, T} # cumulate U value over fast time-steps\n η_c::T # cumulate η value over fast time-steps\n U_s::SVector{2, T} # starting U field value\n η_s::T # starting η field value\n Δu::SVector{2, T} # reconciliation adjustment to u, Δu = 1/H * (U_averaged - ∫u)\n η_diag::T # η from baroclinic model (for diagnostic)\n Δη::T # diagnostic difference: η_barotropic - η_baroclinic\n y::T # y-coordinate of grid\n end\nend\n\nfunction init_state_auxiliary!(\n m::BarotropicModel,\n state_aux::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n m,\n (m, A, tmp, geom) -> ocean_init_aux!(m, m.baroclinic.problem, A, geom),\n state_aux,\n grid,\n direction,\n )\nend\n\nfunction vars_state(m::BarotropicModel, ::Gradient, T)\n @vars begin\n U::SVector{2, T}\n end\nend\n\n@inline function compute_gradient_argument!(\n m::BarotropicModel,\n G::Vars,\n Q::Vars,\n A,\n t,\n)\n G.U = Q.U\n return nothing\nend\n\nfunction vars_state(m::BarotropicModel, ::GradientFlux, T)\n @vars begin\n ν∇U::SMatrix{3, 2, T, 6}\n end\nend\n\n@inline function compute_gradient_flux!(\n m::BarotropicModel,\n D::Vars,\n G::Grad,\n Q::Vars,\n A::Vars,\n t,\n)\n ν = viscosity_tensor(m)\n D.ν∇U = -ν * G.U\n\n return nothing\nend\n\n@inline function viscosity_tensor(bm::BarotropicModel)\n m = bm.baroclinic\n return Diagonal(@SVector [m.νʰ, m.νʰ, 0])\nend\n\nvars_state(m::BarotropicModel, ::UpwardIntegrals, T) = @vars()\nvars_state(m::BarotropicModel, ::DownwardIntegrals, T) = @vars()\n\n@inline function flux_first_order!(\n m::BarotropicModel,\n F::Grad,\n Q::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n @inbounds begin\n U = @SVector [Q.U[1], Q.U[2], 0]\n η = Q.η\n H = m.baroclinic.problem.H\n g = m.baroclinic.grav\n Iʰ = @SMatrix [\n 1 0\n 0 1\n 0 0\n ]\n\n F.η += U\n F.U += g * H * η * Iʰ\n end\nend\n\n@inline function flux_second_order!(\n m::BarotropicModel,\n F::Grad,\n Q::Vars,\n D::Vars,\n HD::Vars,\n A::Vars,\n t::Real,\n)\n # numerical diffusivity for stability\n F.U += D.ν∇U\n\n return nothing\nend\n\n@inline function source!(\n m::BarotropicModel,\n S::Vars,\n Q::Vars,\n D::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n @inbounds begin\n U = Q.U\n\n # f × u\n f = coriolis_force(m.baroclinic, A.y)\n S.U -= @SVector [-f * U[2], f * U[1]]\n\n # vertically integrated baroclinic model tendency\n S.U += A.Gᵁ\n end\nend\n\n@inline wavespeed(m::BarotropicModel, n⁻, _...) =\n abs(SVector(m.baroclinic.cʰ, m.baroclinic.cʰ, m.baroclinic.cᶻ)' * n⁻)\n\n# We want not have jump penalties on η (since not a flux variable)\nfunction update_penalty!(\n ::RusanovNumericalFlux,\n ::BarotropicModel,\n n⁻,\n λ,\n ΔQ::Vars,\n Q⁻,\n A⁻,\n Q⁺,\n A⁺,\n t,\n)\n ΔQ.η = -0\n\n return nothing\nend\n\nboundary_conditions(bm::BarotropicModel) =\n (bm.baroclinic.problem.boundary_conditions[1],)\n\n\n\"\"\"\n boundary_state!(nf, bc, ::BarotropicModel, args...)\n\napplies boundary conditions for this model\ndispatches to a function in OceanBoundaryConditions.jl based on BC type defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n@inline function boundary_state!(nf, bc, bm::BarotropicModel, args...)\n return ocean_model_boundary!(bm, bc, nf, args...)\nend\n" }, { "path": "src/Ocean/SplitExplicit01/Communication.jl", "content": "import ...BalanceLaws:\n tendency_from_slow_to_fast!,\n cummulate_fast_solution!,\n reconcile_from_fast_to_slow!\n\nusing Printf\n\n@inline function set_fast_for_stepping!(\n slow::OceanModel,\n fast::BarotropicModel,\n dgFast,\n Qfast,\n S_fast,\n slow_dt,\n rkC,\n rkW,\n s,\n nStages,\n fast_time_rec,\n fast_steps,\n)\n FT = typeof(slow_dt)\n\n #- inverse ratio of additional fast time steps (for weighted average)\n # --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n add = slow.add_fast_substeps\n\n #- set time-step\n # Warning: only make sense for LS3NRK33Heuns\n # where 12 is lowest common mutiple (LCM) of all RK-Coeff inverse\n fast_dt = fast_time_rec[1]\n steps = fast_dt > 0 ? ceil(Int, slow_dt / fast_dt / FT(12)) : 1\n ntsFull = 12 * steps\n fast_dt = slow_dt / ntsFull\n add = add > 0 ? floor(Int, ntsFull / add) : 0\n\n #- time to start fast time-stepping (fast_time_rec[3]) for this stage:\n if s == nStages\n # Warning: only works with few RK-scheme such as LS3NRK33Heuns\n fast_time_rec[3] = rkW[1]\n fract_dt = (1 - fast_time_rec[3])\n fast_time_rec[3] *= slow_dt\n fast_steps[2] = 1\n else\n fast_time_rec[3] = 0.0\n fract_dt = rkC[s + 1] - fast_time_rec[3]\n fast_steps[2] = 0\n end\n\n #- set number of sub-steps we need\n # will time-average fast over: fast_steps[1] , fast_steps[3]\n # centered on fract_dt*slow_dt which corresponds to advance in time of slow\n steps = ceil(Int, fract_dt * slow_dt / fast_dt)\n add = min(add, steps - 1)\n fast_steps[1] = steps - add\n fast_steps[3] = steps + add\n fast_time_rec[1] = fract_dt * slow_dt / steps\n\n #- select which fast time-step (fast_steps[2]) solution to save for next time-step\n # Warning: only works with few RK-scheme such as LS3NRK33Heuns\n fast_steps[2] *= steps\n if s == 1\n fast_steps[2] = round(Int, ntsFull * rkW[1])\n end\n # @printf(\"Update @ s= %i : frac_dt = %.6f , dt_fast = %.1f , steps= %i , add= %i\\n\",\n # s, fract_dt, fast_time_rec[1], steps, add)\n # println(\" fast_time_rec = \",fast_time_rec)\n # println(\" fast_steps = \",fast_steps)\n\n # set starting point for fast-state solution\n # Warning: only works with few RK-scheme such as LS3NRK33Heuns\n if s == 1\n S_fast.η .= Qfast.η\n S_fast.U .= Qfast.U\n elseif s == nStages\n Qfast.η .= dgFast.state_auxiliary.η_s\n Qfast.U .= dgFast.state_auxiliary.U_s\n else\n Qfast.η .= S_fast.η\n Qfast.U .= S_fast.U\n end\n\n # initialise cumulative arrays\n fast_time_rec[2] = 0.0\n dgFast.state_auxiliary.η_c .= -0\n dgFast.state_auxiliary.U_c .= (@SVector [-0, -0])'\n\n return nothing\nend\n\n@inline function initialize_fast_state!(\n slow::OceanModel,\n fast::BarotropicModel,\n dgSlow,\n dgFast,\n Qslow,\n Qfast,\n slow_dt,\n fast_time_rec,\n fast_steps;\n firstStage = false,\n)\n\n #- inverse ratio of additional fast time steps (for weighted average)\n # --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n add = slow.add_fast_substeps\n\n #- set time-step and number of sub-steps we need\n # will time-average fast over: fast_steps[1] , fast_steps[3]\n # centered on fast_steps[2] which corresponds to advance in time of slow\n fast_dt = fast_time_rec[1]\n if add == 0\n steps = fast_dt > 0 ? ceil(Int, slow_dt / fast_dt) : 1\n fast_steps[1:3] = [1 1 1] * steps\n else\n steps = fast_dt > 0 ? ceil(Int, slow_dt / fast_dt / add) : 1\n fast_steps[2] = add * steps\n fast_steps[1] = (add - 1) * steps\n fast_steps[3] = (add + 1) * steps\n end\n fast_time_rec[1] = slow_dt / fast_steps[2]\n fast_time_rec[2] = 0.0\n # @printf(\"Update: frac_dt = %.1f , dt_fast = %.1f , nsubsteps= %i\\n\",\n # slow_dt,fast_time_rec[1],fast_steps[3])\n # println(\" fast_steps = \",fast_steps)\n\n dgFast.state_auxiliary.η_c .= -0\n dgFast.state_auxiliary.U_c .= (@SVector [-0, -0])'\n\n # set fast-state to previously stored value\n if !firstStage\n Qfast.η .= dgFast.state_auxiliary.η_s\n Qfast.U .= dgFast.state_auxiliary.U_s\n end\n\n return nothing\nend\n\n@inline function initialize_adjustment!(\n slow::OceanModel,\n fast::BarotropicModel,\n dgSlow,\n dgFast,\n Qslow,\n Qfast,\n)\n ## reset tendency adjustment before calling Baroclinic Model\n dgSlow.state_auxiliary.ΔGu .= 0\n\n return nothing\nend\n\n@inline function tendency_from_slow_to_fast!(\n slow::OceanModel,\n fast::BarotropicModel,\n dgSlow,\n dgFast,\n Qslow,\n Qfast,\n dQslow2fast,\n)\n FT = eltype(Qslow)\n\n # integrate the tendency\n tendency_dg = dgSlow.modeldata.tendency_dg\n tend = tendency_dg.balance_law\n grid = dgSlow.grid\n elems = grid.topology.elems\n update_auxiliary_state!(tendency_dg, tend, dQslow2fast, 0, elems)\n\n info = basic_grid_info(dgSlow)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n ## get top value (=integral over full depth) of ∫du\n nb_aux_tnd = number_states(tend, Auxiliary())\n data_tnd = reshape(\n tendency_dg.state_auxiliary.data,\n Nqh,\n Nqk,\n nb_aux_tnd,\n nelemv,\n nelemh,\n )\n index_∫du = varsindex(vars_state(tend, Auxiliary(), FT), :∫du)\n flat_∫du = @view data_tnd[:, end, index_∫du, end, 1:nrealelemh]\n\n ## copy into Gᵁ of dgFast\n nb_aux_fst = number_states(fast, Auxiliary())\n data_fst = reshape(dgFast.state_auxiliary.data, Nqh, nb_aux_fst, nelemh)\n index_Gᵁ = varsindex(vars_state(fast, Auxiliary(), FT), :Gᵁ)\n boxy_Gᵁ = @view data_fst[:, index_Gᵁ, 1:nrealelemh]\n boxy_Gᵁ .= flat_∫du\n\n ## scale by -1/H and copy back to ΔGu\n # note: since tendency_dg.state_auxiliary.∫du is not used after this, could be\n # re-used to store a 3-D copy of \"-Gu\"\n nb_aux_slw = number_states(slow, Auxiliary())\n data_slw = reshape(\n dgSlow.state_auxiliary.data,\n Nqh,\n Nqk,\n nb_aux_slw,\n nelemv,\n nelemh,\n )\n index_ΔGu = varsindex(vars_state(slow, Auxiliary(), FT), :ΔGu)\n boxy_ΔGu = @view data_slw[:, :, index_ΔGu, :, 1:nrealelemh]\n boxy_∫du = @view data_tnd[:, end:end, index_∫du, end:end, 1:nrealelemh]\n boxy_ΔGu .= -boxy_∫du / slow.problem.H\n\n return nothing\nend\n\n@inline function cummulate_fast_solution!(\n fast::BarotropicModel,\n dgFast,\n Qfast,\n fast_time,\n fast_dt,\n substep,\n fast_steps,\n fast_time_rec,\n)\n #- might want to use some of the weighting factors: weights_η & weights_U\n #- should account for case where fast_dt < fast.param.dt\n\n # cumulate Fast solution:\n if substep >= fast_steps[1]\n dgFast.state_auxiliary.U_c .+= Qfast.U\n dgFast.state_auxiliary.η_c .+= Qfast.η\n fast_time_rec[2] += 1.0\n end\n\n # save mid-point solution to start from the next time-step\n if substep == fast_steps[2]\n dgFast.state_auxiliary.U_s .= Qfast.U\n dgFast.state_auxiliary.η_s .= Qfast.η\n end\n\n return nothing\nend\n\n@inline function reconcile_from_fast_to_slow!(\n slow::OceanModel,\n fast::BarotropicModel,\n dgSlow,\n dgFast,\n Qslow,\n Qfast,\n fast_time_rec;\n lastStage = false,\n)\n FT = eltype(Qslow)\n info = basic_grid_info(dgSlow)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n grid = dgSlow.grid\n elems = grid.topology.elems\n\n # need to calculate int_u using integral kernels\n # u_slow := u_slow + (1/H) * (u_fast - \\int_{-H}^{0} u_slow)\n\n ## get time weighted averaged out of cumulative arrays\n dgFast.state_auxiliary.U_c .*= 1 / fast_time_rec[2]\n dgFast.state_auxiliary.η_c .*= 1 / fast_time_rec[2]\n # @printf(\" reconcile_from_fast_to_slow! @ s= %i : time_Count = %6.3f\\n\",\n # 0, fast_time_rec[2])\n # # s, fast_time_rec[2])\n\n ## Compute: \\int_{-H}^{0} u_slow\n\n # integrate vertically horizontal velocity\n flowintegral_dg = dgSlow.modeldata.flowintegral_dg\n flowint = flowintegral_dg.balance_law\n update_auxiliary_state!(flowintegral_dg, flowint, Qslow, 0, elems)\n\n # get top value (=integral over full depth)\n nb_aux_flw = number_states(flowint, Auxiliary())\n data_flw = reshape(\n flowintegral_dg.state_auxiliary.data,\n Nqh,\n Nqk,\n nb_aux_flw,\n nelemv,\n nelemh,\n )\n index_∫u = varsindex(vars_state(flowint, Auxiliary(), FT), :∫u)\n flat_∫u = @view data_flw[:, end, index_∫u, end, 1:nrealelemh]\n\n ## substract ∫u from U and divide by H\n\n # Δu is a place holder for 1/H * (Ū - ∫u)\n Δu = dgFast.state_auxiliary.Δu\n Δu .= dgFast.state_auxiliary.U_c\n\n nb_aux_fst = number_states(fast, Auxiliary())\n data_fst = reshape(dgFast.state_auxiliary.data, Nqh, nb_aux_fst, nelemh)\n index_Δu = varsindex(vars_state(fast, Auxiliary(), FT), :Δu)\n boxy_Δu = @view data_fst[:, index_Δu, 1:nrealelemh]\n boxy_Δu .-= flat_∫u\n boxy_Δu ./= slow.problem.H\n\n ## apply the 2D correction to the 3D solution\n nb_cons_slw = number_states(slow, Prognostic())\n data_slw = reshape(Qslow.data, Nqh, Nqk, nb_cons_slw, nelemv, nelemh)\n index_u = varsindex(vars_state(slow, Prognostic(), FT), :u)\n boxy_u = @view data_slw[:, :, index_u, :, 1:nrealelemh]\n boxy_u .+= reshape(boxy_Δu, Nqh, 1, 2, 1, nrealelemh)\n\n ## save Eta from 3D model into η_diag (aux var of 2D model)\n ## and store difference between η from Barotropic Model and η_diag\n ## Note: since 3D η is not used (just in output), only do this at last stage\n ## (save computation and get Δη diagnose over full time-step)\n if lastStage\n index_η = varsindex(vars_state(slow, Prognostic(), FT), :η)\n boxy_η_3D = @view data_slw[:, :, index_η, :, 1:nrealelemh]\n flat_η = @view data_slw[:, end, index_η, end, 1:nrealelemh]\n index_η_diag = varsindex(vars_state(fast, Auxiliary(), FT), :η_diag)\n boxy_η_diag = @view data_fst[:, index_η_diag, 1:nrealelemh]\n boxy_η_diag .= flat_η\n dgFast.state_auxiliary.Δη .=\n dgFast.state_auxiliary.η_c - dgFast.state_auxiliary.η_diag\n\n ## copy 2D model Eta over to 3D model\n index_η_c = varsindex(vars_state(fast, Auxiliary(), FT), :η_c)\n boxy_η_2D = @view data_fst[:, index_η_c, 1:nrealelemh]\n boxy_η_3D .= reshape(boxy_η_2D, Nqh, 1, 1, 1, nrealelemh)\n\n # reset fast-state to end of time-step value\n Qfast.η .= dgFast.state_auxiliary.η_s\n Qfast.U .= dgFast.state_auxiliary.U_s\n end\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit01/Continuity3dModel.jl", "content": "struct Continuity3dModel{M} <: AbstractOceanModel\n ocean::M\n function Continuity3dModel(ocean::M) where {M}\n return new{M}(ocean)\n end\nend\nvars_state(cm::Continuity3dModel, ::Prognostic, FT) =\n vars_state(cm.ocean, Prognostic(), FT)\n\n# Continuity3dModel is used to compute the horizontal divergence of u\n\nvars_state(cm::Continuity3dModel, ::Auxiliary, T) = @vars()\nvars_state(cm::Continuity3dModel, ::Gradient, T) = @vars()\nvars_state(cm::Continuity3dModel, ::GradientFlux, T) = @vars()\nvars_state(cm::Continuity3dModel, ::UpwardIntegrals, T) = @vars()\ninit_state_auxiliary!(cm::Continuity3dModel, _...) = nothing\ninit_state_prognostic!(cm::Continuity3dModel, _...) = nothing\n@inline flux_second_order!(cm::Continuity3dModel, _...) = nothing\n@inline source!(cm::Continuity3dModel, _...) = nothing\n@inline update_penalty!(::RusanovNumericalFlux, ::Continuity3dModel, _...) =\n nothing\n\n# This allows the balance law framework to compute the horizontal gradient of u\n# (which will be stored back in the field θ)\n@inline function flux_first_order!(\n m::Continuity3dModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n @inbounds begin\n u = state.u # Horizontal components of velocity\n v = @SVector [u[1], u[2], -0]\n\n # ∇ • (v)\n # Just using θ to store w = ∇h • u\n flux.θ += v\n end\n\n return nothing\nend\n\n# This is zero because when taking the horizontal gradient we're piggy-backing\n# on θ and want to ensure we do not use it's jump\n@inline wavespeed(cm::Continuity3dModel, n⁻, _...) = -zero(eltype(n⁻))\n\nboundary_conditions(cm::Continuity3dModel) = (\n cm.ocean.problem.boundary_conditions[1],\n cm.ocean.problem.boundary_conditions[1],\n cm.ocean.problem.boundary_conditions[1],\n)\n\nboundary_state!(\n ::Union{NumericalFluxGradient, NumericalFluxSecondOrder},\n bc,\n cm::Continuity3dModel,\n _...,\n) = nothing\n\n\"\"\"\n boundary_state!(nf, bc, ::Continuity3dModel, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in OceanBoundaryConditions.jl based on BC type defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n\n@inline function boundary_state!(\n nf::NumericalFluxFirstOrder,\n bc,\n cm::Continuity3dModel,\n args...,\n)\n return ocean_model_boundary!(cm, bc, nf, args...)\nend\n" }, { "path": "src/Ocean/SplitExplicit01/IVDCModel.jl", "content": "# Linear model equations, for split-explicit ocean model implicit vertical diffusion\n# convective adjustment step.\n#\n# In this version the operator is tweked to be the indentity for testing\n\n\"\"\"\n IVDCModel{M} <: BalanceLaw\n\n This code defines DG `BalanceLaw` terms for an operator, L, that is evaluated from iterative\n implicit solver to solve an equation of the form\n\n (L + 1/Δt) ϕ^{n+1} = ϕ^{n}/Δt\n\n where L is a vertical diffusion operator with a spatially varying diffusion\n coefficient.\n\n # Usage\n\n parent_model = OceanModel{FT}(prob...)\n linear_model = IVDCModel( parent_model )\n\n\"\"\"\n\n# Create a new child linear model instance, attached to whatever parent\n# BalanceLaw instantiates this.\n# (Not sure we need parent, but maybe we will get some parameters from it)\nstruct IVDCModel{M} <: AbstractOceanModel\n parent_om::M\n function IVDCModel(parent_om::M;) where {M}\n return new{M}(parent_om)\n end\nend\n\n\"\"\"\n Set model state variables and operators\n\"\"\"\n\n# State variable and initial value, just one for now, θ\n\nvars_state(m::IVDCModel, ::Prognostic, FT) = @vars(θ::FT)\n\nfunction init_state_prognostic!(m::IVDCModel, Q::Vars, A::Vars, localgeo, t)\n @inbounds begin\n Q.θ = -0\n end\n return nothing\nend\n\nvars_state(m::IVDCModel, ::Auxiliary, FT) = @vars(θ_init::FT)\nfunction init_state_auxiliary!(m::IVDCModel, A::Vars, _...)\n @inbounds begin\n A.θ_init = -0\n end\n return nothing\nend\n\n# Variables and operations used in differentiating first derivatives\n\nvars_state(m::IVDCModel, ::Gradient, FT) = @vars(∇θ::FT, ∇θ_init::FT,)\n\n@inline function compute_gradient_argument!(\n m::IVDCModel,\n G::Vars,\n Q::Vars,\n A,\n t,\n)\n G.∇θ = Q.θ\n G.∇θ_init = A.θ_init\n\n return nothing\nend\n\n# Variables and operations used in differentiating second derivatives\n\nvars_state(m::IVDCModel, ::GradientFlux, FT) = @vars(κ∇θ::SVector{3, FT})\n\n@inline function compute_gradient_flux!(\n m::IVDCModel,\n D::Vars,\n G::Grad,\n Q::Vars,\n A::Vars,\n t,\n)\n\n κ = diffusivity_tensor(m, G.∇θ_init[3])\n D.κ∇θ = -κ * G.∇θ\n\n return nothing\nend\n\n@inline function diffusivity_tensor(m::IVDCModel, ∂θ∂z)\n κᶻ = m.parent_om.κᶻ * 0.5\n κᶜ = m.parent_om.κᶜ\n ∂θ∂z < 0 ? κ = (@SVector [0, 0, κᶜ]) : κ = (@SVector [0, 0, κᶻ])\n # ∂θ∂z <= 1e-10 ? κ = (@SVector [0, 0, κᶜ]) : κ = (@SVector [0, 0, κᶻ])\n\n return Diagonal(-κ)\nend\n\n# Function to apply I to state variable\n\n@inline function source!(\n m::IVDCModel,\n S::Vars,\n Q::Vars,\n D::Vars,\n A::Vars,\n t,\n direction,\n)\n #ivdc_dt = m.ivdc_dt\n ivdc_dt = m.parent_om.ivdc_dt\n @inbounds begin\n S.θ = Q.θ / ivdc_dt\n # S.θ = 0\n end\n\n return nothing\nend\n\n## Numerical fluxes and boundaries\n\nfunction flux_first_order!(::IVDCModel, _...) end\n\nfunction flux_second_order!(\n ::IVDCModel,\n F::Grad,\n S::Vars,\n D::Vars,\n H::Vars,\n A::Vars,\n t,\n)\n F.θ += D.κ∇θ\n # F.θ = 0\n\nend\n\nfunction wavespeed(m::IVDCModel, n⁻, _...)\n C = abs(SVector(m.parent_om.cʰ, m.parent_om.cʰ, m.parent_om.cᶻ)' * n⁻)\n # C = abs(SVector(m.parent_om.cʰ, m.parent_om.cʰ, 50)' * n⁻)\n # C = abs(SVector(1, 1, 1)' * n⁻)\n ### C = abs(SVector(10, 10, 10)' * n⁻)\n # C = abs(SVector(50, 50, 50)' * n⁻)\n # C = abs(SVector( 0, 0, 0)' * n⁻)\n return C\nend\n\nboundary_conditions(m::IVDCModel) = boundary_conditions(m.parent_om)\nfunction boundary_state!(\n nf::Union{\n NumericalFluxFirstOrder,\n NumericalFluxGradient,\n CentralNumericalFluxGradient,\n },\n bc,\n m::IVDCModel,\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n t,\n _...,\n)\n Q⁺.θ = Q⁻.θ\n\n return nothing\nend\n\n### From - function numerical_boundary_flux_gradient! , DGMethods/NumericalFluxes.jl\n### boundary_state!(\n### numerical_flux,\n### balance_law,\n### state_conservative⁺,\n### state_auxiliary⁺,\n### normal_vector,\n### state_conservative⁻,\n### state_auxiliary⁻,\n### bctype,\n### t,\n### state1⁻,\n### aux1⁻,\n### )\n\nfunction boundary_state!(\n nf::Union{NumericalFluxSecondOrder, CentralNumericalFluxSecondOrder},\n bctype,\n m::IVDCModel,\n Q⁺,\n D⁺,\n HD⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n HD⁻,\n A⁻,\n t,\n _...,\n)\n Q⁺.θ = Q⁻.θ\n D⁺.κ∇θ = n⁻ * -0\n # D⁺.κ∇θ = n⁻ * -0 + 7000\n # D⁺.κ∇θ = -D⁻.κ∇θ\n\n return nothing\nend\n\n### boundary_state!(\n### numerical_flux,\n### balance_law,\n### state_conservative⁺,\n### state_gradient_flux⁺,\n### state_auxiliary⁺,\n### normal_vector,\n### state_conservative⁻,\n### state_gradient_flux⁻,\n### state_auxiliary⁻,\n### bctype,\n### t,\n### state1⁻,\n### diff1⁻,\n### aux1⁻,\n### )\n\n### boundary_flux_second_order!(\n### numerical_flux,\n### balance_law,\n### Grad{S}(flux),\n### state_conservative⁺,\n### state_gradient_flux⁺,\n### state_hyperdiffusive⁺,\n### state_auxiliary⁺,\n### normal_vector,\n### state_conservative⁻,\n### state_gradient_flux⁻,\n### state_hyperdiffusive⁻,\n### state_auxiliary⁻,\n### bctype,\n### t,\n### state1⁻,\n### diff1⁻,\n### aux1⁻,\n### )\n" }, { "path": "src/Ocean/SplitExplicit01/OceanBoundaryConditions.jl", "content": "abstract type OceanBoundaryCondition end\n\n\"\"\"\n Defining dummy structs to dispatch on for boundary conditions.\n\"\"\"\nstruct CoastlineFreeSlip <: OceanBoundaryCondition end\nstruct CoastlineNoSlip <: OceanBoundaryCondition end\nstruct OceanFloorFreeSlip <: OceanBoundaryCondition end\nstruct OceanFloorNoSlip <: OceanBoundaryCondition end\nstruct OceanFloorLinearDrag <: OceanBoundaryCondition end\nstruct OceanSurfaceNoStressNoForcing <: OceanBoundaryCondition end\nstruct OceanSurfaceStressNoForcing <: OceanBoundaryCondition end\nstruct OceanSurfaceNoStressForcing <: OceanBoundaryCondition end\nstruct OceanSurfaceStressForcing <: OceanBoundaryCondition end\n\n# these functions just trim off the extra arguments\nfunction ocean_model_boundary!(\n model::AbstractOceanModel,\n bc,\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n Q⁺,\n A⁺,\n n,\n Q⁻,\n A⁻,\n t,\n _...,\n)\n return ocean_boundary_state!(model, bc, nf, Q⁺, A⁺, n, Q⁻, A⁻, t)\nend\n\nfunction ocean_model_boundary!(\n model::AbstractOceanModel,\n bc,\n nf::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n HD⁺,\n A⁺,\n n,\n Q⁻,\n D⁻,\n HD⁻,\n A⁻,\n t,\n _...,\n)\n return ocean_boundary_state!(model, bc, nf, Q⁺, D⁺, A⁺, n, Q⁻, D⁻, A⁻, t)\nend\n\n\"\"\"\n CoastlineFreeSlip\n\napplies boundary condition ∇u = 0 and ∇θ = 0\n\"\"\"\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineFreeSlip, ::NumericalFluxFirstOrder)\n\napply free slip boundary conditions for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n u⁻ = Q⁻.u\n n = @SVector [n⁻[1], n⁻[2]]\n\n # Q⁺.u = u⁻ - 2 * (n⋅u⁻) * n\n Q⁺.u = u⁻ - 2 * (n ∘ u⁻) * n\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n U⁻ = Q⁻.U\n n = @SVector [n⁻[1], n⁻[2]]\n\n # Q⁺.U = U⁻ - 2 * (n⋅U⁻) * n\n Q⁺.U = U⁻ - 2 * (n ∘ U⁻) * n\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineFreeSlip, ::NumericalFluxGradient)\n\napply free slip boundary conditions for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n u⁻ = Q⁻.u\n ud⁻ = A⁻.u_d\n n = @SVector [n⁻[1], n⁻[2]]\n\n # Q⁺.u = u⁻ - (n⋅u⁻) * n\n Q⁺.u = u⁻ - (n ∘ u⁻) * n\n A⁺.u_d = ud⁻ - (n ∘ ud⁻) * n\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n U⁻ = Q⁻.U\n n = @SVector [n⁻[1], n⁻[2]]\n\n # Q⁺.U = U⁻ - (n⋅U⁻) * n\n Q⁺.U = U⁻ - (n ∘ U⁻) * n\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineFreeSlip, ::NumericalFluxSecondOrder)\n\napply free slip boundary conditions for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineFreeSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇U = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n CoastlineNoSlip\n\napplies boundary condition u = 0 and ∇θ = 0\n\"\"\"\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineNoSlip, ::NumericalFluxFirstOrder)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineNoSlip,\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n Q⁺.u = -Q⁻.u\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineNoSlip,\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n Q⁺.U = -Q⁻.U\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineNoSlip, ::NumericalFluxGradient)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineNoSlip,\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n FT = eltype(Q⁺)\n Q⁺.u = SVector(-zero(FT), -zero(FT))\n A⁺.u_d = SVector(-zero(FT), -zero(FT))\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineNoSlip,\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n FT = eltype(Q⁺)\n Q⁺.U = SVector(-zero(FT), -zero(FT))\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::CoastlineNoSlip, ::NumericalFluxSecondOrder)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::CoastlineNoSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇u = D⁻.ν∇u\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n@inline function ocean_boundary_state!(\n ::BarotropicModel,\n ::CoastlineNoSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇U = D⁻.ν∇U\n\n return nothing\nend\n\n\"\"\"\n OceanFloorFreeSlip\n\napplies boundary condition ∇u = 0 and ∇θ = 0\n\"\"\"\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanFloorFreeSlip, ::NumericalFluxFirstOrder)\n\napply free slip boundary conditions for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::OceanFloorFreeSlip,\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n A⁺.w = -A⁻.w\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanFloorFreeSlip, ::NumericalFluxGradient)\n\napply free slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::OceanFloorFreeSlip,\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n FT = eltype(Q⁺)\n A⁺.w = -zero(FT)\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanFloorFreeSlip, ::NumericalFluxSecondOrder)\n\napply free slip boundary conditions for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::OceanFloorFreeSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n OceanFloorNoSlip\n\napplies boundary condition u = 0 and ∇θ = 0\n\"\"\"\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::Union{OceanFloorNoSlip, OceanFloorLinearDrag}, ::NumericalFluxFirstOrder)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::Union{OceanFloorNoSlip, OceanFloorLinearDrag},\n ::NumericalFluxFirstOrder,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n Q⁺.u = -Q⁻.u\n A⁺.w = -A⁻.w\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::Union{OceanFloorNoSlip, OceanFloorLinearDrag}, ::NumericalFluxGradient)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::Union{OceanFloorNoSlip, OceanFloorLinearDrag},\n ::NumericalFluxGradient,\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n FT = eltype(Q⁺)\n Q⁺.u = SVector(-zero(FT), -zero(FT))\n A⁺.w = -zero(FT)\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanFloorNoSlip, ::NumericalFluxSecondOrder)\n\napply no slip boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::OceanFloorNoSlip,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇u = D⁻.ν∇u\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n OceanFloorLinearDrag\n\napplies boundary condition u = 0 with linear drag on viscous-flux and ∇θ = 0\n\"\"\"\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanFloorLinearDrag, ::NumericalFluxSecondOrder)\n\napply linear drag boundary condition for velocity\napply no penetration boundary for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n m::AbstractOceanModel,\n ::OceanFloorLinearDrag,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n u, v = Q⁻.u\n\n D⁺.ν∇u = -m.problem.Cd_lin * @SMatrix [-0 -0; -0 -0; u v]\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::Union{OceanSurface*}, ::Union{NumericalFluxFirstOrder, NumericalFluxGradient})\napplying neumann boundary conditions, so don't need to do anything for these numerical fluxes\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::Union{\n OceanSurfaceNoStressNoForcing,\n OceanSurfaceStressNoForcing,\n OceanSurfaceNoStressForcing,\n OceanSurfaceStressForcing,\n },\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n Q⁺,\n A⁺,\n n⁻,\n Q⁻,\n A⁻,\n t,\n)\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanSurfaceNoStressNoForcing, ::NumericalFluxSecondOrder)\n\napply no flux boundary condition for velocity\napply no flux boundary condition for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n ::AbstractOceanModel,\n ::OceanSurfaceNoStressNoForcing,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanSurfaceStressNoForcing, ::NumericalFluxSecondOrder)\n\napply wind-stress boundary condition for velocity\napply no flux boundary condition for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n m::AbstractOceanModel,\n ::OceanSurfaceStressNoForcing,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n τᶻ = velocity_flux(m.problem, A⁻.y, m.ρₒ)\n D⁺.ν∇u = n⁻ * (@SVector [-τᶻ, -0])'\n D⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanSurfaceNoStressForcing, ::NumericalFluxSecondOrder)\n\napply no flux boundary condition for velocity\napply forcing boundary condition for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n m::AbstractOceanModel,\n ::OceanSurfaceNoStressForcing,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n σᶻ = temperature_flux(m.problem, A⁻.y, Q⁻.θ)\n D⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n D⁺.κ∇θ = -n⁻ * σᶻ\n\n return nothing\nend\n\n\"\"\"\n ocean_boundary_state!(::AbstractOceanModel, ::OceanSurfaceStressForcing, ::NumericalFluxSecondOrder)\n\napply wind-stress boundary condition for velocity\napply forcing boundary condition for temperature\n\"\"\"\n@inline function ocean_boundary_state!(\n m::AbstractOceanModel,\n ::OceanSurfaceStressForcing,\n ::NumericalFluxSecondOrder,\n Q⁺,\n D⁺,\n A⁺,\n n⁻,\n Q⁻,\n D⁻,\n A⁻,\n t,\n)\n τᶻ = velocity_flux(m.problem, A⁻.y, m.ρₒ)\n σᶻ = temperature_flux(m.problem, A⁻.y, Q⁻.θ)\n D⁺.ν∇u = n⁻ * (@SVector [-τᶻ, -0])'\n D⁺.κ∇θ = -n⁻ * σᶻ\n\n return nothing\nend\n\n@inline velocity_flux(p::AbstractOceanProblem, y, ρ) =\n -(p.τₒ / ρ) * cos(y * π / p.Lʸ)\n\n@inline function temperature_flux(p::AbstractOceanProblem, y, θ)\n θʳ = p.θᴱ * (1 - y / p.Lʸ)\n return p.λʳ * (θʳ - θ)\nend\n" }, { "path": "src/Ocean/SplitExplicit01/OceanModel.jl", "content": "struct OceanModel{P, T} <: AbstractOceanModel\n problem::P\n grav::T\n ρₒ::T\n cʰ::T\n cᶻ::T\n add_fast_substeps::T\n numImplSteps::T\n ivdc_dt::T\n αᵀ::T\n νʰ::T\n νᶻ::T\n κʰ::T\n κᶻ::T\n κᶜ::T\n fₒ::T\n β::T\n function OceanModel{FT}(\n problem;\n grav = FT(10), # m/s^2\n ρₒ = FT(1000), # kg/m^3\n cʰ = FT(0), # m/s\n cᶻ = FT(0), # m/s\n add_fast_substeps = 0,\n numImplSteps = 0,\n ivdc_dt = FT(1),\n αᵀ = FT(2e-4), # 1/K\n νʰ = FT(5e3), # m^2/s\n νᶻ = FT(5e-3), # m^2/s\n κʰ = FT(1e3), # m^2/s\n κᶻ = FT(1e-4), # vertical diffusivity, m^2/s\n κᶜ = FT(1e-4), # convective adjustment vertical diffusivity, m^2/s\n fₒ = FT(1e-4), # Hz\n β = FT(1e-11), # Hz/m\n ) where {FT <: AbstractFloat}\n return new{typeof(problem), FT}(\n problem,\n grav,\n ρₒ,\n cʰ,\n cᶻ,\n add_fast_substeps,\n numImplSteps,\n ivdc_dt,\n αᵀ,\n νʰ,\n νᶻ,\n κʰ,\n κᶻ,\n κᶜ,\n fₒ,\n β,\n )\n end\nend\n\nfunction calculate_dt(grid, ::OceanModel, _...)\n #=\n minΔx = min_node_distance(grid, HorizontalDirection())\n minΔz = min_node_distance(grid, VerticalDirection())\n\n CFL_gravity = minΔx / model.cʰ\n CFL_diffusive = minΔz^2 / (1000 * model.κᶻ)\n CFL_viscous = minΔz^2 / model.νᶻ\n\n dt = 1 // 2 * minimum([CFL_gravity, CFL_diffusive, CFL_viscous])\n =#\n # FT = eltype(grid)\n # dt = FT(1)\n\n return nothing\nend\n\n\"\"\"\n OceanDGModel()\n\nhelper function to add required filtering\nnot used in the Driver+Config setup\n\"\"\"\nfunction OceanDGModel(\n bl::OceanModel,\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux;\n kwargs...,\n)\n N = polynomialorders(grid)\n Nvert = N[end]\n vert_filter = CutoffFilter(grid, Nvert - 1)\n exp_filter = ExponentialFilter(grid, 1, 8)\n\n flowintegral_dg = DGModel(\n FlowIntegralModel(bl),\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux,\n )\n\n tendency_dg = DGModel(\n TendencyIntegralModel(bl),\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux,\n )\n\n conti3d_dg = DGModel(\n Continuity3dModel(bl),\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux,\n )\n FT = eltype(grid)\n conti3d_Q = init_ode_state(conti3d_dg, FT(0); init_on_cpu = true)\n\n ivdc_dg = DGModel(\n IVDCModel(bl),\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux;\n direction = VerticalDirection(),\n )\n ivdc_Q = init_ode_state(ivdc_dg, FT(0); init_on_cpu = true) # Not sure this is needed since we set values later,\n # but we'll do it just in case!\n\n ivdc_RHS = init_ode_state(ivdc_dg, FT(0); init_on_cpu = true) # Not sure this is needed since we set values later,\n # but we'll do it just in case!\n\n ivdc_bgm_solver = BatchedGeneralizedMinimalResidual(\n ivdc_dg,\n ivdc_Q;\n max_subspace_size = 10,\n )\n\n modeldata = (\n vert_filter = vert_filter,\n exp_filter = exp_filter,\n flowintegral_dg = flowintegral_dg,\n tendency_dg = tendency_dg,\n conti3d_dg = conti3d_dg,\n conti3d_Q = conti3d_Q,\n ivdc_dg = ivdc_dg,\n ivdc_Q = ivdc_Q,\n ivdc_RHS = ivdc_RHS,\n ivdc_bgm_solver = ivdc_bgm_solver,\n )\n\n return DGModel(\n bl,\n grid,\n numfluxnondiff,\n numfluxdiff,\n gradnumflux;\n kwargs...,\n modeldata = modeldata,\n )\nend\n\nfunction vars_state(m::OceanModel, ::Prognostic, T)\n @vars begin\n u::SVector{2, T}\n η::T\n θ::T\n end\nend\n\nfunction init_state_prognostic!(m::OceanModel, Q::Vars, A::Vars, localgeo, t)\n return ocean_init_state!(m.problem, Q, A, localgeo, t)\nend\n\nfunction vars_state(m::OceanModel, ::Auxiliary, T)\n @vars begin\n w::T\n pkin::T # kinematic pressure: ∫(-g αᵀ θ)\n wz0::T # w at z=0\n u_d::SVector{2, T} # velocity deviation from vertical mean\n ΔGu::SVector{2, T}\n y::T # y-coordinate of the box\n end\nend\n\nfunction init_state_auxiliary!(\n m::OceanModel,\n state_aux::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n m,\n (m, A, tmp, geom) -> ocean_init_aux!(m, m.problem, A, geom),\n state_aux,\n grid,\n direction,\n )\nend\n\nfunction vars_state(m::OceanModel, ::Gradient, T)\n @vars begin\n u::SVector{2, T}\n ud::SVector{2, T}\n θ::T\n end\nend\n\n@inline function compute_gradient_argument!(\n m::OceanModel,\n G::Vars,\n Q::Vars,\n A,\n t,\n)\n G.u = Q.u\n G.ud = A.u_d\n G.θ = Q.θ\n\n return nothing\nend\n\nfunction vars_state(m::OceanModel, ::GradientFlux, T)\n @vars begin\n ν∇u::SMatrix{3, 2, T, 6}\n κ∇θ::SVector{3, T}\n end\nend\n\n@inline function compute_gradient_flux!(\n m::OceanModel,\n D::Vars,\n G::Grad,\n Q::Vars,\n A::Vars,\n t,\n)\n ν = viscosity_tensor(m)\n # D.ν∇u = ν * G.u\n D.ν∇u =\n -@SMatrix [\n m.νʰ*G.ud[1, 1] m.νʰ*G.ud[1, 2]\n m.νʰ*G.ud[2, 1] m.νʰ*G.ud[2, 2]\n m.νᶻ*G.u[3, 1] m.νᶻ*G.u[3, 2]\n ]\n\n κ = diffusivity_tensor(m, G.θ[3])\n D.κ∇θ = -κ * G.θ\n\n return nothing\nend\n\n@inline viscosity_tensor(m::OceanModel) = Diagonal(@SVector [m.νʰ, m.νʰ, m.νᶻ])\n\n@inline function diffusivity_tensor(m::OceanModel, ∂θ∂z)\n\n if m.numImplSteps > 0\n κ = (@SVector [m.κʰ, m.κʰ, m.κᶻ * 0.5])\n else\n ∂θ∂z < 0 ? κ = (@SVector [m.κʰ, m.κʰ, m.κᶜ]) :\n κ = (@SVector [m.κʰ, m.κʰ, m.κᶻ])\n end\n\n return Diagonal(κ)\nend\n\n\"\"\"\n vars_integral(::OceanModel)\n\nlocation to store integrands for bottom up integrals\n∇hu = the horizontal divegence of u, e.g. dw/dz\n\"\"\"\nfunction vars_state(m::OceanModel, ::UpwardIntegrals, T)\n @vars begin\n ∇hu::T\n buoy::T\n # ∫u::SVector{2, T}\n end\nend\n\n\"\"\"\n integral_load_auxiliary_state!(::OceanModel)\n\ncopy w to var_integral\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case OceanModel\nI -> array of integrand variables\nQ -> array of state variables\nA -> array of aux variables\n\"\"\"\n@inline function integral_load_auxiliary_state!(\n m::OceanModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.∇hu = A.w # borrow the w value from A...\n I.buoy = m.grav * m.αᵀ * Q.θ # buoyancy to integrate vertically from top (=reverse)\n # I.∫u = Q.u\n\n return nothing\nend\n\n\"\"\"\n integral_set_auxiliary_state!(::OceanModel)\n\ncopy integral results back out to aux\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case OceanModel\nA -> array of aux variables\nI -> array of integrand variables\n\"\"\"\n@inline function integral_set_auxiliary_state!(m::OceanModel, A::Vars, I::Vars)\n A.w = I.∇hu\n A.pkin = -I.buoy\n # A.∫u = I.∫u\n\n return nothing\nend\n\n\"\"\"\n vars_reverse_integral(::OceanModel)\n\nlocation to store integrands for top down integrals\nαᵀθ = density perturbation\n\"\"\"\nfunction vars_state(m::OceanModel, ::DownwardIntegrals, T)\n @vars begin\n buoy::T\n end\nend\n\n\"\"\"\n reverse_integral_load_auxiliary_state!(::OceanModel)\n\ncopy αᵀθ to var_reverse_integral\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case OceanModel\nI -> array of integrand variables\nA -> array of aux variables\n\"\"\"\n@inline function reverse_integral_load_auxiliary_state!(\n m::OceanModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.buoy = A.pkin\n\n return nothing\nend\n\n\"\"\"\n reverse_integral_set_auxiliary_state!(::OceanModel)\n\ncopy reverse integral results back out to aux\nthis computation is done pointwise at each nodal point\n\narguments:\nm -> model in this case OceanModel\nA -> array of aux variables\nI -> array of integrand variables\n\"\"\"\n@inline function reverse_integral_set_auxiliary_state!(\n m::OceanModel,\n A::Vars,\n I::Vars,\n)\n A.pkin = I.buoy\n\n return nothing\nend\n\n@inline function flux_first_order!(\n m::OceanModel,\n F::Grad,\n Q::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n @inbounds begin\n u = Q.u # Horizontal components of velocity\n θ = Q.θ\n w = A.w # vertical velocity\n pkin = A.pkin\n v = @SVector [u[1], u[2], w]\n Iʰ = @SMatrix [\n 1 -0\n -0 1\n -0 -0\n ]\n\n # ∇h • (g η)\n #- jmc: put back this term to check\n # η = Q.η\n # F.u += m.grav * η * Iʰ\n\n # ∇ • (u θ)\n F.θ += v * θ\n\n # ∇h • pkin\n F.u += pkin * Iʰ\n\n # ∇h • (v ⊗ u)\n # F.u += v * u'\n end\n\n return nothing\nend\n\n@inline function flux_second_order!(\n m::OceanModel,\n F::Grad,\n Q::Vars,\n D::Vars,\n HD::Vars,\n A::Vars,\n t::Real,\n)\n #- vertical viscosity only (horizontal fluxes in horizontal model)\n # F.u -= @SVector([0, 0, 1]) * D.ν∇u[3, :]'\n #- all 3 direction viscous flux for horizontal momentum tendency\n F.u += D.ν∇u\n\n F.θ += D.κ∇θ\n\n return nothing\nend\n\n@inline function source!(\n m::OceanModel,\n S::Vars,\n Q::Vars,\n D::Vars,\n A::Vars,\n t::Real,\n direction,\n)\n @inbounds begin\n u = Q.u # Horizontal components of velocity\n ud = A.u_d # Horizontal velocity deviation from vertical mean\n\n # f × u\n f = coriolis_force(m, A.y)\n # S.u -= @SVector [-f * u[2], f * u[1]]\n S.u -= @SVector [-f * ud[2], f * ud[1]]\n\n #- borotropic tendency adjustment\n S.u += A.ΔGu\n\n # switch this to S.η if you comment out the fast mode in MultistateMultirateRungeKutta\n S.η += A.wz0\n end\n\n return nothing\nend\n\n@inline coriolis_force(m::OceanModel, y) = m.fₒ + m.β * y\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n m::OceanModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n FT = eltype(Q)\n A = dg.state_auxiliary\n MD = dg.modeldata\n\n # `update_auxiliary_state!` gets called twice, once for the real elements\n # and once for the ghost elements. Only apply the filters to the real elems.\n if elems == dg.grid.topology.realelems\n # required to ensure that after integration velocity field is divergence free\n vert_filter = MD.vert_filter\n # Q[1] = u[1] = u, Q[2] = u[2] = v\n apply!(Q, (1, 2), dg.grid, vert_filter, direction = VerticalDirection())\n\n exp_filter = MD.exp_filter\n # Q[4] = θ\n apply!(Q, (4,), dg.grid, exp_filter, direction = VerticalDirection())\n end\n\n #----------\n # Compute Divergence of Horizontal Flow field using \"conti3d_dg\" DGmodel\n\n conti3d_dg = dg.modeldata.conti3d_dg\n # ct3d_Q = dg.modeldata.conti3d_Q\n # ct3d_dQ = similar(ct3d_Q)\n # fill!(ct3d_dQ, 0)\n #- Instead, use directly conti3d_Q to store dQ (since we will not update the state)\n ct3d_dQ = dg.modeldata.conti3d_Q\n\n ct3d_bl = conti3d_dg.balance_law\n\n # call \"conti3d_dg\" DGmodel\n # note: with \"increment = false\", just return tendency (no state update)\n p = nothing\n conti3d_dg(ct3d_dQ, Q, p, t; increment = false)\n\n # Copy from ct3d_dQ.θ which is realy ∇h•u into A.w (which will be integrated)\n function f!(::OceanModel, dQ, A, t)\n @inbounds begin\n A.w = dQ.θ\n end\n end\n update_auxiliary_state!(f!, dg, m, ct3d_dQ, t, elems)\n #----------\n\n info = basic_grid_info(dg)\n Nqh, Nqk = info.Nqh, info.Nqk\n nelemv, nelemh = info.nvertelem, info.nhorzelem\n nrealelemh = info.nhorzrealelem\n\n # compute integrals for w and pkin\n indefinite_stack_integral!(dg, m, Q, A, t, elems) # bottom -> top\n reverse_indefinite_stack_integral!(dg, m, Q, A, t, elems) # top -> bottom\n\n # copy down wz0\n # We are unable to use vars (ie A.w) for this because this operation will\n # return a SubArray, and adapt (used for broadcasting along reshaped arrays)\n # has a limited recursion depth for the types allowed.\n nb_aux_m = number_states(m, Auxiliary())\n data_m = reshape(A.data, Nqh, Nqk, nb_aux_m, nelemv, nelemh)\n\n # project w(z=0) down the stack\n index_w = varsindex(vars_state(m, Auxiliary(), FT), :w)\n index_wz0 = varsindex(vars_state(m, Auxiliary(), FT), :wz0)\n flat_wz0 = @view data_m[:, end:end, index_w, end:end, 1:nrealelemh]\n boxy_wz0 = @view data_m[:, :, index_wz0, :, 1:nrealelemh]\n boxy_wz0 .= flat_wz0\n\n # Compute Horizontal Flow deviation from vertical mean\n\n flowintegral_dg = dg.modeldata.flowintegral_dg\n flowint = flowintegral_dg.balance_law\n update_auxiliary_state!(flowintegral_dg, flowint, Q, 0, elems)\n\n ## get top value (=integral over full depth)\n nb_aux_flw = number_states(flowint, Auxiliary())\n data_flw = reshape(\n flowintegral_dg.state_auxiliary.data,\n Nqh,\n Nqk,\n nb_aux_flw,\n nelemv,\n nelemh,\n )\n index_∫u = varsindex(vars_state(flowint, Auxiliary(), FT), :∫u)\n\n flat_∫u = @view data_flw[:, end:end, index_∫u, end:end, 1:nrealelemh]\n\n ## make a copy of horizontal velocity\n A.u_d .= Q.u\n\n ## and remove vertical mean velocity\n index_ud = varsindex(vars_state(m, Auxiliary(), FT), :u_d)\n boxy_ud = @view data_m[:, :, index_ud, :, 1:nrealelemh]\n boxy_ud .-= flat_∫u / m.problem.H\n\n return true\nend\n\n@inline wavespeed(m::OceanModel, n⁻, _...) =\n abs(SVector(m.cʰ, m.cʰ, m.cᶻ)' * n⁻)\n\n# We want not have jump penalties on η (since not a flux variable)\nfunction update_penalty!(\n ::RusanovNumericalFlux,\n ::OceanModel,\n n⁻,\n λ,\n ΔQ::Vars,\n Q⁻,\n A⁻,\n Q⁺,\n A⁺,\n t,\n)\n ΔQ.η = -0\n\n return nothing\nend\n\nboundary_conditions(ocean::OceanModel) = ocean.problem.boundary_conditions\n\n\"\"\"\n boundary_state!(nf, bc, ::OceanModel, args...)\n\napplies boundary conditions for this model\ndispatches to a function in OceanBoundaryConditions.jl based on BC type defined by a problem such as SimpleBoxProblem.jl\n\"\"\"\n\n@inline function boundary_state!(nf, bc, ocean::OceanModel, args...)\n return ocean_model_boundary!(ocean, bc, nf, args...)\nend\n" }, { "path": "src/Ocean/SplitExplicit01/SplitExplicitLSRK2nMethod.jl", "content": "export SplitExplicitLSRK2nSolver\n\nusing KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\nusing StaticArrays\nusing ...SystemSolvers\nusing ...MPIStateArrays: array_device, realview\nusing ...GenericCallbacks\n\nusing ...ODESolvers:\n AbstractODESolver, LowStorageRungeKutta2N, update!, updatedt!, getdt\nimport ...ODESolvers: dostep!\n\nusing ...BalanceLaws:\n # initialize_fast_state!,\n # initialize_adjustment!,\n tendency_from_slow_to_fast!,\n cummulate_fast_solution!,\n reconcile_from_fast_to_slow!\n\nLSRK2N = LowStorageRungeKutta2N\n\n@doc \"\"\"\n SplitExplicitLSRK2nSolver(slow_solver, fast_solver; dt, t0 = 0, coupled = true)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q_fast} = f_fast(Q_fast, Q_slow, t)\n \\\\dot{Q_slow} = f_slow(Q_slow, Q_fast, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis method performs an operator splitting to timestep the Sea-Surface elevation\nand vertically averaged horizontal velocity of the model at a faster rate than\nthe full model, using LowStorageRungeKutta2N time-stepping.\n\n\"\"\" SplitExplicitLSRK2nSolver\nmutable struct SplitExplicitLSRK2nSolver{SS, FS, RT, MSA} <: AbstractODESolver\n \"slow solver\"\n slow_solver::SS\n \"fast solver\"\n fast_solver::FS\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"storage for transfer tendency\"\n dQ2fast::MSA\n\n function SplitExplicitLSRK2nSolver(\n slow_solver::LSRK2N,\n fast_solver,\n Q = nothing;\n dt = getdt(slow_solver),\n t0 = slow_solver.t,\n ) where {AT <: AbstractArray}\n SS = typeof(slow_solver)\n FS = typeof(fast_solver)\n RT = real(eltype(slow_solver.dQ))\n\n dQ2fast = similar(slow_solver.dQ)\n dQ2fast .= -0.0\n MSA = typeof(dQ2fast)\n return new{SS, FS, RT, MSA}(\n slow_solver,\n fast_solver,\n RT(dt),\n RT(t0),\n 0,\n dQ2fast,\n )\n end\nend\n\nfunction dostep!(\n Qvec,\n split::SplitExplicitLSRK2nSolver{SS},\n param,\n time::Real,\n) where {SS <: LSRK2N}\n slow = split.slow_solver\n fast = split.fast_solver\n\n Qslow = Qvec.slow\n Qfast = Qvec.fast\n\n dQslow = slow.dQ\n dQ2fast = split.dQ2fast\n\n slow_bl = slow.rhs!.balance_law\n fast_bl = fast.rhs!.balance_law\n\n groupsize = 256\n\n slow_dt = getdt(slow)\n fast_dt_in = getdt(fast)\n\n for slow_s in 1:length(slow.RKA)\n # Current slow state time\n slow_stage_time = time + slow.RKC[slow_s] * slow_dt\n\n # Fractional time for slow stage\n if slow_s == length(slow.RKA)\n fract_dt = (1 - slow.RKC[slow_s]) * slow_dt\n else\n fract_dt = (slow.RKC[slow_s + 1] - slow.RKC[slow_s]) * slow_dt\n end\n\n # Initialize fast model and set time-step and number of substeps we need\n # Note: to reproduce previous Fast output, 1) remove \"firstStage = ...\" line\n fast_steps = [0 0 0]\n FT = typeof(slow_dt)\n fast_time_rec = [fast_dt_in FT(0)]\n initialize_fast_state!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n fract_dt,\n fast_time_rec,\n fast_steps;\n firstStage = (slow_s == 1),\n )\n # Initialize tentency adjustment before evaluation of slow mode\n initialize_adjustment!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n )\n\n # Evaluate the slow mode\n # --> save tendency for the fast\n slow.rhs!(dQ2fast, Qslow, param, slow_stage_time, increment = false)\n\n # vertically integrate slow tendency to advance fast equation\n # and use vertical mean for slow model (negative source)\n # ---> work with dQ2fast as input\n tendency_from_slow_to_fast!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n dQ2fast,\n )\n\n # Compute (and RK update) slow tendency\n slow.rhs!(dQslow, Qslow, param, slow_stage_time, increment = true)\n\n # Update (RK-stage) slow state\n event = Event(array_device(Qslow))\n event = update!(array_device(Qslow), groupsize)(\n realview(dQslow),\n realview(Qslow),\n slow.RKA[slow_s % length(slow.RKA) + 1],\n slow.RKB[slow_s],\n slow_dt,\n nothing,\n nothing,\n nothing;\n ndrange = length(realview(Qslow)),\n dependencies = (event,),\n )\n wait(array_device(Qslow), event)\n\n # Determine number of substeps we need\n fast_dt = fast_time_rec[1]\n nsubsteps = fast_steps[3]\n\n updatedt!(fast, fast_dt)\n\n for substep in 1:nsubsteps\n fast_time = slow_stage_time + (substep - 1) * fast_dt\n dostep!(Qfast, fast, param, fast_time)\n\n # cumulate fast solution\n cummulate_fast_solution!(\n fast_bl,\n fast.rhs!,\n Qfast,\n fast_time,\n fast_dt,\n substep,\n fast_steps,\n fast_time_rec,\n )\n end\n\n # Reconcile slow equation using fast equation\n # Note: to reproduce previous Fast output, 2) remove \"lastStage = ...\" line\n reconcile_from_fast_to_slow!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n fast_time_rec;\n lastStage = (slow_s == length(slow.RKA)),\n )\n\n end\n\n # reset fast time-step to original value\n updatedt!(fast, fast_dt_in)\n\n # now do implicit mixing step\n nImplSteps = slow_bl.numImplSteps\n if nImplSteps > 0\n # 1. get implicit mising model, model state variable array and solver handles\n ivdc_dg = slow.rhs!.modeldata.ivdc_dg\n ivdc_bl = ivdc_dg.balance_law\n ivdc_Q = slow.rhs!.modeldata.ivdc_Q\n ivdc_solver = slow.rhs!.modeldata.ivdc_bgm_solver\n # ivdc_solver_dt = getdt(ivdc_solver) # would work if solver time-step was set\n # FT = typeof(slow_dt)\n # ivdc_solver_dt = slow_dt / FT(nImplSteps) # just recompute time-step\n ivdc_solver_dt = ivdc_bl.parent_om.ivdc_dt\n # println(\"ivdc_solver_dt = \",ivdc_solver_dt )\n # 2. setup start RHS, initial guess and values for computing mixing coeff\n ivdc_Q.θ .= Qslow.θ\n ivdc_RHS = slow.rhs!.modeldata.ivdc_RHS\n ivdc_RHS.θ .= Qslow.θ\n ivdc_RHS.θ .= ivdc_RHS.θ ./ ivdc_solver_dt\n ivdc_dg.state_auxiliary.θ_init .= ivdc_Q.θ\n # 3. Invoke iterative solver\n\n println(\"BEFORE maximum(ivdc_Q.θ[:]): \", maximum(ivdc_Q.realdata[:]))\n println(\"BEFORE minimum(ivdc_Q.θ[:]): \", minimum(ivdc_Q.realdata[:]))\n\n lm!(y, x) = ivdc_dg(y, x, nothing, 0; increment = false)\n solve_tot = 0\n iter_tot = 0\n for i in 1:nImplSteps\n solve_time = @elapsed iters =\n linearsolve!(lm!, nothing, ivdc_solver, ivdc_Q, ivdc_RHS)\n solve_tot = solve_tot + solve_time\n iter_tot = iter_tot + iters\n # Set new RHS and initial values\n ivdc_RHS.θ .= ivdc_Q.θ ./ ivdc_solver_dt\n ivdc_dg.state_auxiliary.θ_init .= ivdc_Q.θ\n end\n println(\"solver iters, time: \", iter_tot, \", \", solve_tot)\n\n println(\"AFTER maximum(ivdc_Q.θ[:]): \", maximum(ivdc_Q.realdata[:]))\n println(\"AFTER minimum(ivdc_Q.θ[:]): \", minimum(ivdc_Q.realdata[:]))\n\n # exit()\n # Now update\n Qslow.θ .= ivdc_Q.θ\n\n end\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit01/SplitExplicitLSRK3nMethod.jl", "content": "export SplitExplicitLSRK3nSolver\n\nusing KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\nusing StaticArrays\nusing ...SystemSolvers\nusing ...MPIStateArrays: array_device, realview\nusing ...GenericCallbacks\n#using Printf\n\nusing ...ODESolvers:\n AbstractODESolver, LowStorageRungeKutta3N, update!, updatedt!, getdt\nimport ...ODESolvers: dostep!\n\nusing ...BalanceLaws:\n tendency_from_slow_to_fast!,\n cummulate_fast_solution!,\n reconcile_from_fast_to_slow!\n\nLSRK3N = LowStorageRungeKutta3N\n\n@doc \"\"\"\n SplitExplicitLSRK3nSolver(slow_solver, fast_solver; dt, t0 = 0, coupled = true)\n\nThis is a time stepping object for explicitly time stepping the differential\nequation given by the right-hand-side function `f` with the state `Q`, i.e.,\n\n```math\n \\\\dot{Q_fast} = f_fast(Q_fast, Q_slow, t)\n \\\\dot{Q_slow} = f_slow(Q_slow, Q_fast, t)\n```\n\nwith the required time step size `dt` and optional initial time `t0`. This\ntime stepping object is intended to be passed to the `solve!` command.\n\nThis method performs an operator splitting to timestep the Sea-Surface elevation\nand vertically averaged horizontal velocity of the model at a faster rate than\nthe full model, using LowStorageRungeKutta3N time-stepping.\n\n\"\"\" SplitExplicitLSRK3nSolver\nmutable struct SplitExplicitLSRK3nSolver{SS, FS, RT, MSA, MSB} <:\n AbstractODESolver\n \"slow solver\"\n slow_solver::SS\n \"fast solver\"\n fast_solver::FS\n \"time step\"\n dt::RT\n \"time\"\n t::RT\n \"elapsed time steps\"\n steps::Int\n \"storage for transfer tendency\"\n dQ2fast::MSA\n \"saving original fast state\"\n S_fast::MSB\n\n function SplitExplicitLSRK3nSolver(\n slow_solver::LSRK3N,\n fast_solver,\n Q = nothing;\n dt = getdt(slow_solver),\n t0 = slow_solver.t,\n ) where {AT <: AbstractArray}\n SS = typeof(slow_solver)\n FS = typeof(fast_solver)\n RT = real(eltype(slow_solver.dQ))\n\n dQ2fast = similar(slow_solver.dQ)\n dQ2fast .= -0.0\n #- Warning: Number of fast-solution to save (here only 1, in S_fast) should be as\n # large as number of non zero Butcher Coeff (including Weight \"b\") below 2 diagonal,\n # i.e., with ns= Number of Stages and a(ns+1,:) = b(:), all non zero a(i,j)|_{i > j+1}.\n # Saving only 1 fast-solution workd for LS3NRK33Heuns.\n S_fast = similar(fast_solver.dQ)\n S_fast .= -0.0\n MSA = typeof(dQ2fast)\n MSB = typeof(S_fast)\n return new{SS, FS, RT, MSA, MSB}(\n slow_solver,\n fast_solver,\n RT(dt),\n RT(t0),\n 0,\n dQ2fast,\n S_fast,\n )\n end\nend\n\nfunction dostep!(\n Qvec,\n split::SplitExplicitLSRK3nSolver{SS},\n param,\n time::Real,\n) where {SS <: LSRK3N}\n slow = split.slow_solver\n fast = split.fast_solver\n\n Qslow = Qvec.slow\n Qfast = Qvec.fast\n\n dQslow = slow.dQ\n dRslow = slow.dR\n dQ2fast = split.dQ2fast\n S_fast = split.S_fast\n\n slow_bl = slow.rhs!.balance_law\n fast_bl = fast.rhs!.balance_law\n\n rv_Q = realview(Qslow)\n rv_dQ = realview(dQslow)\n rv_dR = realview(dRslow)\n groupsize = 256\n\n slow_dt = getdt(slow)\n fast_dt_in = getdt(fast)\n rkA = slow.RKA\n rkB = slow.RKB\n rkC = slow.RKC\n rkW = slow.RKW\n nStages = length(rkC)\n\n rv_dR .= -0\n for s in 1:nStages\n # Current slow state time\n slow_stage_time = time + rkC[s] * slow_dt\n # @printf(\"-- main dostep! stage s=%3i , t= %10.2f\\n\",s,slow_stage_time)\n\n # Initialize fast model and set time-step and number of substeps we need\n fast_steps = [0 0 0]\n FT = typeof(slow_dt)\n fast_time_rec = [fast_dt_in FT(0) FT(0)]\n set_fast_for_stepping!(\n slow_bl,\n fast_bl,\n fast.rhs!,\n Qfast,\n S_fast,\n slow_dt,\n rkC,\n rkW,\n s,\n nStages,\n fast_time_rec,\n fast_steps,\n )\n # Initialize tentency adjustment before evaluation of slow mode\n initialize_adjustment!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n )\n\n # Evaluate the slow mode\n # --> save tendency for the fast\n slow.rhs!(dQ2fast, Qslow, param, slow_stage_time, increment = false)\n\n # vertically integrate slow tendency to advance fast equation\n # and use vertical mean for slow model (negative source)\n # ---> work with dQ2fast as input\n tendency_from_slow_to_fast!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n dQ2fast,\n )\n\n # Compute (and RK update) slow tendency\n slow.rhs!(dQslow, Qslow, param, slow_stage_time, increment = true)\n\n # Update (RK-stage) slow state\n event = Event(array_device(Qslow))\n event = update!(array_device(Qslow), groupsize)(\n rv_dQ,\n rv_dR,\n rv_Q,\n rkA[s % nStages + 1, 1],\n rkA[s % nStages + 1, 2],\n rkB[s, 1],\n rkB[s, 2],\n slow_dt,\n nothing,\n nothing,\n nothing;\n ndrange = length(rv_Q),\n dependencies = (event,),\n )\n wait(array_device(Qslow), event)\n\n # Determine number of substeps we need\n fast_dt = fast_time_rec[1]\n nsubsteps = fast_steps[3]\n\n updatedt!(fast, fast_dt)\n\n for substep in 1:nsubsteps\n fast_time = time + fast_time_rec[3] + (substep - 1) * fast_dt\n # @printf(\"-- main dostep! substep=%3i , t= %10.2f\\n\",substep,fast_time)\n dostep!(Qfast, fast, param, fast_time)\n\n # cumulate fast solution\n cummulate_fast_solution!(\n fast_bl,\n fast.rhs!,\n Qfast,\n fast_time,\n fast_dt,\n substep,\n fast_steps,\n fast_time_rec,\n )\n end\n\n # reconcile slow equation using fast equation\n reconcile_from_fast_to_slow!(\n slow_bl,\n fast_bl,\n slow.rhs!,\n fast.rhs!,\n Qslow,\n Qfast,\n fast_time_rec;\n lastStage = (s == nStages),\n )\n\n end\n\n # reset fast time-step to original value\n updatedt!(fast, fast_dt_in)\n\n # now do implicit mixing step\n nImplSteps = slow_bl.numImplSteps\n if nImplSteps > 0\n # 1. get implicit mising model, model state variable array and solver handles\n ivdc_dg = slow.rhs!.modeldata.ivdc_dg\n ivdc_bl = ivdc_dg.balance_law\n ivdc_Q = slow.rhs!.modeldata.ivdc_Q\n ivdc_solver = slow.rhs!.modeldata.ivdc_bgm_solver\n # ivdc_solver_dt = getdt(ivdc_solver) # would work if solver time-step was set\n # FT = typeof(slow_dt)\n # ivdc_solver_dt = slow_dt / FT(nImplSteps) # just recompute time-step\n ivdc_solver_dt = ivdc_bl.parent_om.ivdc_dt\n # println(\"ivdc_solver_dt = \",ivdc_solver_dt )\n # 2. setup start RHS, initial guess and values for computing mixing coeff\n ivdc_Q.θ .= Qslow.θ\n ivdc_RHS = slow.rhs!.modeldata.ivdc_RHS\n ivdc_RHS.θ .= Qslow.θ\n ivdc_RHS.θ .= ivdc_RHS.θ ./ ivdc_solver_dt\n ivdc_dg.state_auxiliary.θ_init .= ivdc_Q.θ\n # 3. Invoke iterative solver\n\n println(\"BEFORE maximum(ivdc_Q.θ[:]): \", maximum(ivdc_Q.realdata[:]))\n println(\"BEFORE minimum(ivdc_Q.θ[:]): \", minimum(ivdc_Q.realdata[:]))\n\n lm!(y, x) = ivdc_dg(y, x, nothing, 0; increment = false)\n solve_tot = 0\n iter_tot = 0\n for i in 1:nImplSteps\n solve_time = @elapsed iters =\n linearsolve!(lm!, nothing, ivdc_solver, ivdc_Q, ivdc_RHS)\n solve_tot = solve_tot + solve_time\n iter_tot = iter_tot + iters\n # Set new RHS and initial values\n ivdc_RHS.θ .= ivdc_Q.θ ./ ivdc_solver_dt\n ivdc_dg.state_auxiliary.θ_init .= ivdc_Q.θ\n end\n println(\"solver iters, time: \", iter_tot, \", \", solve_tot)\n\n println(\"AFTER maximum(ivdc_Q.θ[:]): \", maximum(ivdc_Q.realdata[:]))\n println(\"AFTER minimum(ivdc_Q.θ[:]): \", minimum(ivdc_Q.realdata[:]))\n\n # exit()\n # Now update\n Qslow.θ .= ivdc_Q.θ\n\n end\n\n return nothing\nend\n" }, { "path": "src/Ocean/SplitExplicit01/SplitExplicitModel.jl", "content": "module SplitExplicit01\n\nexport OceanDGModel,\n OceanModel,\n Continuity3dModel,\n HorizontalModel,\n BarotropicModel,\n AbstractOceanProblem\n\n#using Printf\nusing StaticArrays\nusing LinearAlgebra: I, dot, Diagonal\n\nusing ...VariableTemplates\nusing ...MPIStateArrays\nusing ...DGMethods: init_ode_state, basic_grid_info\nusing ...Mesh.Filters: CutoffFilter, apply!, ExponentialFilter\nusing ...Mesh.Grids:\n polynomialorders,\n dimensionality,\n dofs_per_element,\n VerticalDirection,\n HorizontalDirection,\n min_node_distance\n\nusing ...BalanceLaws\n#import ...BalanceLaws: nodal_update_auxiliary_state!\n\nusing ...DGMethods.NumericalFluxes:\n NumericalFluxFirstOrder,\n NumericalFluxGradient,\n NumericalFluxSecondOrder,\n RusanovNumericalFlux,\n CentralNumericalFluxFirstOrder,\n CentralNumericalFluxGradient,\n CentralNumericalFluxSecondOrder\n\nusing ..Ocean: AbstractOceanProblem\n\nimport ...DGMethods.NumericalFluxes:\n update_penalty!, numerical_flux_second_order!, NumericalFluxFirstOrder\n\nimport ...DGMethods: LocalGeometry, DGModel, calculate_dt\n\nimport ...BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n parameter_set,\n boundary_conditions,\n boundary_state!,\n update_auxiliary_state!,\n update_auxiliary_state_gradient!,\n compute_gradient_argument!,\n init_state_auxiliary!,\n init_state_prognostic!,\n compute_gradient_flux!,\n indefinite_stack_integral!,\n reverse_indefinite_stack_integral!,\n integral_load_auxiliary_state!,\n integral_set_auxiliary_state!,\n reverse_integral_load_auxiliary_state!,\n reverse_integral_set_auxiliary_state!\n\nimport ...SystemSolvers: BatchedGeneralizedMinimalResidual, linearsolve!\n\n×(a::SVector, b::SVector) = StaticArrays.cross(a, b)\n∘(a::SVector, b::SVector) = StaticArrays.dot(a, b)\n\nabstract type AbstractOceanModel <: BalanceLaw end\n\nfunction ocean_init_aux! end\nfunction ocean_init_state! end\nfunction ocean_model_boundary! end\nfunction set_fast_for_stepping! end\nfunction initialize_fast_state! end\nfunction initialize_adjustment! end\n\ninclude(\"SplitExplicitLSRK2nMethod.jl\")\ninclude(\"SplitExplicitLSRK3nMethod.jl\")\ninclude(\"OceanModel.jl\")\ninclude(\"Continuity3dModel.jl\")\ninclude(\"VerticalIntegralModel.jl\")\ninclude(\"BarotropicModel.jl\")\ninclude(\"IVDCModel.jl\")\ninclude(\"Communication.jl\")\ninclude(\"OceanBoundaryConditions.jl\")\n\nend\n" }, { "path": "src/Ocean/SplitExplicit01/VerticalIntegralModel.jl", "content": "struct TendencyIntegralModel{M} <: AbstractOceanModel\n ocean::M\n function TendencyIntegralModel(ocean::M) where {M}\n return new{M}(ocean)\n end\nend\nvars_state(tm::TendencyIntegralModel, ::Prognostic, FT) =\n vars_state(tm.ocean, Prognostic(), FT)\nvars_state(tm::TendencyIntegralModel, ::GradientFlux, FT) = @vars()\n\nfunction vars_state(m::TendencyIntegralModel, ::Auxiliary, T)\n @vars begin\n ∫du::SVector{2, T}\n end\nend\n\nfunction vars_state(m::TendencyIntegralModel, ::UpwardIntegrals, T)\n @vars begin\n ∫du::SVector{2, T}\n end\nend\n\n@inline function integral_load_auxiliary_state!(\n m::TendencyIntegralModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.∫du = A.∫du\n\n return nothing\nend\n\n@inline function integral_set_auxiliary_state!(\n m::TendencyIntegralModel,\n A::Vars,\n I::Vars,\n)\n A.∫du = I.∫du\n\n return nothing\nend\n\ninit_state_auxiliary!(tm::TendencyIntegralModel, A::Vars, _...) = nothing\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n tm::TendencyIntegralModel,\n dQ::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n A = dg.state_auxiliary\n\n # copy tendency vector to aux state for integration\n function f!(::TendencyIntegralModel, dQ, A, t)\n @inbounds begin\n A.∫du = @SVector [dQ.u[1], dQ.u[2]]\n end\n\n return nothing\n end\n update_auxiliary_state!(f!, dg, tm, dQ, t)\n\n # compute integral for Gᵁ\n indefinite_stack_integral!(dg, tm, dQ, A, t, elems) # bottom -> top\n\n return true\nend\n\n#-------------------------------------------------------------------------------\nstruct FlowIntegralModel{M} <: AbstractOceanModel\n ocean::M\n function FlowIntegralModel(ocean::M) where {M}\n return new{M}(ocean)\n end\nend\nvars_state(fm::FlowIntegralModel, ::Prognostic, FT) =\n vars_state(fm.ocean, Prognostic(), FT)\nvars_state(fm::FlowIntegralModel, ::GradientFlux, FT) = @vars()\n\nfunction vars_state(m::FlowIntegralModel, ::Auxiliary, T)\n @vars begin\n ∫u::SVector{2, T}\n end\nend\n\nfunction vars_state(m::FlowIntegralModel, ::UpwardIntegrals, T)\n @vars begin\n ∫u::SVector{2, T}\n end\nend\n\n@inline function integral_load_auxiliary_state!(\n m::FlowIntegralModel,\n I::Vars,\n Q::Vars,\n A::Vars,\n)\n I.∫u = Q.u\n\n return nothing\nend\n\n@inline function integral_set_auxiliary_state!(\n m::FlowIntegralModel,\n A::Vars,\n I::Vars,\n)\n A.∫u = I.∫u\n\n return nothing\nend\n\ninit_state_auxiliary!(fm::FlowIntegralModel, A::Vars, _...) = nothing\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n fm::FlowIntegralModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n A = dg.state_auxiliary\n\n # compute vertical integral of u\n indefinite_stack_integral!(dg, fm, Q, A, t, elems) # bottom -> top\n\n return true\nend\n" }, { "path": "src/Ocean/SuperModels.jl", "content": "module SuperModels\n\nusing MPI\n\nusing ClimateMachine\n\nusing ClimateMachine: Settings\nusing ...BalanceLaws: parameter_set\n\nusing ...DGMethods.NumericalFluxes\n\nusing ..HydrostaticBoussinesq:\n HydrostaticBoussinesqModel, NonLinearAdvectionTerm, Forcing\n\nusing ..OceanProblems: InitialValueProblem, InitialConditions\nusing ..Ocean: FreeSlip, Impenetrable, Insulating, OceanBC, Penetrable\nusing ..CartesianFields: SpectralElementField\n\nusing ...Mesh.Filters: CutoffFilter, ExponentialFilter\nusing ...Mesh.Grids: polynomialorders, DiscontinuousSpectralElementGrid\n\nusing ClimateMachine:\n LS3NRK33Heuns,\n OceanBoxGCMConfigType,\n OceanBoxGCMSpecificInfo,\n DriverConfiguration\n\nimport ClimateMachine: SolverConfiguration\n\n#####\n##### It's super good\n#####\n\nstruct HydrostaticBoussinesqSuperModel{D, G, E, S, F, N, T, C}\n domain::D\n grid::G\n equations::E\n state::S\n fields::F\n numerical_fluxes::N\n timestepper::T\n solver_configuration::C\nend\n\n\"\"\"\n HydrostaticBoussinesqSuperModel(; domain, time_step, parameters,\n initial_conditions = InitialConditions(),\n advection = (momentum = NonLinearAdvectionTerm(), tracers = NonLinearAdvectionTerm()),\n turbulence_closure = (νʰ=0, νᶻ=0, κʰ=0, κᶻ=0),\n coriolis = (f₀=0, β=0),\n rusanov_wave_speeds = (cʰ=0, cᶻ=0),\n buoyancy = (αᵀ=0,)\n numerical_fluxes = ( first_order = RusanovNumericalFlux(),\n second_order = CentralNumericalFluxSecondOrder(),\n gradient = CentralNumericalFluxGradient())\n timestepper = ClimateMachine.ExplicitSolverType(solver_method=LS3NRK33Heuns),\n )\n\nBuilds a `SuperModel` that solves the Hydrostatic Boussinesq equations.\n\"\"\"\nfunction HydrostaticBoussinesqSuperModel(;\n domain,\n parameters,\n time_step, # We don't want to have to provide this here, but neverthless it's required.\n initial_conditions = InitialConditions(),\n advection = (\n momentum = NonLinearAdvectionTerm(),\n tracers = NonLinearAdvectionTerm(),\n ),\n turbulence_closure = (νʰ = 0, νᶻ = 0, κʰ = 0, κᶻ = 0),\n coriolis = (f₀ = 0, β = 0),\n rusanov_wave_speeds = (cʰ = 0, cᶻ = 0),\n buoyancy = (αᵀ = 0,),\n forcing = Forcing(),\n numerical_fluxes = (\n first_order = RusanovNumericalFlux(),\n second_order = CentralNumericalFluxSecondOrder(),\n gradient = CentralNumericalFluxGradient(),\n ),\n timestepper = ClimateMachine.ExplicitSolverType(\n solver_method = LS3NRK33Heuns,\n ),\n filters = nothing,\n modeldata = NamedTuple(),\n array_type = Settings.array_type,\n mpicomm = MPI.COMM_WORLD,\n init_on_cpu = true,\n boundary_tags = ((0, 0), (0, 0), (1, 2)),\n boundary_conditions = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n ),\n)\n\n #####\n ##### Build the grid\n #####\n\n # Change global setting if its set here\n Settings.array_type = array_type\n\n grid = DiscontinuousSpectralElementGrid(\n domain;\n boundary_tags = boundary_tags,\n mpicomm = mpicomm,\n array_type = array_type,\n )\n\n FT = eltype(domain)\n\n #####\n ##### Construct generic problem type InitialValueProblem\n #####\n\n problem = InitialValueProblem{FT}(\n dimensions = (domain.L.x, domain.L.y, domain.L.z),\n initial_conditions = initial_conditions,\n boundary_conditions = boundary_conditions,\n )\n\n #####\n ##### Build HydrostaticBoussinesqModel/Equations\n #####\n\n equations = HydrostaticBoussinesqModel{eltype(domain)}(\n parameters,\n problem,\n momentum_advection = advection.momentum,\n tracer_advection = advection.tracers,\n forcing = forcing,\n cʰ = convert(FT, rusanov_wave_speeds.cʰ),\n cᶻ = convert(FT, rusanov_wave_speeds.cᶻ),\n αᵀ = convert(FT, buoyancy.αᵀ),\n νʰ = convert(FT, turbulence_closure.νʰ),\n νᶻ = convert(FT, turbulence_closure.νᶻ),\n κʰ = convert(FT, turbulence_closure.κʰ),\n κᶻ = convert(FT, turbulence_closure.κᶻ),\n fₒ = convert(FT, coriolis.f₀),\n β = FT(coriolis.β),\n )\n\n ####\n #### \"modeldata\"\n ####\n #### OceanModels require filters (?). If one was not provided, we build a default.\n ####\n\n # Default vertical filter and horizontal exponential filter:\n if isnothing(filters)\n filters = (\n vert_filter = CutoffFilter(grid, polynomialorders(grid)),\n exp_filter = ExponentialFilter(grid, 1, 8),\n )\n end\n\n modeldata = merge(modeldata, filters)\n\n ####\n #### We build a DriverConfiguration here for the purposes of building\n #### a SolverConfiguration. Then we throw it away.\n ####\n\n driver_configuration = DriverConfiguration(\n OceanBoxGCMConfigType(),\n \"\",\n (domain.Np, domain.Np),\n eltype(domain),\n array_type,\n parameter_set(equations),\n equations,\n MPI.COMM_WORLD,\n grid,\n numerical_fluxes.first_order,\n numerical_fluxes.second_order,\n numerical_fluxes.gradient,\n nothing,\n nothing, # filter\n OceanBoxGCMSpecificInfo(),\n )\n\n ####\n #### Pass through the SolverConfiguration interface so that we use\n #### the checkpointing infrastructure\n ####\n\n solver_configuration = ClimateMachine.SolverConfiguration(\n zero(FT),\n convert(FT, time_step),\n driver_configuration,\n init_on_cpu = init_on_cpu,\n ode_dt = convert(FT, time_step),\n ode_solver_type = timestepper,\n Courant_number = 0.4,\n modeldata = modeldata,\n )\n\n state = solver_configuration.Q\n\n u = SpectralElementField(domain, grid, state, 1)\n v = SpectralElementField(domain, grid, state, 2)\n η = SpectralElementField(domain, grid, state, 3)\n θ = SpectralElementField(domain, grid, state, 4)\n\n fields = (u = u, v = v, η = η, θ = θ)\n\n return HydrostaticBoussinesqSuperModel(\n domain,\n grid,\n equations,\n state,\n fields,\n numerical_fluxes,\n timestepper,\n solver_configuration,\n )\nend\n\ncurrent_time(model::HydrostaticBoussinesqSuperModel) =\n model.solver_configuration.solver.t\nΔt(model::HydrostaticBoussinesqSuperModel) =\n model.solver_configuration.solver.dt\ncurrent_step(model::HydrostaticBoussinesqSuperModel) =\n model.solver_configuration.solver.steps\n\nend # module\n" }, { "path": "src/Utilities/SingleStackUtils/SingleStackUtils.jl", "content": "module SingleStackUtils\n\nexport get_vars_from_nodal_stack,\n get_vars_from_element_stack,\n get_horizontal_variance,\n get_horizontal_mean,\n reduce_nodal_stack,\n reduce_element_stack,\n horizontally_average!,\n dict_of_nodal_states,\n NodalStack,\n single_stack_diagnostics\n\nusing OrderedCollections\nusing UnPack\nusing StaticArrays\nimport KernelAbstractions: CPU\n\nusing ..BalanceLaws\nusing ..DGMethods\nusing ..DGMethods.Grids\nusing ..MPIStateArrays\nusing ..VariableTemplates\n\n\"\"\"\n get_vars_from_nodal_stack(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n i::Int = 1,\n j::Int = 1,\n exclude::Vector{String} = String[],\n interp = false,\n ) where {T, dim, N}\n\nReturn a dictionary whose keys are the `flattenednames()` of the variables\nspecified in `vars` (as returned by e.g. `vars_state`), and\nwhose values are arrays of the values for that variable along the vertical\ndimension in `Q`. Only a single element is expected in the horizontal as\nthis is intended for the single stack configuration and `i` and `j` identify\nthe horizontal nodal coordinates.\n\nVariables listed in `exclude` are skipped.\n\"\"\"\nfunction get_vars_from_nodal_stack(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n i::Int = 1,\n j::Int = 1,\n exclude::Vector{String} = String[],\n interp = false,\n) where {T, dim, N}\n\n # extract grid information and bring `Q` to the host if needed\n FT = eltype(Q)\n Nq = N .+ 1\n # Code assumes the same polynomial order in all horizontal directions\n @inbounds begin\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nqk = dim == 2 ? 1 : Nq[dim]\n end\n Np = dofs_per_element(grid)\n state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)\n\n # set up the dictionary to be returned\n var_names = flattenednames(vars)\n stack_vals = OrderedDict()\n num_vars = varsize(vars)\n vars_wanted = Int[]\n @inbounds for vi in 1:num_vars\n if !(var_names[vi] in exclude)\n stack_vals[var_names[vi]] = FT[]\n push!(vars_wanted, vi)\n end\n end\n elemtobndy = convert(Array, grid.elemtobndy)\n vmap⁻ = convert(Array, grid.vmap⁻)\n vmap⁺ = convert(Array, grid.vmap⁺)\n vgeo = convert(Array, grid.vgeo)\n # extract values from `state_data`\n @inbounds for ev in vrange, k in 1:Nqk, v in vars_wanted\n if interp && k == 1 && elemtobndy[5, ev] == 0\n # Get face degree of freedom number\n n = i + Nq1 * ((j - 1))\n # get the element numbers\n ev⁻ = ev\n # Get neighboring id data\n id⁻, id⁺ = vmap⁻[n, 5, ev⁻], vmap⁺[n, 5, ev⁻]\n ev⁺ = ((id⁺ - 1) ÷ Np) + 1\n # get the volume degree of freedom numbers\n vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1\n\n J⁻, J⁺ = vgeo[vid⁻, Grids._M, ev⁻], vgeo[vid⁺, Grids._M, ev⁺]\n state_local = J⁻ * state_data[vid⁻, v, ev⁻]\n state_local += J⁺ * state_data[vid⁺, v, ev⁺]\n state_local /= (J⁻ + J⁺)\n push!(stack_vals[var_names[v]], state_local)\n elseif interp && k == Nqk && elemtobndy[6, ev] == 0\n # Get face degree of freedom number\n n = i + Nq1 * ((j - 1))\n # get the element numbers\n ev⁻ = ev\n # Get neighboring id data\n id⁻, id⁺ = vmap⁻[n, 6, ev⁻], vmap⁺[n, 6, ev⁻]\n # periodic and need to handle this point (otherwise handled above)\n if id⁺ == id⁻\n vid⁻ = ((id⁻ - 1) % Np) + 1\n\n state_local = state_data[vid⁻, v, ev⁻]\n push!(stack_vals[var_names[v]], state_local)\n end\n else\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n state_local = state_data[ijk, v, ev]\n push!(stack_vals[var_names[v]], state_local)\n end\n end\n\n return stack_vals\nend\n\n\"\"\"\n get_vars_from_element_stack(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n ) where {T, dim, N}\n\nReturn an array of [`get_vars_from_nodal_stack()`](@ref)s whose dimensions\nare the number of nodal points per element in the horizontal plane.\n\nVariables listed in `exclude` are skipped.\n\"\"\"\nfunction get_vars_from_element_stack(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds Nq1 = Nq[1]\n @inbounds Nq2 = Nq[2]\n\n return [\n get_vars_from_nodal_stack(\n grid,\n Q,\n vars,\n vrange = vrange,\n i = i,\n j = j,\n exclude = exclude,\n interp = interp,\n ) for i in 1:Nq1, j in 1:Nq2\n ]\nend\n\n\"\"\"\n get_horizontal_mean(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n ) where {T, dim, N}\n\nReturn a dictionary whose keys are the `flattenednames()` of the variables\nspecified in `vars` (as returned by e.g. `vars_state`), and\nwhose values are arrays of the horizontal averages for that variable along\nthe vertical dimension in `Q`. Only a single element is expected in the\nhorizontal as this is intended for the single stack configuration.\n\nVariables listed in `exclude` are skipped.\n\"\"\"\nfunction get_horizontal_mean(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds Nq1 = Nq[1]\n @inbounds Nq2 = Nq[2]\n\n vars_avg = OrderedDict()\n vars_sq = OrderedDict()\n for i in 1:Nq1\n for j in 1:Nq2\n vars_nodal = get_vars_from_nodal_stack(\n grid,\n Q,\n vars,\n vrange = vrange,\n i = i,\n j = j,\n exclude = exclude,\n interp = interp,\n )\n vars_avg = merge(+, vars_avg, vars_nodal)\n end\n end\n map!(x -> x ./ Nq1 / Nq1, values(vars_avg))\n return vars_avg\nend\n\n\"\"\"\n get_horizontal_variance(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n ) where {T, dim, N}\n\nReturn a dictionary whose keys are the `flattenednames()` of the variables\nspecified in `vars` (as returned by e.g. `vars_state`), and\nwhose values are arrays of the horizontal variance for that variable along\nthe vertical dimension in `Q`. Only a single element is expected in the\nhorizontal as this is intended for the single stack configuration.\n\nVariables listed in `exclude` are skipped.\n\"\"\"\nfunction get_horizontal_variance(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars;\n vrange::UnitRange = 1:size(Q, 3),\n exclude::Vector{String} = String[],\n interp = false,\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds Nq1 = Nq[1]\n @inbounds Nq2 = Nq[2]\n\n vars_avg = OrderedDict()\n vars_sq = OrderedDict()\n for i in 1:Nq1\n for j in 1:Nq2\n vars_nodal = get_vars_from_nodal_stack(\n grid,\n Q,\n vars,\n vrange = vrange,\n i = i,\n j = j,\n exclude = exclude,\n interp = interp,\n )\n vars_nodal_sq = OrderedDict(vars_nodal)\n map!(x -> x .^ 2, values(vars_nodal_sq))\n vars_avg = merge(+, vars_avg, vars_nodal)\n vars_sq = merge(+, vars_sq, vars_nodal_sq)\n end\n end\n map!(x -> (x ./ Nq1 / Nq1) .^ 2, values(vars_avg))\n map!(x -> x ./ Nq1 / Nq1, values(vars_sq))\n vars_var = merge(-, vars_sq, vars_avg)\n return vars_var\nend\n\n\"\"\"\n reduce_nodal_stack(\n op::Function,\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars::NamedTuple,\n var::String;\n vrange::UnitRange = 1:size(Q, 3),\n ) where {T, dim, N}\n\nReduce `var` from `vars` within `Q` over all nodal points in the specified\n`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is\nthe final value returned by `op` and `z` is the index within `vrange` where the\n`result` was determined.\n\"\"\"\nfunction reduce_nodal_stack(\n op::Function,\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars::Type,\n var::String;\n vrange::UnitRange = 1:size(Q, 3),\n i::Int = 1,\n j::Int = 1,\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds begin\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nqk = dim == 2 ? 1 : Nq[dim]\n end\n\n var_names = flattenednames(vars)\n var_ind = findfirst(s -> s == var, var_names)\n if var_ind === nothing\n return\n end\n\n state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)\n z = vrange.start\n FT = eltype(state_data)\n\n # Initialize result with identity operation for operator\n result = if op isa typeof(+)\n zero(FT)\n elseif op isa typeof(*)\n one(FT)\n elseif op isa typeof(min)\n floatmax(FT)\n elseif op isa typeof(max)\n -floatmax(FT)\n else\n error(\"unknown operator: $op\")\n end\n\n for ev in vrange\n for k in 1:Nqk\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n new_result = op(result, state_data[ijk, var_ind, ev])\n if !isequal(new_result, result)\n result = new_result\n z = ev\n end\n end\n end\n\n return (result, z)\nend\n\n\"\"\"\n reduce_element_stack(\n op::Function,\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars::NamedTuple,\n var::String;\n vrange::UnitRange = 1:size(Q, 3),\n ) where {T, dim, N}\n\nReduce `var` from `vars` within `Q` over all nodal points in the specified\n`vrange` of elements with `op`. Return a tuple `(result, z)` where `result` is\nthe final value returned by `op` and `z` is the index within `vrange` where the\n`result` was determined.\n\"\"\"\nfunction reduce_element_stack(\n op::Function,\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n vars::Type,\n var::String;\n vrange::UnitRange = 1:size(Q, 3),\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds Nq1 = Nq[1]\n @inbounds Nq2 = Nq[2]\n\n return [\n reduce_nodal_stack(\n op,\n grid,\n Q,\n vars,\n var,\n vrange = vrange,\n i = i,\n j = j,\n ) for i in 1:Nq1, j in 1:Nq2\n ]\nend\n\n\"\"\"\n horizontally_average!(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n i_vars,\n ) where {T, dim, N}\n\nHorizontally average variables, from variable\nindexes `i_vars`, in `MPIStateArray` `Q`.\n\n!!! note\n These are not proper horizontal averages-- the main\n purpose of this method is to ensure that there are\n no horizontal fluxes for a single stack configuration.\n\"\"\"\nfunction horizontally_average!(\n grid::DiscontinuousSpectralElementGrid{T, dim, N},\n Q::MPIStateArray,\n i_vars,\n) where {T, dim, N}\n\n Nq = N .+ 1\n @inbounds begin\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nqk = dim == 2 ? 1 : Nq[dim]\n end\n\n ArrType = typeof(Q.data)\n state_data = array_device(Q) isa CPU ? Q.realdata : Array(Q.realdata)\n\n for ev in 1:size(state_data, 3), k in 1:Nqk, i_v in i_vars\n Q_sum = 0\n for i in 1:Nq1, j in 1:Nq2\n Q_sum += state_data[i + Nq1 * ((j - 1) + Nq2 * (k - 1)), i_v, ev]\n end\n Q_ave = Q_sum / (Nq1 * Nq2)\n for i in 1:Nq1, j in 1:Nq2\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n state_data[ijk, i_v, ev] = Q_ave\n end\n end\n Q.realdata .= ArrType(state_data)\nend\n\nget_data(solver_config, ::Prognostic) = solver_config.Q\nget_data(solver_config, ::Auxiliary) = solver_config.dg.state_auxiliary\nget_data(solver_config, ::GradientFlux) = solver_config.dg.state_gradient_flux\n\n\"\"\"\n dict_of_nodal_states(\n solver_config,\n state_types = (Prognostic(), Auxiliary());\n aux_excludes = [],\n interp = false,\n )\n\nA dictionary of single stack prognostic and auxiliary\nvariables at the `i=1`,`j=1` node given\n - `solver_config` a `SolverConfiguration`\n - `aux_excludes` a vector of strings containing the\n variables to exclude from the auxiliary state.\n\"\"\"\nfunction dict_of_nodal_states(\n solver_config,\n state_types = (Prognostic(), Auxiliary());\n aux_excludes = String[],\n interp = false,\n)\n FT = eltype(solver_config.Q)\n all_state_vars = []\n for st in state_types\n state_vars = get_vars_from_nodal_stack(\n solver_config.dg.grid,\n get_data(solver_config, st),\n vars_state(solver_config.dg.balance_law, st, FT),\n exclude = st isa Auxiliary ? aux_excludes : String[],\n interp = interp,\n )\n push!(all_state_vars, state_vars...)\n end\n return OrderedDict(all_state_vars...)\nend\n\n# A container for holding various\n# global or point-wise states:\nstruct States{P, A, D, HD}\n prog::P\n aux::A\n diffusive::D\n hyperdiffusive::HD\nend\n\n\"\"\"\n NodalStack(\n bl::BalanceLaw,\n grid::DiscontinuousSpectralElementGrid,\n prognostic,\n auxiliary,\n diffusive,\n hyperdiffusive;\n i = 1,\n j = 1,\n interp = true,\n )\n\nA struct whose `iterate(::NodalStack)` traverses\nthe nodal stack and returns a NamedTuple of\npoint-wise fields (`Vars`).\n\n# Example\n```julia\nfor state_local in NodalStack(\n bl,\n grid,\n prognostic, # global field along nodal stack\n auxiliary,\n diffusive,\n hyperdiffusive\n )\nprog = state_local.prog # point-wise field along nodal stack\nend\n```\n\n## TODO: Make `prognostic`, `auxiliary`, `diffusive`, `hyperdiffusive` optional\n\n# Arguments\n - `bl` the balance law\n - `grid` the discontinuous spectral element grid\n - `prognostic` the global prognostic state\n - `auxiliary` the global auxiliary state\n - `diffusive` the global diffusive state (gradient-flux)\n - `hyperdiffusive` the global hyperdiffusive state\n - `i,j` the `i,j`'th nodal stack (in the horizontal directions)\n - `interp` a bool indicating whether to\n interpolate the duplicate Gauss-Lebotto\n points at the element faces.\n\n!!! warn\n Before iterating, the data is transferred from the\n device (GPU) to the host (CPU), as this is intended\n for debugging / diagnostics usage.\n\"\"\"\nstruct NodalStack{N, BL, G, S, VR, TI, TJ, CI, IN}\n bl::BL\n grid::G\n states::S\n vrange::VR\n i::TI\n j::TJ\n cart_ind::CI\n interp::IN\n function NodalStack(\n bl::BalanceLaw,\n grid::DiscontinuousSpectralElementGrid;\n prognostic,\n auxiliary,\n diffusive,\n hyperdiffusive,\n i = 1,\n j = 1,\n interp = true,\n )\n states = States(prognostic, auxiliary, diffusive, hyperdiffusive)\n vrange = 1:size(prognostic, 3)\n grid_info = basic_grid_info(grid)\n @unpack Nqk = grid_info\n if last(polynomialorders(grid)) == 0\n interp = false\n end\n # Store cartesian indices, so we can map the iter_state\n # to the cartesian space `Q[i, var, j]`\n if interp\n cart_ind = CartesianIndices(((Nqk - 1), size(prognostic, 3)))\n else\n cart_ind = CartesianIndices((Nqk, size(prognostic, 3)))\n end\n args = (bl, grid, states, vrange, i, j, cart_ind, interp)\n BL, G, S, VR, TI, TJ, CI, IN = typeof.(args)\n if interp\n len = size(prognostic, 3) * (Nqk - 1) + 1\n else\n len = size(prognostic, 3) * Nqk\n end\n new{len, BL, G, S, VR, TI, TJ, CI, IN}(args...)\n end\nend\n\nBase.length(gs::NodalStack{N}) where {N} = N\n\n# Helper function\nget_state(v, state, vid⁻, vid⁺, ev⁻, ev⁺, J⁻, J⁺) =\n (J⁻ * state[vid⁻, v, ev⁻] + J⁺ * state[vid⁺, v, ev⁺]) / (J⁻ + J⁺)\n\nto_cpu(state) =\n array_device(state) isa CPU ? state.realdata : Array(state.realdata)\n\nfunction interp_top(state, args, n_vars, bl, st)\n vs = Vars{vars_state(bl, st, eltype(state))}\n if n_vars ≠ 0\n return vs(map(v -> get_state(v, state, args...), 1:n_vars))\n else\n return nothing\n end\nend\n\nfunction interp_bot(state, vid⁻, ev⁻, n_vars, bl, st)\n vs = Vars{vars_state(bl, st, eltype(state))}\n if n_vars ≠ 0\n return vs(map(v -> state[vid⁻, v, ev⁻], 1:n_vars))\n else\n return nothing\n end\nend\n\nfunction no_interp(state, ijk, ev, n_vars, bl, st)\n vs = Vars{vars_state(bl, st, eltype(state))}\n if n_vars ≠ 0\n return vs(map(v -> state[ijk, v, ev], 1:n_vars))\n else\n return nothing\n end\nend\n\nfunction Base.iterate(gs::NodalStack, iter_state = 1)\n iter_state > length(gs) && return nothing\n\n # extract grid information\n grid = gs.grid\n FT = eltype(grid)\n grid_info = basic_grid_info(grid)\n @unpack N, Nq, Np, Nqk = grid_info\n @inbounds Nq1, Nq2 = Nq[1], Nq[2]\n states = gs.states\n bl = gs.bl\n interp = gs.interp\n\n # bring `Q` to the host if needed\n\n prognostic = to_cpu(states.prog)\n auxiliary = to_cpu(states.aux)\n diffusive = to_cpu(states.diffusive)\n hyperdiffusive = to_cpu(states.hyperdiffusive)\n\n n_vars_prog = size(prognostic, 2)\n n_vars_aux = size(auxiliary, 2)\n n_vars_diff = size(diffusive, 2)\n n_vars_hd = size(hyperdiffusive, 2)\n\n elemtobndy = convert(Array, grid.elemtobndy)\n vmap⁻ = convert(Array, grid.vmap⁻)\n vmap⁺ = convert(Array, grid.vmap⁺)\n vgeo = convert(Array, grid.vgeo)\n i, j = gs.i, gs.j\n if iter_state == length(gs)\n ijk_cart = (Nqk, size(states.prog, 3))\n else\n ijk_cart = Tuple(gs.cart_ind[iter_state])\n end\n ev = ijk_cart[2]\n k = ijk_cart[1]\n iter_state⁺ = iter_state + 1\n if interp && k == 1 && elemtobndy[5, ev] == 0\n # Get face degree of freedom number\n n = i + Nq1 * ((j - 1))\n # get the element numbers\n ev⁻ = ev\n # Get neighboring id data\n id⁻, id⁺ = vmap⁻[n, 5, ev⁻], vmap⁺[n, 5, ev⁻]\n ev⁺ = ((id⁺ - 1) ÷ Np) + 1\n # get the volume degree of freedom numbers\n vid⁻, vid⁺ = ((id⁻ - 1) % Np) + 1, ((id⁺ - 1) % Np) + 1\n J⁻, J⁺ = vgeo[vid⁻, Grids._M, ev⁻], vgeo[vid⁺, Grids._M, ev⁺]\n\n args = (vid⁻, vid⁺, ev⁻, ev⁺, J⁻, J⁺)\n#! format: off\n prog = interp_top(prognostic, args, n_vars_prog, bl, Prognostic())\n aux = interp_top(auxiliary, args, n_vars_aux, bl, Auxiliary())\n ∇flux = interp_top(diffusive, args, n_vars_diff, bl, GradientFlux())\n hyperdiff = interp_top(hyperdiffusive, args, n_vars_hd, bl, Hyperdiffusive())\n#! format: on\n\n return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)\n\n elseif interp && k == Nqk && elemtobndy[6, ev] == 0\n # Get face degree of freedom number\n n = i + Nq1 * ((j - 1))\n # get the element numbers\n ev⁻ = ev\n # Get neighboring id data\n id⁻, id⁺ = vmap⁻[n, 6, ev⁻], vmap⁺[n, 6, ev⁻]\n # periodic and need to handle this point (otherwise handled above)\n if id⁺ == id⁻\n vid⁻ = ((id⁻ - 1) % Np) + 1\n\n#! format: off\n prog = interp_bot(prognostic, vid⁻, ev⁻, n_vars_prog, bl, Prognostic())\n aux = interp_bot(auxiliary, vid⁻, ev⁻, n_vars_aux, bl, Auxiliary())\n ∇flux = interp_bot(diffusive, vid⁻, ev⁻, n_vars_diff, bl, GradientFlux())\n hyperdiff = interp_bot( hyperdiffusive, vid⁻, ev⁻, n_vars_hd, bl, Hyperdiffusive())\n#! format: on\n\n return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)\n else\n error(\"uncaught case in iterate(::NodalStack)\")\n end\n else\n ijk = i + Nq1 * ((j - 1) + Nq2 * (k - 1))\n\n#! format: off\n prog = no_interp(prognostic, ijk, ev, n_vars_prog, bl, Prognostic())\n aux = no_interp(auxiliary, ijk, ev, n_vars_aux, bl, Auxiliary())\n ∇flux = no_interp(diffusive, ijk, ev, n_vars_diff, bl, GradientFlux())\n hyperdiff = no_interp(hyperdiffusive, ijk, ev, n_vars_hd, bl, Hyperdiffusive())\n#! format: on\n\n return ((; prog, aux, ∇flux, hyperdiff), iter_state⁺)\n end\nend\n\ninclude(\"single_stack_diagnostics.jl\")\n\nend # module\n" }, { "path": "src/Utilities/SingleStackUtils/single_stack_diagnostics.jl", "content": "using ..Orientations\nimport ..VariableTemplates: flattened_named_tuple\nusing ..VariableTemplates\n\n# Sometimes `NodalStack` returns local states\n# that is `nothing`. Here, we return `nothing`\n# to preserve the keys (e.g., `hyperdiff`)\n# when misssing.\nflattened_named_tuple(v::Nothing, ft::FlattenType = FlattenArr()) = nothing\n\n\"\"\"\n single_stack_diagnostics(\n grid::DiscontinuousSpectralElementGrid,\n bl::BalanceLaw,\n t::Real,\n direction;\n kwargs...,\n )\n\n# Arguments\n - `grid` the grid\n - `bl` the balance law\n - `t` time\n - `direction` direction\n - `kwargs` keyword arguments, passed to [`NodalStack`](@ref).\n\nAn array of nested NamedTuples, containing results of\n - `z` - altitude\n - `prog` - the prognostic state\n - `aux` - the auxiliary state\n - `∇flux` - the gradient-flux (diffusive) state\n - `hyperdiff` - the hyperdiffusive state\n\nand all the nested NamedTuples, merged together,\nfrom the `precompute` methods.\n\"\"\"\nfunction single_stack_diagnostics(\n grid::DiscontinuousSpectralElementGrid,\n bl::BalanceLaw,\n t::Real,\n direction;\n kwargs...,\n)\n return [\n begin\n @unpack prog, aux, ∇flux, hyperdiff = local_states\n diffusive = ∇flux\n state = prog\n hyperdiffusive = hyperdiff\n\n _args_fx1 = (; state, aux, t, direction)\n _args_fx2 = (; state, aux, t, diffusive, hyperdiffusive)\n _args_src = (; state, aux, t, direction, diffusive)\n\n cache_fx1 = precompute(bl, _args_fx1, Flux{FirstOrder}())\n cache_fx2 = precompute(bl, _args_fx2, Flux{SecondOrder}())\n cache_src = precompute(bl, _args_src, Source())\n\n z = altitude(bl, aux)\n\n nt = (;\n z = altitude(bl, aux),\n prog = flattened_named_tuple(prog), # Vars -> flattened NamedTuples\n aux = flattened_named_tuple(aux), # Vars -> flattened NamedTuples\n ∇flux = flattened_named_tuple(∇flux), # Vars -> flattened NamedTuples\n hyperdiff = flattened_named_tuple(hyperdiff), # Vars -> flattened NamedTuples\n cache_fx1,\n cache_fx2,\n cache_src,\n )\n # Flatten top level:\n flattened_named_tuple(nt)\n end for local_states in NodalStack(bl, grid; kwargs...)\n ]\nend\n" }, { "path": "src/Utilities/TicToc/TicToc.jl", "content": "\"\"\"\n TicToc -- timing measurement\n\nLow-overhead time interval measurement via minimally invasive macros.\n\n\"\"\"\nmodule TicToc\n\nusing Printf\n\nexport @tic, @toc, tictoc\n\n# explicitly enable due to issues with pre-compilation\nconst tictoc_enabled = false\n\n# disable to reduce overhead\nconst tictoc_track_memory = true\n\nif tictoc_track_memory\n\n mutable struct TimingInfo\n ncalls::Int\n time::UInt64\n allocd::Int64\n gctime::UInt64\n curr::UInt64\n mem::Base.GC_Num\n end\n TimingInfo() = TimingInfo(\n 0,\n 0,\n 0,\n 0,\n 0,\n Base.GC_Num(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0),\n )\n\nelse # !tictoc_track_memory\n\n mutable struct TimingInfo\n ncalls::Int\n time::UInt64\n curr::UInt64\n end\n TimingInfo() = TimingInfo(0, 0, 0)\n\nend # if tictoc_track_memory\n\nconst timing_infos = TimingInfo[]\nconst timing_info_names = Symbol[]\nconst atexit_function_registered = Ref(false)\n\n# `@tic` helper\nfunction _tic(nm)\n @static if !tictoc_enabled\n return quote end\n end\n exti = Symbol(\"tictoc__\", nm)\n global timing_info_names\n if exti in timing_info_names\n err_ex = quote\n throw(ArgumentError(\"$(nm) already used in @tic\"))\n end\n else\n err_ex = quote end\n end\n push!(timing_info_names, exti)\n @static if tictoc_track_memory\n quote\n $(err_ex)\n global $(exti)\n $(exti).curr = time_ns()\n $(exti).mem = Base.gc_num()\n end\n else\n quote\n $(err_ex)\n global $(exti)\n $(exti).curr = time_ns()\n end\n end\nend\n\n\"\"\"\n @tic nm\n\nIndicate the start of the interval `nm`.\n\"\"\"\nmacro tic(args...)\n na = length(args)\n if na != 1\n throw(ArgumentError(\"wrong number of arguments in @tic\"))\n end\n ex = args[1]\n if !isa(ex, Symbol)\n throw(ArgumentError(\"need a name argument to @tic\"))\n end\n return _tic(ex)\nend\n\n# `@toc` helper\nfunction _toc(nm)\n @static if !tictoc_enabled\n return quote end\n end\n exti = Symbol(\"tictoc__\", nm)\n @static if tictoc_track_memory\n quote\n global $(exti)\n $(exti).time += time_ns() - $(exti).curr\n $(exti).ncalls += 1\n local diff = Base.GC_Diff(Base.gc_num(), $(exti).mem)\n $(exti).allocd += diff.allocd\n $(exti).gctime += diff.total_time\n end\n else\n quote\n global $(exti)\n $(exti).time += time_ns() - $(exti).curr\n $(exti).ncalls += 1\n end\n end\nend\n\n\"\"\"\n @toc nm\n\nIndicate the end of the interval `nm`.\n\"\"\"\nmacro toc(args...)\n na = length(args)\n if na != 1\n throw(ArgumentError(\"wrong number of arguments in @toc\"))\n end\n ex = args[1]\n if !isa(ex, Symbol)\n throw(ArgumentError(\"need a name argument to @toc\"))\n end\n return _toc(ex)\nend\n\n\"\"\"\n print_timing_info()\n\n`atexit()` function, writes all information about every interval to\n`stdout`.\n\"\"\"\nfunction print_timing_info()\n println(\"TicToc timing information\")\n @static if tictoc_track_memory\n println(\"name,ncalls,tottime(ns),allocbytes,gctime\")\n else\n println(\"name,ncalls,tottime(ns)\")\n end\n for i in 1:length(timing_info_names)\n @static if tictoc_track_memory\n s = @sprintf(\n \"%s,%d,%d,%d,%d\",\n timing_info_names[i],\n timing_infos[i].ncalls,\n timing_infos[i].time,\n timing_infos[i].allocd,\n timing_infos[i].gctime\n )\n else\n s = @sprintf(\n \"%s,%d,%d\",\n timing_info_names[i],\n timing_infos[i].ncalls,\n timing_infos[i].time\n )\n end\n println(s)\n end\nend\n\n\"\"\"\n tictoc()\n\nCall at program start (only once!) to set up the globals used by the\nmacros and to register the at-exit callback.\n\"\"\"\nfunction tictoc()\n @static if !tictoc_enabled\n return 0\n end\n global timing_info_names\n for nm in timing_info_names\n exti = Symbol(nm)\n isdefined(@__MODULE__, exti) && continue\n expr = quote\n const $exti = $TimingInfo()\n end\n eval(Expr(:toplevel, expr))\n push!(timing_infos, getfield(@__MODULE__, exti))\n end\n if parse(Int, get(ENV, \"TICTOC_PRINT_RESULTS\", \"0\")) == 1\n if !atexit_function_registered[]\n atexit(print_timing_info)\n atexit_function_registered[] = true\n end\n end\n return length(timing_info_names)\nend\n\nend # module\n" }, { "path": "src/Utilities/VariableTemplates/VariableTemplates.jl", "content": "module VariableTemplates\n\nexport varsize, Vars, Grad, @vars, varsindex, varsindices\n\nusing StaticArrays\nusing LinearAlgebra\n\n\"\"\"\n varsindex(S, p::Symbol, [sp::Symbol...])\n\nReturn a range of indices corresponding to the property `p` and\n(optionally) its subproperties `sp` based on the template type `S`.\n\n# Examples\n```julia-repl\njulia> S = @vars(x::Float64, y::Float64)\njulia> varsindex(S, :y)\n2:2\n\njulia> S = @vars(x::Float64, y::@vars(α::Float64, β::SVector{3, Float64}))\njulia> varsindex(S, :y, :β)\n3:5\n```\n\"\"\"\nfunction varsindex(::Type{S}, insym::Symbol) where {S <: NamedTuple}\n offset = 0\n for varsym in fieldnames(S)\n T = fieldtype(S, varsym)\n if T <: Real\n offset += 1\n varrange = offset:offset\n elseif T <: SHermitianCompact\n LT = StaticArrays.lowertriangletype(T)\n N = length(LT)\n varrange = offset .+ (1:N)\n offset += N\n elseif T <: StaticArray\n N = length(T)\n varrange = offset .+ (1:N)\n offset += N\n else\n varrange = offset .+ (1:varsize(T))\n offset += varsize(T)\n end\n if insym == varsym\n return varrange\n end\n end\n error(\"symbol '$insym' not found\")\nend\n\n# return `Symbol`s unchanged.\nwrap_val(sym) = sym\n# wrap integer in `Val`\nwrap_val(i::Int) = Val(i)\n\n# We enforce that calls to `varsindex` on\n# an `NTuple` must be unrapped in `Val`.\n# This is enforced to synchronize failures\n# on the CPU and GPU, rather than allowing\n# CPU-working and GPU-breaking versions.\n# This means that users _must_ wrap `sym`\n# in `Val`, which can be done with `wrap_val`\n# above.\nBase.@propagate_inbounds function varsindex(\n ::Type{S},\n sym::Symbol,\n rest...,\n) where {S <: Union{NamedTuple, Tuple}}\n vi = varsindex(fieldtype(S, sym), rest...)\n return varsindex(S, sym)[vi]\nend\nBase.@propagate_inbounds function varsindex(\n ::Type{S},\n ::Val{i},\n rest...,\n) where {i, S <: Union{NamedTuple, Tuple}}\n et = eltype(S)\n offset = (i - 1) * varsize(et)\n vi = varsindex(et, rest...)\n return (vi.start + offset):(vi.stop + offset)\nend\n\nBase.@propagate_inbounds function varsindex(\n ::Type{S},\n ::Val{i},\n) where {i, S <: SArray}\n return i:i\nend\n\n\"\"\"\n varsindices(S, ps::Tuple)\n varsindices(S, ps...)\n\nReturn a tuple of indices corresponding to the properties\nspecified by `ps` based on the template type `S`. Properties\ncan be specified using either symbols or strings.\n\n# Examples\n```julia-repl\njulia> S = @vars(x::Float64, y::Float64, z::Float64)\njulia> varsindices(S, (:x, :z))\n(1, 3)\n\njulia> S = @vars(x::Float64, y::@vars(α::Float64, β::SVector{3, Float64}))\njulia> varsindices(S, \"x\", \"y.β\")\n(1, 3, 4, 5)\n```\n\"\"\"\nfunction varsindices(::Type{S}, vars::Tuple) where {S <: NamedTuple}\n indices = Int[]\n for var in vars\n splitvar = split(string(var), '.')\n append!(indices, collect(varsindex(S, map(Symbol, splitvar)...)))\n end\n Tuple(indices)\nend\nvarsindices(::Type{S}, vars...) where {S <: NamedTuple} = varsindices(S, vars)\n\n\"\"\"\n varsize(S)\n\nThe number of elements specified by the template type `S`.\n\"\"\"\nvarsize(::Type{T}) where {T <: Real} = 1\nvarsize(::Type{Tuple{}}) = 0\nvarsize(::Type{NamedTuple{(), Tuple{}}}) = 0\nvarsize(::Type{SArray{S, T, N} where L}) where {S, T, N} = prod(S.parameters)\n\ninclude(\"var_names.jl\")\n\n# TODO: should be possible to get rid of @generated\n@generated function varsize(::Type{S}) where {S}\n types = fieldtypes(S)\n isempty(types) ? 0 : sum(varsize, types)\nend\n\nfunction process_vars!(syms, typs, expr)\n if expr isa LineNumberNode\n return\n elseif expr isa Expr && expr.head == :block\n for arg in expr.args\n process_vars!(syms, typs, arg)\n end\n return\n elseif expr.head == :(::)\n push!(syms, expr.args[1])\n push!(typs, expr.args[2])\n return\n else\n error(\"Invalid expression\")\n end\nend\n\n\"\"\"\n @vars(var1::Type1, var2::Type2)\n\nA convenient syntax for describing a `NamedTuple` type.\n\n# Example\n```julia\njulia> @vars(a::Float64, b::Float64)\nNamedTuple{(:a, :b),Tuple{Float64,Float64}}\n```\n\"\"\"\nmacro vars(args...)\n syms = Any[]\n typs = Any[]\n for arg in args\n process_vars!(syms, typs, arg)\n end\n :(NamedTuple{$(tuple(syms...)), Tuple{$(esc.(typs)...)}})\nend\n\nstruct GetVarError <: Exception\n sym::Symbol\nend\nstruct SetVarError <: Exception\n sym::Symbol\nend\n\nabstract type AbstractVars{S, A, offset} end\n\n\"\"\"\n Vars{S,A,offset}(array::A)\n\nDefines property overloading for `array` using the type `S` as a template. `offset` is used to shift the starting element of the array.\n\"\"\"\nstruct Vars{S, A, offset} <: AbstractVars{S, A, offset}\n array::A\nend\nVars{S}(array) where {S} = Vars{S, typeof(array), 0}(array)\n\nBase.parent(v::AbstractVars) = getfield(v, :array)\nBase.eltype(v::AbstractVars) = eltype(parent(v))\nBase.propertynames(::AbstractVars{S}) where {S} = fieldnames(S)\nBase.similar(v::AbstractVars) = typeof(v)(similar(parent(v)))\n\n@generated function Base.getproperty(\n v::Vars{S, A, offset},\n sym::Symbol,\n) where {S, A, offset}\n expr = quote\n Base.@_inline_meta\n array = parent(v)\n end\n for k in fieldnames(S)\n T = fieldtype(S, k)\n if T <: Real\n retexpr = :($T(array[$(offset + 1)]))\n offset += 1\n elseif T <: SHermitianCompact\n LT = StaticArrays.lowertriangletype(T)\n N = length(LT)\n retexpr = :($T($LT($([:(array[$(offset + i)]) for i in 1:N]...))))\n offset += N\n elseif T <: StaticArray\n N = length(T)\n retexpr = :($T($([:(array[$(offset + i)]) for i in 1:N]...)))\n offset += N\n else\n retexpr = :(Vars{$T, A, $offset}(array))\n offset += varsize(T)\n end\n push!(expr.args, :(\n if sym == $(QuoteNode(k))\n return @inbounds $retexpr\n end\n ))\n end\n push!(expr.args, :(throw(GetVarError(sym))))\n expr\nend\n\n@generated function Base.setproperty!(\n v::Vars{S, A, offset},\n sym::Symbol,\n val,\n) where {S, A, offset}\n expr = quote\n Base.@_inline_meta\n array = parent(v)\n end\n for k in fieldnames(S)\n T = fieldtype(S, k)\n if T <: Real\n retexpr = :(array[$(offset + 1)] = val)\n offset += 1\n elseif T <: SHermitianCompact\n LT = StaticArrays.lowertriangletype(T)\n N = length(LT)\n retexpr = :(\n array[($(offset + 1)):($(offset + N))] .=\n $T(val).lowertriangle\n )\n offset += N\n elseif T <: StaticArray\n N = length(T)\n retexpr = :(array[($(offset + 1)):($(offset + N))] .= val[:])\n offset += N\n else\n offset += varsize(T)\n continue\n end\n push!(expr.args, :(\n if sym == $(QuoteNode(k))\n return @inbounds $retexpr\n end\n ))\n end\n push!(expr.args, :(throw(SetVarError(sym))))\n expr\nend\n\n\"\"\"\n Grad{S,A,offset}(array::A)\n\nDefines property overloading along slices of the second dimension of `array` using the type `S` as a template. `offset` is used to shift the starting element of the array.\n\"\"\"\nstruct Grad{S, A, offset} <: AbstractVars{S, A, offset}\n array::A\nend\nGrad{S}(array) where {S} = Grad{S, typeof(array), 0}(array)\n\n@generated function Base.getproperty(\n v::Grad{S, A, offset},\n sym::Symbol,\n) where {S, A, offset}\n if A <: SubArray\n M = size(fieldtype(A, 1), 1)\n else\n M = size(A, 1)\n end\n expr = quote\n Base.@_inline_meta\n array = parent(v)\n end\n for k in fieldnames(S)\n T = fieldtype(S, k)\n if T <: Real\n retexpr = :(SVector{$M, $T}(\n $([:(array[$i, $(offset + 1)]) for i in 1:M]...),\n ))\n offset += 1\n elseif T <: StaticArray\n N = length(T)\n retexpr = :(SMatrix{$M, $N, $(eltype(T))}(\n $([:(array[$i, $(offset + j)]) for i in 1:M, j in 1:N]...),\n ))\n offset += N\n else\n retexpr = :(Grad{$T, A, $offset}(array))\n offset += varsize(T)\n end\n push!(expr.args, :(\n if sym == $(QuoteNode(k))\n return @inbounds $retexpr\n end\n ))\n end\n push!(expr.args, :(throw(GetVarError(sym))))\n expr\nend\n\n@generated function Base.setproperty!(\n v::Grad{S, A, offset},\n sym::Symbol,\n val::AbstractArray,\n) where {S, A, offset}\n if A <: SubArray\n M = size(fieldtype(A, 1), 1)\n else\n M = size(A, 1)\n end\n expr = quote\n Base.@_inline_meta\n array = parent(v)\n end\n for k in fieldnames(S)\n T = fieldtype(S, k)\n if T <: Real\n retexpr = :(array[:, $(offset + 1)] = val)\n offset += 1\n elseif T <: StaticArray\n N = length(T)\n retexpr = :(\n array[\n :,\n # static range is used here to force dispatch to\n # StaticArrays setindex! because generic setindex! is slow\n StaticArrays.SUnitRange($(offset + 1), $(offset + N)),\n ] = val\n )\n offset += N\n else\n offset += varsize(T)\n continue\n end\n push!(expr.args, :(\n if sym == $(QuoteNode(k))\n return @inbounds $retexpr\n end\n ))\n end\n push!(expr.args, :(throw(SetVarError(sym))))\n expr\nend\n\nexport unroll_map, @unroll_map\n\"\"\"\n @unroll_map(f::F, N::Int, args...) where {F}\n unroll_map(f::F, N::Int, args...) where {F}\n\nUnroll N-expressions and wrap arguments in `Val`.\n\"\"\"\n@generated function unroll_map(f::F, ::Val{N}, args...) where {F, N}\n quote\n Base.@_inline_meta\n Base.Cartesian.@nexprs $N i -> f(Val(i), args...)\n end\nend\nmacro unroll_map(func, N, args...)\n @assert func.head == :(->)\n body = func.args[2]\n pushfirst!(body.args, :(Base.@_inline_meta))\n quote\n $unroll_map($(esc(func)), Val($(esc(N))), $(esc(args))...)\n end\nend\n\nexport vuntuple\n\"\"\"\n vuntuple(f::F, N::Int)\n\nVal-Unroll ntuple: wrap `ntuple`\narguments in `Val` for unrolling.\n\"\"\"\nvuntuple(f::F, N::Int) where {F} = ntuple(i -> f(Val(i)), Val(N))\n\n# Inside unroll_map expressions, all indexes `i`\n# are wrapped in `Val`, so we must redirect\n# these methods:\nBase.@propagate_inbounds Base.getindex(t::Tuple, ::Val{i}) where {i} =\n Base.getindex(t, i)\nBase.@propagate_inbounds Base.getindex(a::SArray, ::Val{i}) where {i} =\n Base.getindex(a, i)\n\nBase.@propagate_inbounds function Base.getindex(\n v::Vars{NTuple{N, T}, A, offset},\n ::Val{i},\n) where {N, T, A, offset, i} # 1 <= i <= N\n return Vars{T, A, offset + (i - 1) * varsize(T)}(parent(v))\nend\n\nBase.@propagate_inbounds function Base.getindex(\n v::Grad{NTuple{N, T}, A, offset},\n ::Val{i},\n) where {N, T, A, offset, i} # 1 <= i <= N\n return Grad{T, A, offset + (i - 1) * varsize(T)}(parent(v))\nend\n\n\"\"\"\n getpropertyorindex\n\nAn interchangeably and nested-friendly\n`getproperty`/`getindex`.\n\"\"\"\nfunction getpropertyorindex end\n\n# Redirect to Base getproperty/getindex:\nBase.@propagate_inbounds getpropertyorindex(t::Tuple, ::Val{i}) where {i} =\n Base.getindex(t, i)\nBase.@propagate_inbounds getpropertyorindex(\n a::AbstractArray,\n ::Val{i},\n) where {i} = Base.getindex(a, i)\nBase.@propagate_inbounds getpropertyorindex(v::AbstractVars, s::Symbol) =\n Base.getproperty(v, s)\nBase.@propagate_inbounds getpropertyorindex(\n v::AbstractVars,\n ::Val{i},\n) where {i} = Base.getindex(v, Val(i))\n\n# Only one element left:\nBase.@propagate_inbounds getpropertyorindex(\n v::AbstractVars,\n t::Tuple{A},\n) where {A} = getpropertyorindex(v, t[1])\nBase.@propagate_inbounds getpropertyorindex(\n a::AbstractArray,\n t::Tuple{A},\n) where {A} = getpropertyorindex(a, t[1])\n\n# Peel first element from tuple and recurse:\nBase.@propagate_inbounds getpropertyorindex(v::AbstractVars, t::Tuple) =\n getpropertyorindex(getpropertyorindex(v, t[1]), Tuple(t[2:end]))\n\n# Redirect to getpropertyorindex:\nBase.@propagate_inbounds Base.getproperty(v::AbstractVars, tup_chain::Tuple) =\n getpropertyorindex(v, tup_chain)\nBase.@propagate_inbounds Base.getindex(v::AbstractVars, tup_chain::Tuple) =\n getpropertyorindex(v, tup_chain)\n\ninclude(\"flattened_tup_chain.jl\")\n\nfunction Base.show(io::IO, v::AbstractVars)\n s = \"$(nameof(typeof(v))) object:\\n\"\n for tup in flattened_tup_chain(v, RetainArr())\n name = join(tup, \"_\")\n val = getproperty(v, wrap_val.(tup))\n s *= \" $name = $val\\n\"\n end\n print(io, s)\nend\n\nend # module\n" }, { "path": "src/Utilities/VariableTemplates/flattened_tup_chain.jl", "content": "using LinearAlgebra\n\nexport flattened_tup_chain, flattened_named_tuple, flattened_tuple\nexport FlattenType, FlattenArr, RetainArr\n\nabstract type FlattenType end\n\n\"\"\"\n FlattenArr\n\nFlatten arrays in `flattened_tup_chain`\nand `flattened_named_tuple`.\n\"\"\"\nstruct FlattenArr <: FlattenType end\n\n\"\"\"\n RetainArr\n\nDo _not_ flatten arrays in `flattened_tup_chain`\nand `flattened_named_tuple`.\n\"\"\"\nstruct RetainArr <: FlattenType end\n\n# The Vars instance has many empty entries.\n# Keeping all of the keys results in many\n# duplicated values. So, it's best we\n# \"prune\" the tree by removing the keys:\nflattened_tup_chain(\n ::Type{NamedTuple{(), Tuple{}}},\n ::FlattenType = FlattenArr();\n prefix = (Symbol(),),\n) = ()\n\nflattened_tup_chain(\n ::Type{T},\n ::FlattenType;\n prefix = (Symbol(),),\n) where {T <: Real} = (prefix,)\n\nflattened_tup_chain(\n ::Type{T},\n ::RetainArr;\n prefix = (Symbol(),),\n) where {T <: SArray} = (prefix,)\nflattened_tup_chain(\n ::Type{T},\n ::FlattenArr;\n prefix = (Symbol(),),\n) where {T <: SArray} = ntuple(i -> (prefix..., i), length(T))\n\nflattened_tup_chain(\n ::Type{T},\n ::RetainArr;\n prefix = (Symbol(),),\n) where {T <: SHermitianCompact} = (prefix,)\nflattened_tup_chain(\n ::Type{T},\n ::FlattenType;\n prefix = (Symbol(),),\n) where {T <: SHermitianCompact} =\n ntuple(i -> (prefix..., i), length(StaticArrays.lowertriangletype(T)))\n\nflattened_tup_chain(\n ::Type{T},\n ::RetainArr;\n prefix = (Symbol(),),\n) where {N, TA, T <: Diagonal{N, TA}} = (prefix,)\nflattened_tup_chain(\n ::Type{T},\n ::FlattenArr;\n prefix = (Symbol(),),\n) where {N, TA, T <: Diagonal{N, TA}} = ntuple(i -> (prefix..., i), length(TA))\n\nflattened_tup_chain(::Type{T}, ::FlattenType; prefix = (Symbol(),)) where {T} =\n (prefix,)\n\n\"\"\"\n flattened_tup_chain(::Type{T}) where {T <: Union{NamedTuple,NTuple}}\n\nAn array of tuples, containing symbols\nand integers for every combination of\neach field in the `Vars` array.\n\"\"\"\nfunction flattened_tup_chain(\n ::Type{T},\n ft::FlattenType = FlattenArr();\n prefix = (Symbol(),),\n) where {T <: Union{NamedTuple, NTuple}}\n map(1:fieldcount(T)) do i\n Ti = fieldtype(T, i)\n name = fieldname(T, i)\n sname = name isa Int ? name : Symbol(name)\n flattened_tup_chain(\n Ti,\n ft;\n prefix = prefix == (Symbol(),) ? (sname,) : (prefix..., sname),\n )\n end |>\n Iterators.flatten |>\n collect\nend\nflattened_tup_chain(\n ::AbstractVars{S},\n ft::FlattenType = FlattenArr(),\n) where {S} = flattened_tup_chain(S, ft)\n\n\"\"\"\n flattened_named_tuple\n\nA flattened NamedTuple, given a\n`Vars` or nested `NamedTuple` instance.\n\n# Example:\n\n```julia\nusing Test\nusing ClimateMachine.VariableTemplates\nnt = (x = 1, a = (y = 2, z = 3, b = ((a = 1,), (a = 2,), (a = 3,))));\nfnt = flattened_named_tuple(nt);\n@test keys(fnt) == (:x, :a_y, :a_z, :a_b_1_a, :a_b_2_a, :a_b_3_a)\n@test length(fnt) == 6\n@test fnt.x == 1\n@test fnt.a_y == 2\n@test fnt.a_z == 3\n@test fnt.a_b_1_a == 1\n@test fnt.a_b_2_a == 2\n@test fnt.a_b_3_a == 3\n```\n\"\"\"\nfunction flattened_named_tuple end\n\nfunction flattened_named_tuple(v::AbstractVars, ft::FlattenType = FlattenArr())\n ftc = flattened_tup_chain(v, ft)\n keys_ = Symbol.(join.(ftc, :_))\n vals = map(x -> getproperty(v, wrap_val.(x)), ftc)\n length(keys_) == length(vals) || error(\"key-value mismatch\")\n return (; zip(keys_, vals)...)\nend\n\nfunction flattened_named_tuple(nt::NamedTuple, ft::FlattenType = FlattenArr())\n ftc = flattened_tup_chain(typeof(nt), ft)\n keys_ = Symbol.(join.(ftc, :_))\n vals = flattened_tuple(ft, nt)\n length(keys_) == length(vals) || error(\"key-value mismatch\")\n return (; zip(keys_, vals)...)\nend\n\nflattened_tuple(::FlattenArr, a::AbstractArray) = tuple(a...)\nflattened_tuple(::RetainArr, a::AbstractArray) = tuple(a)\n\nflattened_tuple(::FlattenArr, a::Diagonal) = tuple(a.diag...)\nflattened_tuple(::RetainArr, a::Diagonal) = tuple(a.diag)\n\nflattened_tuple(::FlattenArr, a::SHermitianCompact) = tuple(a.lowertriangle...)\nflattened_tuple(::RetainArr, a::SHermitianCompact) = tuple(a.lowertriangle)\n\n# when we splat an empty tuple `b` into `flattened_tuple(ft, b...)`\nflattened_tuple(::FlattenType) = ()\n\n# for structs\nflattened_tuple(::FlattenType, a) = (a,)\n\n# Divide and conquer:\nflattened_tuple(ft::FlattenType, a, b...) =\n tuple(flattened_tuple(ft, a)..., flattened_tuple(ft, b...)...)\n\nflattened_tuple(ft::FlattenType, a::Tuple) = flattened_tuple(ft, a...)\n\nflattened_tuple(ft::FlattenType, a::NamedTuple) = flattened_tuple(ft, Tuple(a))\n" }, { "path": "src/Utilities/VariableTemplates/var_names.jl", "content": "export flattenednames\n\nflattenednames(nt::Type{NTuple{N, T}}; prefix = \"\") where {N, T} =\n Iterators.flatten([\n flattenednames(T; prefix = \"$(prefix)[$i]\") for i in 1:N\n ]) |> collect\nflattenednames(::Type{NamedTuple{(), Tuple{}}}; prefix = \"\") = ()\nflattenednames(::Type{T}; prefix = \"\") where {T <: Real} = (prefix,)\nflattenednames(::Type{T}; prefix = \"\") where {T <: SArray} =\n ntuple(i -> \"$prefix[$i]\", length(T))\nfunction flattenednames(::Type{T}; prefix = \"\") where {T <: SHermitianCompact}\n N = size(T, 1)\n [[\"$prefix[$i,$j]\" for i in j:N] for j in 1:N] |>\n Iterators.flatten |>\n collect\nend\nfunction flattenednames(::Type{T}; prefix = \"\") where {T <: NamedTuple}\n map(1:fieldcount(T)) do i\n Ti = fieldtype(T, i)\n name = fieldname(T, i)\n flattenednames(\n Ti,\n prefix = prefix == \"\" ? string(name) : string(prefix, '.', name),\n )\n end |>\n Iterators.flatten |>\n collect\nend\n" }, { "path": "test/Arrays/basics.jl", "content": "using Test, MPI\n\nusing ClimateMachine\nusing ClimateMachine.MPIStateArrays\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\nconst mpicomm = MPI.COMM_WORLD\n\n@testset \"MPIStateArray basics\" begin\n Q = MPIStateArray{Float32}(mpicomm, ArrayType, 4, 6, 8)\n\n @test eltype(Q) == Float32\n @test size(Q) == (4, 6, 8)\n\n fillval = 0.5f0\n fill!(Q, fillval)\n\n ClimateMachine.gpu_allowscalar(true)\n @test Q[1] == fillval\n @test Q[2, 3, 4] == fillval\n @test Q[end] == fillval\n\n @test Array(Q) == fill(fillval, 4, 6, 8)\n\n Q[2, 3, 4] = 2fillval\n @test Q[2, 3, 4] != fillval\n ClimateMachine.gpu_allowscalar(false)\n\n Qp = copy(Q)\n\n @test typeof(Qp) == typeof(Q)\n @test eltype(Qp) == eltype(Q)\n @test size(Qp) == size(Q)\n @test Array(Qp) == Array(Q)\n\n Qp = similar(Q)\n\n @test typeof(Qp) == typeof(Q)\n @test eltype(Qp) == eltype(Q)\n @test size(Qp) == size(Q)\n\n copyto!(Qp, Q)\n @test Array(Qp) == Array(Q)\nend\n" }, { "path": "test/Arrays/broadcasting.jl", "content": "using Test, MPI\n\nusing ClimateMachine\nusing ClimateMachine.MPIStateArrays\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\nconst mpicomm = MPI.COMM_WORLD\n\n@testset \"MPIStateArray broadcasting\" begin\n let\n localsize = (4, 6, 8)\n A = rand(Float32, localsize)\n B = rand(Float32, localsize)\n\n QA = MPIStateArray{Float32}(mpicomm, ArrayType, localsize...)\n QB = similar(QA)\n\n QA .= A\n QB .= B\n\n @test Array(QA) == A\n @test Array(QB) == B\n\n QC = QA .+ QB\n @test typeof(QC) == typeof(QA)\n C = Array(QC)\n @test C == A .+ B\n\n QC = QA .+ sqrt.(QB)\n C = Array(QC)\n @test C ≈ A .+ sqrt.(B)\n\n QC = QA .+ sqrt.(QB) .* exp.(QA .- QB .^ 2)\n C = Array(QC)\n @test C ≈ A .+ sqrt.(B) .* exp.(A .- B .^ 2)\n\n # writing to an existing array instead of creating a new one\n fill!(QC, 0)\n QC .= QA .+ sqrt.(QB) .* exp.(QA .- QB .^ 2)\n C = Array(QC)\n @test C ≈ A .+ sqrt.(B) .* exp.(A .- B .^ 2)\n end\n\n let\n numelems = 12\n realelems = 1:7\n ghostelems = 8:12\n\n QA = MPIStateArray{Int}(\n mpicomm,\n ArrayType,\n 1,\n 1,\n numelems,\n realelems = realelems,\n ghostelems = ghostelems,\n )\n QB = similar(QA)\n\n fill!(QA, 1)\n fill!(QB, 3)\n\n QB .= QA .+ QB\n\n @test all(Array(QB)[realelems] .== 4)\n @test all(Array(QB)[ghostelems] .== 3)\n end\nend\n" }, { "path": "test/Arrays/mpi_comm.jl", "content": "using Test\nusing MPI\nusing ClimateMachine\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.Mesh.BrickMesh\nusing Pkg\nusing KernelAbstractions\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\nconst comm = MPI.COMM_WORLD\n\n\nfunction main()\n\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n\n @assert csize == 3\n\n if crank == 0\n numreal = 4\n numghost = 3\n\n nabrtorank = [1, 2]\n\n sendelems = [1, 2, 3, 4, 1, 4]\n nabrtorecv = [1:2, 3:3]\n nabrtosend = [1:4, 5:6]\n\n vmaprecv = [\n 37,\n 38,\n 39,\n 40,\n 42,\n 43,\n 44,\n 45,\n 46,\n 49,\n 52,\n 53,\n 54,\n 57,\n 60,\n 61,\n 62,\n 63,\n ]\n vmapsend =\n [3, 6, 9, 10, 11, 12, 19, 22, 25, 34, 35, 36, 1, 2, 3, 28, 31, 34]\n\n nabrtovmaprecv = [1:13, 14:18]\n nabrtovmapsend = [1:12, 13:18]\n\n expectedghostdata = [\n 1001,\n 1002,\n 1003,\n 1004,\n 1006,\n 1007,\n 1008,\n 1009,\n 1010,\n 1013,\n 1016,\n 1017,\n 1018,\n 2003,\n 2006,\n 2007,\n 2008,\n 2009,\n ]\n elseif crank == 1\n numreal = 2\n numghost = 4\n\n nabrtorank = [0]\n\n sendelems = [1, 2]\n nabrtorecv = [1:4]\n nabrtosend = [1:2]\n\n vmaprecv = [21, 24, 27, 28, 29, 30, 37, 40, 43, 52, 53, 54]\n vmapsend = [1, 2, 3, 4, 6, 7, 8, 9, 10, 13, 16, 17, 18]\n\n nabrtovmaprecv = [1:length(vmaprecv)]\n nabrtovmapsend = [1:length(vmapsend)]\n\n expectedghostdata = [3, 6, 9, 10, 11, 12, 19, 22, 25, 34, 35, 36]\n elseif crank == 2\n numreal = 1\n numghost = 2\n\n nabrtorank = [0]\n\n sendelems = [1]\n nabrtorecv = [1:2]\n nabrtosend = [1:1]\n\n vmaprecv = [10, 11, 12, 19, 22, 25]\n vmapsend = [3, 6, 7, 8, 9]\n\n nabrtovmaprecv = [1:length(vmaprecv)]\n nabrtovmapsend = [1:length(vmapsend)]\n\n expectedghostdata = [1, 2, 3, 28, 31, 34]\n end\n\n numelem = numreal + numghost\n\n realelems = 1:numreal\n ghostelems = numreal .+ (1:numghost)\n\n weights = Array{Int64}(undef, (0, 0, 0))\n\n A = MPIStateArray{Int64}(\n comm,\n ArrayType,\n 9,\n 2,\n numelem,\n realelems,\n ghostelems,\n ArrayType(vmaprecv),\n ArrayType(vmapsend),\n nabrtorank,\n nabrtovmaprecv,\n nabrtovmapsend,\n ArrayType(weights),\n )\n\n Q = Array(A.data)\n Q .= -1\n shift = 100\n Q[:, 1, realelems] .=\n reshape((crank * 1000) .+ (1:(9 * numreal)), 9, numreal)\n Q[:, 2, realelems] .=\n reshape((crank * 1000) .+ shift .+ (1:(9 * numreal)), 9, numreal)\n copyto!(A.data, Q)\n\n event = Event(array_device(A))\n event = MPIStateArrays.begin_ghost_exchange!(A; dependencies = event)\n event = MPIStateArrays.end_ghost_exchange!(A; dependencies = event)\n wait(array_device(A), event)\n\n Q = Array(A.data)\n @test all(expectedghostdata .== Q[:, 1, :][:][vmaprecv])\n @test all(shift .+ expectedghostdata .== Q[:, 2, :][:][vmaprecv])\n\nend\n\nmain()\n" }, { "path": "test/Arrays/reductions.jl", "content": "using Test, MPI\nusing LinearAlgebra\n\nusing ClimateMachine\nusing ClimateMachine.MPIStateArrays\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\nconst mpicomm = MPI.COMM_WORLD\n\nmpisize = MPI.Comm_size(mpicomm)\nmpirank = MPI.Comm_rank(mpicomm)\n\n@testset \"MPIStateArray reductions\" begin\n\n localsize = (4, 6, 8)\n\n A = Array{Float32}(reshape(1:prod(localsize), localsize))\n globalA = vcat([A for _ in 1:mpisize]...)\n\n QA = MPIStateArray{Float32}(mpicomm, ArrayType, localsize...)\n QA .= A\n\n\n @test norm(QA, 1) ≈ norm(globalA, 1)\n @test norm(QA) ≈ norm(globalA)\n @test norm(QA, Inf) ≈ norm(globalA, Inf)\n\n @test norm(QA; dims = (1, 3)) ≈ mapslices(norm, globalA; dims = (1, 3))\n @test norm(QA, 1; dims = (1, 3)) ≈\n mapslices(S -> norm(S, 1), globalA, dims = (1, 3))\n @test norm(QA, Inf; dims = (1, 3)) ≈\n mapslices(S -> norm(S, Inf), globalA, dims = (1, 3))\n\n B = Array{Float32}(reshape(reverse(1:prod(localsize)), localsize))\n globalB = vcat([B for _ in 1:mpisize]...)\n\n QB = similar(QA)\n QB .= B\n\n @test isapprox(euclidean_distance(QA, QB), norm(globalA .- globalB))\n @test isapprox(dot(QA, QB), dot(globalA, globalB))\n\n C = fill(Float32(mpirank + 1), localsize)\n globalC = vcat([fill(i, localsize) for i in 1:mpisize]...)\n QC = similar(QA)\n QC .= C\n\n @test sum(QC) == sum(globalC)\n @test Array(sum(QC; dims = (1, 3))) == sum(globalC; dims = (1, 3))\n @test maximum(QC) == maximum(globalC)\n @test Array(maximum(QC; dims = (1, 3))) == maximum(globalC; dims = (1, 3))\n @test minimum(QC) == minimum(globalC)\n @test Array(minimum(QC; dims = (1, 3))) == minimum(globalC; dims = (1, 3))\nend\n" }, { "path": "test/Arrays/runtests.jl", "content": "using Test\ninclude(joinpath(\"..\", \"testhelpers.jl\"))\n\n@testset \"MPIStateArrays reductions\" begin\n runmpi(joinpath(@__DIR__, \"basics.jl\"))\n runmpi(joinpath(@__DIR__, \"broadcasting.jl\"))\n runmpi(joinpath(@__DIR__, \"reductions.jl\"))\n runmpi(joinpath(@__DIR__, \"reductions.jl\"), ntasks = 3)\n runmpi(joinpath(@__DIR__, \"varsindex.jl\"))\nend\n" }, { "path": "test/Arrays/varsindex.jl", "content": "using Test, MPI\n\nusing ClimateMachine\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.MPIStateArrays: getstateview\nusing ClimateMachine.VariableTemplates: @vars, varsindex, varsindices\nusing StaticArrays\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\nconst mpicomm = MPI.COMM_WORLD\n\nconst V = @vars begin\n a::Float32\n b::SVector{3, Float32}\n c::SMatrix{3, 8, Float32}\n d::Float32\n e::@vars begin\n a::Float32\n b::SVector{3, Float32}\n d::Float32\n end\nend\n\nconst VNT = @vars begin\n a::Float32\n e::Tuple{ntuple(3) do i\n @vars(b::SVector{3, Float32})\n end...}\nend\n\n@testset \"MPIStateArray varsindex\" begin\n # check with invalid vars size\n @test_throws ErrorException MPIStateArray{Float32, V}(\n mpicomm,\n ArrayType,\n 4,\n 1,\n 8,\n )\n\n Q = MPIStateArray{Float32, V}(mpicomm, ArrayType, 4, 34, 8)\n @test Q.a === view(MPIStateArrays.realview(Q), :, 1:1, :)\n @test Q.b === view(MPIStateArrays.realview(Q), :, 2:4, :)\n @test Q.c === view(MPIStateArrays.realview(Q), :, 5:28, :)\n @test Q.d === view(MPIStateArrays.realview(Q), :, 29:29, :)\n @test Q.e === view(MPIStateArrays.realview(Q), :, 30:34, :)\n\n @test getstateview(Q, \"a\") === view(MPIStateArrays.realview(Q), :, 1:1, :)\n @test getstateview(Q, \"b\") === view(MPIStateArrays.realview(Q), :, 2:4, :)\n @test getstateview(Q, \"c\") === view(MPIStateArrays.realview(Q), :, 5:28, :)\n @test getstateview(Q, \"d\") === view(MPIStateArrays.realview(Q), :, 29:29, :)\n @test getstateview(Q, \"e\") === view(MPIStateArrays.realview(Q), :, 30:34, :)\n @test getstateview(Q, \"e.a\") ===\n view(MPIStateArrays.realview(Q), :, 30:30, :)\n @test getstateview(Q, \"e.b\") ===\n view(MPIStateArrays.realview(Q), :, 31:33, :)\n @test getstateview(Q, \"e.d\") ===\n view(MPIStateArrays.realview(Q), :, 34:34, :)\n\n @test getstateview(Q, :(a)) === view(MPIStateArrays.realview(Q), :, 1:1, :)\n @test getstateview(Q, :(b)) === view(MPIStateArrays.realview(Q), :, 2:4, :)\n @test getstateview(Q, :(c)) === view(MPIStateArrays.realview(Q), :, 5:28, :)\n @test getstateview(Q, :(d)) ===\n view(MPIStateArrays.realview(Q), :, 29:29, :)\n @test getstateview(Q, :(e)) ===\n view(MPIStateArrays.realview(Q), :, 30:34, :)\n @test getstateview(Q, :(e.a)) ===\n view(MPIStateArrays.realview(Q), :, 30:30, :)\n @test getstateview(Q, :(e.b)) ===\n view(MPIStateArrays.realview(Q), :, 31:33, :)\n @test getstateview(Q, :(e.d)) ===\n view(MPIStateArrays.realview(Q), :, 34:34, :)\n\n @test_throws ErrorException Q.aa\n @test_throws ErrorException getstateview(Q, \"aa\")\n\n P = similar(Q)\n @test P.a === view(MPIStateArrays.realview(P), :, 1:1, :)\n @test P.b === view(MPIStateArrays.realview(P), :, 2:4, :)\n @test P.c === view(MPIStateArrays.realview(P), :, 5:28, :)\n @test P.d === view(MPIStateArrays.realview(P), :, 29:29, :)\n @test P.e === view(MPIStateArrays.realview(P), :, 30:34, :)\n\n @test getstateview(P, \"a\") === view(MPIStateArrays.realview(P), :, 1:1, :)\n @test getstateview(P, \"b\") === view(MPIStateArrays.realview(P), :, 2:4, :)\n @test getstateview(P, \"c\") === view(MPIStateArrays.realview(P), :, 5:28, :)\n @test getstateview(P, \"d\") === view(MPIStateArrays.realview(P), :, 29:29, :)\n @test getstateview(P, \"e\") === view(MPIStateArrays.realview(P), :, 30:34, :)\n @test getstateview(P, \"e.a\") ===\n view(MPIStateArrays.realview(P), :, 30:30, :)\n @test getstateview(P, \"e.b\") ===\n view(MPIStateArrays.realview(P), :, 31:33, :)\n @test getstateview(P, \"e.d\") ===\n view(MPIStateArrays.realview(P), :, 34:34, :)\n\n @test getstateview(P, :(a)) === view(MPIStateArrays.realview(P), :, 1:1, :)\n @test getstateview(P, :(b)) === view(MPIStateArrays.realview(P), :, 2:4, :)\n @test getstateview(P, :(c)) === view(MPIStateArrays.realview(P), :, 5:28, :)\n @test getstateview(P, :(d)) ===\n view(MPIStateArrays.realview(P), :, 29:29, :)\n @test getstateview(P, :(e)) ===\n view(MPIStateArrays.realview(P), :, 30:34, :)\n @test getstateview(P, :(e.a)) ===\n view(MPIStateArrays.realview(P), :, 30:30, :)\n @test getstateview(P, :(e.b)) ===\n view(MPIStateArrays.realview(P), :, 31:33, :)\n @test getstateview(P, :(e.d)) ===\n view(MPIStateArrays.realview(P), :, 34:34, :)\n\n A = MPIStateArray{Float32}(mpicomm, ArrayType, 4, 29, 8)\n @test_throws ErrorException A.a\n @test_throws ErrorException getstateview(A, \"a\")\nend\n\n@testset \"MPIStateArray show_not_finite_fields\" begin\n post_msg = \"are not finite (has NaNs or Inf)\"\n Q = MPIStateArray{Float32, V}(mpicomm, ArrayType, 4, 34, 8)\n Q .= 1\n Qv = view(MPIStateArrays.realview(Q), :, 1:1, :)\n Qv .= NaN\n msg = \"Field(s) (a) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\n\n Qv = view(MPIStateArrays.realview(Q), :, 2:2, :)\n Qv .= NaN\n msg = \"Field(s) (a, and b[1]) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\n\n Qv = view(MPIStateArrays.realview(Q), :, 31:31, :)\n Qv .= NaN\n msg = \"Field(s) (a, b[1], and e.b[1]) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\n\n Q .= 1\n @test show_not_finite_fields(Q) == nothing\nend\n\n@testset \"MPIStateArray show_not_finite_fields - ntuple vars\" begin\n post_msg = \"are not finite (has NaNs or Inf)\"\n Q = MPIStateArray{Float32, VNT}(mpicomm, ArrayType, 4, 10, 8)\n Q .= 1\n Qv = view(MPIStateArrays.realview(Q), :, 1:1, :)\n Qv .= NaN\n msg = \"Field(s) (a) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\n\n Qv = view(MPIStateArrays.realview(Q), :, 2:2, :)\n Qv .= NaN\n msg = \"Field(s) (a, and e[1].b[1]) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\n\n Qv = view(MPIStateArrays.realview(Q), :, 6:6, :)\n Qv .= NaN\n msg = \"Field(s) (a, e[1].b[1], and e[2].b[2]) \" * post_msg\n @test_logs (:warn, msg) show_not_finite_fields(Q)\nend\n" }, { "path": "test/Atmos/EDMF/Artifacts.toml", "content": "[PyCLES_output]\ngit-tree-sha1 = \"61af161b398cb0daabeb2eb1d3c57c7ba3629514\"\n" }, { "path": "test/Atmos/EDMF/bomex_edmf.jl", "content": "using ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.DGMethods\nusing ClimateMachine.SystemSolvers\nimport ClimateMachine.DGMethods: custom_filter!\nusing ClimateMachine.Mesh.Filters: apply!\nusing ClimateMachine.BalanceLaws: vars_state\nusing JLD2, FileIO\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"bomex_model.jl\"))\ninclude(joinpath(\"helper_funcs\", \"diagnostics_configuration.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n # GCM setting - Initialize the grid mean profiles of prognostic variables (ρ,e_int,q_tot,u,v,w)\n z = altitude(m, aux)\n\n # SCM setting - need to have separate cases coded and called from a folder - see what LES does\n # a moist_thermo state is used here to convert the input θ,q_tot to e_int, q_tot profile\n e_int = internal_energy(m, state, aux)\n param_set = parameter_set(m)\n if moisture_model(m) isa DryModel\n ρq_tot = FT(0)\n ts = PhaseDry(param_set, e_int, state.ρ)\n else\n ρq_tot = gm.moisture.ρq_tot\n ts = PhaseEquil_ρeq(param_set, state.ρ, e_int, ρq_tot / state.ρ)\n end\n T = air_temperature(ts)\n p = air_pressure(ts)\n q = PhasePartition(ts)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = ρq_tot * a_min\n end\n\n # initialize environment covariance with zero for now\n if z <= FT(2500)\n en.ρatke = gm.ρ * (FT(1) - z / FT(3000))\n else\n en.ρatke = FT(0)\n end\n en.ρaθ_liq_cv = FT(1e-5) / max(z, FT(10))\n en.ρaq_tot_cv = FT(1e-5) / max(z, FT(10))\n en.ρaθ_liq_q_tot_cv = FT(1e-7) / max(z, FT(10))\n return nothing\nend;\n\nstruct ZeroVerticalVelocityFilter <: AbstractCustomFilter end\nfunction custom_filter!(::ZeroVerticalVelocityFilter, bl, state, aux)\n state.ρu = SVector(state.ρu[1], state.ρu[2], 0)\nend\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # DG polynomial order\n N = 4\n nelem_vert = 20\n\n # Prescribe domain parameters\n zmax = FT(3000)\n\n t0 = FT(0)\n\n # Simulation time\n timeend = FT(400)\n CFLmax = FT(1.2)\n\n config_type = SingleStackConfigType\n\n ode_solver_type = ClimateMachine.IMEXSolverType(\n implicit_model = AtmosAcousticGravityLinearModel,\n implicit_solver = SingleColumnLU,\n solver_method = ARK2GiraldoKellyConstantinescu,\n split_explicit_implicit = true,\n discrete_splitting = false,\n )\n\n N_updrafts = 1\n N_quad = 3\n turbconv = EDMF(FT, N_updrafts, N_quad, param_set)\n\n model =\n bomex_model(FT, config_type, zmax, surface_flux; turbconv = turbconv)\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"BOMEX_EDMF\",\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = zmax,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :moisture, :ρq_tot),\n )\n\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config =\n config_diagnostics(driver_config, timeend; interval = \"50ssecs\")\n\n cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n (\"moisture.ρq_tot\", turbconv_filters(turbconv)...),\n solver_config.dg.grid,\n TMARFilter(),\n )\n Filters.apply!(\n ZeroVerticalVelocityFilter(),\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n solver_config.Q,\n solver_config.dg.state_auxiliary,\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 10\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n # Save to disc (for debugging):\n # @save \"bomex_edmf_nstep=$nstep.jld2\" diag_vs_z\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n check_cons = (\n ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.001)),\n ClimateMachine.ConservationCheck(\"energy.ρe\", \"3000steps\", FT(0.0025)),\n )\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cbtmarfilter, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n\n# add a command line argument to specify the kind of surface flux\n# TODO: this will move to the future namelist functionality\nbomex_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(bomex_args, \"BOMEX\")\n@add_arg_table! bomex_args begin\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk\"\n arg_type = String\n default = \"prescribed\"\nend\n\ncl_args = ClimateMachine.init(\n parse_clargs = true,\n custom_clargs = bomex_args,\n output_dir = get(ENV, \"CLIMATEMACHINE_SETTINGS_OUTPUT_DIR\", \"output\"),\n fix_rng_seed = true,\n)\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\ninclude(joinpath(@__DIR__, \"report_mse_bomex.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/closures/entr_detr.jl", "content": "#### Entrainment-Detrainment kernels\n\nfunction entr_detr(\n bl::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n ts_up,\n ts_en,\n env,\n buoy,\n) where {FT}\n turbconv = turbconv_model(bl)\n EΔ_up = vuntuple(n_updrafts(turbconv)) do i\n entr_detr(bl, turbconv.entr_detr, state, aux, ts_up, ts_en, env, buoy, i)\n end\n E_dyn, Δ_dyn, E_trb = ntuple(i -> map(x -> x[i], EΔ_up), 3)\n return E_dyn, Δ_dyn, E_trb\nend\n\n\"\"\"\n entr_detr(\n m::AtmosModel{FT},\n entr::EntrainmentDetrainment,\n state::Vars,\n aux::Vars,\n ts_up,\n ts_en,\n env,\n buoy,\n i,\n ) where {FT}\n\nReturns the dynamic entrainment and detrainment rates,\nas well as the turbulent entrainment rate, following\nCohen et al. (JAMES, 2020), given:\n - `m`, an `AtmosModel`\n - `entr`, an `EntrainmentDetrainment` model\n - `state`, state variables\n - `aux`, auxiliary variables\n - `ts_up`, updraft thermodynamic states\n - `ts_en`, environment thermodynamic states\n - `env`, NamedTuple of environment variables\n - `buoy`, NamedTuple of environment and updraft buoyancies\n - `i`, index of the updraft\n\"\"\"\nfunction entr_detr(\n m::AtmosModel{FT},\n entr::EntrainmentDetrainment,\n state::Vars,\n aux::Vars,\n ts_up,\n ts_en,\n env,\n buoy,\n i,\n) where {FT}\n\n # Alias convention:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n en_aux = aux.turbconv.environment\n up_aux = aux.turbconv.updraft\n turbconv = turbconv_model(m)\n N_up = n_updrafts(turbconv)\n ρ_inv = 1 / gm.ρ\n a_up_i = up[i].ρa * ρ_inv\n lim_E = entr.lim_ϵ\n lim_amp = entr.lim_amp\n w_min = entr.w_min\n # precompute vars\n w_up_i = fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)\n sqrt_tke = sqrt(max(en.ρatke, 0) * ρ_inv / env.a)\n # ensure far from zero\n Δw = filter_w(w_up_i - env.w, w_min)\n w_up_i = filter_w(w_up_i, w_min)\n Δb = buoy.up[i] - buoy.en\n D_E, D_δ, M_δ, M_E = nondimensional_exchange_functions(\n m,\n entr,\n state,\n aux,\n ts_up,\n ts_en,\n env,\n buoy,\n i,\n )\n\n # I am commenting this out for now, to make sure there is no slowdown here\n Λ_w = abs(Δb / Δw)\n Λ_tke = entr.c_λ * abs(Δb / (max(en.ρatke * ρ_inv, 0) + w_min))\n λ = lamb_smooth_minimum(\n SVector(Λ_w, Λ_tke),\n turbconv.mix_len.smin_ub,\n turbconv.mix_len.smin_rm,\n )\n\n # compute entrainment/detrainment components\n # TO DO: Add updraft height dependency (non-local)\n E_trb = 2 * up[i].ρa * entr.c_t * sqrt_tke / turbconv.pressure.H_up_min\n E_dyn = up[i].ρa * λ * (D_E + M_E)\n Δ_dyn = up[i].ρa * λ * (D_δ + M_δ)\n\n E_dyn = max(E_dyn, FT(0))\n Δ_dyn = max(Δ_dyn, FT(0))\n E_trb = max(E_trb, FT(0))\n return E_dyn, Δ_dyn, E_trb\nend;\n" }, { "path": "test/Atmos/EDMF/closures/mixing_length.jl", "content": "#### Mixing length model kernels\n\"\"\"\n mixing_length(\n m::AtmosModel{FT},\n ml::MixingLengthModel,\n args,\n Δ::Tuple,\n Et::Tuple,\n ts_gm,\n ts_en,\n env,\n ) where {FT}\n\nReturns the mixing length used in the diffusive turbulence closure, given:\n - `m`, an `AtmosModel`\n - `ml`, a `MixingLengthModel`\n - `args`, the top-level arguments\n - `Δ`, the detrainment rate\n - `Et`, the turbulent entrainment rate\n - `ts_gm`, grid-mean thermodynamic states\n - `ts_en`, environment thermodynamic states\n - `env`, NamedTuple of environment variables\n\"\"\"\nfunction mixing_length(\n m::AtmosModel{FT},\n ml::MixingLengthModel,\n args,\n Δ::Tuple,\n Et::Tuple,\n Shear²,\n ts_gm,\n ts_en,\n env,\n) where {FT}\n @unpack state, aux, diffusive, t = args\n # TODO: use functions: obukhov_length, ustar, ϕ_m\n turbconv = turbconv_model(m)\n # Alias convention:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n gm_aux = aux\n N_up = n_updrafts(turbconv)\n\n z = altitude(m, aux)\n param_set = parameter_set(m)\n _grav::FT = grav(param_set)\n ρinv = 1 / gm.ρ\n\n tke_en = max(en.ρatke, 0) * ρinv / env.a\n\n # buoyancy related functions\n # compute obukhov_length and ustar from SurfaceFlux.jl here\n ustar = turbconv.surface.ustar\n obukhov_length = turbconv.surface.obukhov_length\n\n ∂b∂z, Nˢ_eff = compute_buoyancy_gradients(m, args, ts_gm, ts_en)\n Grad_Ri = ∇Richardson_number(∂b∂z, Shear², 1 / ml.max_length, ml.Ri_c)\n Pr_t = turbulent_Prandtl_number(ml.Pr_n, Grad_Ri, ml.ω_pr)\n\n # compute L1\n Nˢ_fact = (sign(Nˢ_eff - eps(FT)) + 1) / 2\n coeff = min(ml.c_b * sqrt(tke_en) / Nˢ_eff, ml.max_length)\n L_Nˢ = coeff * Nˢ_fact + ml.max_length * (FT(1) - Nˢ_fact)\n\n # compute L2 - law of the wall\n # TODO: use zLL from altitude\n surf_vals = subdomain_surface_values(m, gm, gm_aux, turbconv.surface.zLL)\n\n L_W = ml.κ * max(z, 5) / (sqrt(turbconv.surface.κ_star²) * ml.c_m)\n if obukhov_length < -eps(FT)\n L_W *= min((FT(1) - ml.a2 * z / obukhov_length)^ml.a1, 1 / ml.κ)\n end\n\n # compute L3 - entrainment detrainment sources\n # Production/destruction terms\n a = ml.c_m * (Shear² - ∂b∂z / Pr_t) * sqrt(tke_en)\n # Dissipation term\n b = FT(0)\n a_up = vuntuple(i -> up[i].ρa * ρinv, N_up)\n w_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)\n end\n b = sum(\n ntuple(N_up) do i\n Δ[i] / gm.ρ / env.a *\n ((w_up[i] - env.w) * (w_up[i] - env.w) / 2 - tke_en) -\n (w_up[i] - env.w) * Et[i] / gm.ρ * env.w / env.a\n end,\n )\n\n c_neg = ml.c_d * tke_en * sqrt(tke_en)\n if abs(a) > ml.random_minval && 4 * a * c_neg > -b^2\n l_entdet =\n max(-b / FT(2) / a + sqrt(b^2 + 4 * a * c_neg) / 2 / a, FT(0))\n elseif abs(a) < eps(FT) && abs(b) > eps(FT)\n l_entdet = c_neg / b\n else\n l_entdet = FT(0)\n end\n L_tke = l_entdet\n\n if L_Nˢ < eps(FT) || L_Nˢ > ml.max_length\n L_Nˢ = ml.max_length\n end\n if L_W < eps(FT) || L_W > ml.max_length\n L_W = ml.max_length\n end\n if L_tke < eps(FT) || L_tke > ml.max_length\n L_tke = ml.max_length\n end\n\n l_mix =\n lamb_smooth_minimum(SVector(L_Nˢ, L_W, L_tke), ml.smin_ub, ml.smin_rm)\n return l_mix, ∂b∂z, Pr_t\nend;\n" }, { "path": "test/Atmos/EDMF/closures/pressure.jl", "content": "#### Pressure model kernels\n\nfunction perturbation_pressure(bl::AtmosModel{FT}, args, env, buoy) where {FT}\n dpdz = vuntuple(n_updrafts(turbconv_model(bl))) do i\n perturbation_pressure(bl, turbconv_model(bl).pressure, args, env, buoy, i)\n end\n return dpdz\nend\n\n\"\"\"\n perturbation_pressure(\n m::AtmosModel{FT},\n press::PressureModel,\n args,\n env,\n buoy,\n i,\n ) where {FT}\n\nReturns the value of perturbation pressure gradient\nfor updraft i following He et al. (JAMES, 2020), given:\n\n - `m`, an `AtmosModel`\n - `press`, a `PressureModel`\n - `args`, top-level arguments\n - `env`, NamedTuple of environment variables\n - `buoy`, NamedTuple of environment and updraft buoyancies\n - `i`, index of the updraft\n\"\"\"\nfunction perturbation_pressure(\n m::AtmosModel{FT},\n press::PressureModel,\n args,\n env,\n buoy,\n i,\n) where {FT}\n @unpack state, diffusive, aux = args\n # Alias convention:\n up = state.turbconv.updraft\n up_aux = aux.turbconv.updraft\n up_dif = diffusive.turbconv.updraft\n\n w_up_i = fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)\n\n nh_press_buoy = press.α_b * buoy.up[i]\n nh_pressure_adv = -press.α_a * w_up_i * up_dif[i].∇w[3]\n # TO DO: Add updraft height dependency (non-local)\n nh_pressure_drag =\n press.α_d * (w_up_i - env.w) * abs(w_up_i - env.w) / press.H_up_min\n\n dpdz = nh_press_buoy + nh_pressure_adv + nh_pressure_drag\n\n return dpdz\nend;\n" }, { "path": "test/Atmos/EDMF/closures/surface_functions.jl", "content": "#### Surface model kernels\n\nusing Statistics\n\n\"\"\"\n subdomain_surface_values(\n atmos::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n zLL::FT,\n ) where {FT}\n\nReturns the surface values of updraft area fraction, updraft\nliquid water potential temperature (`θ_liq`), updraft total\nwater specific humidity (`q_tot`), environmental variances of\n`θ_liq` and `q_tot`, environmental covariance of `θ_liq` with\n`q_tot`, and environmental TKE, given:\n\n - `atmos`, an `AtmosModel`\n - `state`, state variables\n - `aux`, auxiliary variables\n - `zLL`, height of the lowest nodal level\n\"\"\"\nfunction subdomain_surface_values(\n atmos::AtmosModel,\n state::Vars,\n aux::Vars,\n zLL,\n)\n turbconv = turbconv_model(atmos)\n subdomain_surface_values(turbconv.surface, turbconv, atmos, state, aux, zLL)\nend\n\nfunction subdomain_surface_values(\n surf::SurfaceModel,\n turbconv::EDMF{FT},\n atmos::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n zLL::FT,\n) where {FT}\n\n N_up = n_updrafts(turbconv)\n gm = state\n # TODO: change to new_thermo_state\n ts = recover_thermo_state(atmos, state, aux)\n q = PhasePartition(ts)\n _cp_m = cp_m(ts)\n lv = latent_heat_vapor(ts)\n Π = exner(ts)\n ρ_inv = 1 / gm.ρ\n upd_surface_std = turbconv.surface.upd_surface_std\n\n θ_liq_surface_flux = surf.shf / Π / _cp_m\n q_tot_surface_flux = surf.lhf / lv\n # these value should be given from the SurfaceFluxes.jl once it is merged\n oblength = turbconv.surface.obukhov_length\n ustar = turbconv.surface.ustar\n\n unstable = oblength < -eps(FT)\n fact = unstable ? (1 - surf.ψϕ_stab * zLL / oblength)^(-FT(2 // 3)) : 1\n tke_fact = unstable ? cbrt(zLL / oblength * zLL / oblength) : 0\n ustar² = ustar^2\n θ_liq_cv = 4 * (θ_liq_surface_flux * θ_liq_surface_flux) / (ustar²) * fact\n q_tot_cv = 4 * (q_tot_surface_flux * q_tot_surface_flux) / (ustar²) * fact\n θ_liq_q_tot_cv =\n 4 * (θ_liq_surface_flux * q_tot_surface_flux) / (ustar²) * fact\n tke = ustar² * (surf.κ_star² + tke_fact)\n\n a_up_surf = ntuple(i -> FT(surf.a / N_up), N_up)\n e_int = internal_energy(atmos, state, aux)\n ts_new = new_thermo_state(atmos, state, aux)\n θ_liq = liquid_ice_pottemp(ts_new)\n\n θ_liq_up_surf = ntuple(N_up) do i\n θ_liq + upd_surface_std[i] * sqrt(max(θ_liq_cv, 0))\n end\n\n ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot\n q_tot_up_surf = ntuple(N_up) do i\n ρq_tot * ρ_inv + upd_surface_std[i] * sqrt(max(q_tot_cv, 0))\n end\n\n return (;\n a_up_surf,\n θ_liq_up_surf,\n q_tot_up_surf,\n θ_liq_cv,\n q_tot_cv,\n θ_liq_q_tot_cv,\n tke,\n )\nend;\n\nfunction subdomain_surface_values(\n surf::NeutralDrySurfaceModel,\n turbconv::EDMF{FT},\n atmos::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n zLL::FT,\n) where {FT}\n\n N_up = n_updrafts(turbconv)\n θ_liq_cv = FT(0)\n q_tot_cv = FT(0)\n θ_liq_q_tot_cv = FT(0)\n tke = surf.κ_star² * surf.ustar * surf.ustar\n\n ts_new = new_thermo_state(atmos, state, aux)\n a_up_surf = ntuple(i -> FT(surf.a / N_up), N_up)\n q_tot_up_surf = ntuple(i -> FT(0), N_up)\n θ_liq_up_surf = ntuple(i -> liquid_ice_pottemp(ts_new), N_up)\n\n return (;\n a_up_surf,\n θ_liq_up_surf,\n q_tot_up_surf,\n θ_liq_cv,\n q_tot_cv,\n θ_liq_q_tot_cv,\n tke,\n )\nend;\n\n\"\"\"\n percentile_bounds_mean_norm(\n low_percentile::FT,\n high_percentile::FT,\n n_samples::Int,\n ) where {FT <: AbstractFloat}\n\nReturns the mean of all instances of a standard Gaussian random\nvariable that have a CDF higher than low_percentile and lower\nthan high_percentile, given a total of n_samples of the standard\nGaussian, given:\n - `low_percentile`, lower limit of the CDF\n - `high_percentile`, higher limit of the CDF\n - `n_samples`, the total number of samples drawn from the Gaussian\n\"\"\"\nfunction percentile_bounds_mean_norm(\n low_percentile::FT,\n high_percentile::FT,\n n_samples::Int,\n) where {FT <: AbstractFloat}\n x = rand(Normal(), n_samples)\n xp_low = quantile(Normal(), low_percentile)\n xp_high = quantile(Normal(), high_percentile)\n filter!(y -> xp_low < y < xp_high, x)\n return Statistics.mean(x)\nend\n" }, { "path": "test/Atmos/EDMF/closures/turbulence_functions.jl", "content": "#### Turbulence model kernels\n\n\"\"\"\n thermo_variables(ts::ThermodynamicState)\n\nA NamedTuple of thermodynamic variables,\ncomputed from the given thermodynamic state.\n\"\"\"\nfunction thermo_variables(ts::ThermodynamicState)\n return (\n θ_dry = dry_pottemp(ts),\n θ_liq = liquid_ice_pottemp(ts),\n q_tot = total_specific_humidity(ts),\n T = air_temperature(ts),\n R_m = gas_constant_air(ts),\n q_vap = vapor_specific_humidity(ts),\n q_liq = liquid_specific_humidity(ts),\n q_ice = ice_specific_humidity(ts),\n )\nend\n\n\"\"\"\n compute_buoyancy_gradients(\n m::AtmosModel{FT},\n args,\n ts_gm,\n ts_en\n ) where {FT}\nReturns the environmental buoyancy gradient following Tan et al. (JAMES, 2018)\nand the effective environmental static stability following\nLopez-Gomez et al. (JAMES, 2020), given:\n - `m`, an `AtmosModel`\n - `args`, top-level arguments\n - `ts_gm`, grid-mean thermodynamic state\n - `ts_en`, environment thermodynamic state\n\"\"\"\nfunction compute_buoyancy_gradients(\n m::AtmosModel{FT},\n args,\n ts_gm,\n ts_en,\n) where {FT}\n @unpack state, aux, diffusive = args\n # Alias convention:\n gm = state\n en_dif = diffusive.turbconv.environment\n N_up = n_updrafts(turbconv_model(m))\n param_set = parameter_set(m)\n\n _grav::FT = grav(param_set)\n _R_d::FT = R_d(param_set)\n _R_v::FT = R_v(param_set)\n ε_v::FT = 1 / molmass_ratio(param_set)\n p = air_pressure(ts_gm)\n\n q_tot_en = total_specific_humidity(ts_en)\n θ_liq_en = liquid_ice_pottemp(ts_en)\n lv = latent_heat_vapor(ts_en)\n T = air_temperature(ts_en)\n Π = exner(ts_en)\n q_liq = liquid_specific_humidity(ts_en)\n _cp_m = cp_m(ts_en)\n θ_virt = virtual_pottemp(ts_en)\n\n (ts_dry, ts_cloudy, cloud_frac) =\n compute_subdomain_statistics(m, args, ts_gm, ts_en)\n cloudy = thermo_variables(ts_cloudy)\n dry = thermo_variables(ts_dry)\n\n prefactor = _grav * (_R_d * gm.ρ / p * Π)\n\n ∂b∂θl_dry = prefactor * (1 + (ε_v - 1) * dry.q_tot)\n ∂b∂qt_dry = prefactor * dry.θ_liq * (ε_v - 1)\n\n if cloud_frac > FT(0)\n num =\n prefactor *\n (1 + ε_v * (1 + lv / _R_v / cloudy.T) * cloudy.q_vap - cloudy.q_tot)\n den = 1 + lv * lv / _cp_m / _R_v / cloudy.T / cloudy.T * cloudy.q_vap\n ∂b∂θl_cloudy = num / den\n ∂b∂qt_cloudy =\n (lv / _cp_m / cloudy.T * ∂b∂θl_cloudy - prefactor) * cloudy.θ_dry\n else\n ∂b∂θl_cloudy = FT(0)\n ∂b∂qt_cloudy = FT(0)\n end\n\n ∂b∂θl = (cloud_frac * ∂b∂θl_cloudy + (1 - cloud_frac) * ∂b∂θl_dry)\n ∂b∂qt = (cloud_frac * ∂b∂qt_cloudy + (1 - cloud_frac) * ∂b∂qt_dry)\n\n # Partial buoyancy gradients\n ∂b∂z_θl = en_dif.∇θ_liq[3] * ∂b∂θl\n ∂b∂z_qt = en_dif.∇q_tot[3] * ∂b∂qt\n ∂b∂z = ∂b∂z_θl + ∂b∂z_qt\n\n # Computation of buoyancy frequency based on θ_lv\n ∂θvl∂θ_liq = 1 + (ε_v - 1) * q_tot_en\n ∂θvl∂qt = (ε_v - 1) * θ_liq_en\n # apply chain-rule\n ∂θvl∂z = ∂θvl∂θ_liq * en_dif.∇θ_liq[3] + ∂θvl∂qt * en_dif.∇q_tot[3]\n\n ∂θv∂θvl = exp(lv * q_liq / _cp_m / T)\n λ_stb = cloud_frac\n\n Nˢ_eff =\n _grav / θ_virt *\n ((1 - λ_stb) * en_dif.∇θv[3] + λ_stb * ∂θvl∂z * ∂θv∂θvl)\n return ∂b∂z, Nˢ_eff\nend;\n\n\"\"\"\n ∇Richardson_number(\n ∂b∂z::FT,\n Shear²::FT,\n minval::FT,\n Ri_c::FT,\n ) where {FT}\n\nReturns the gradient Richardson number, given:\n - `∂b∂z`, the vertical buoyancy gradient\n - `Shear²`, the squared vertical gradient of horizontal velocity\n - `maxval`, maximum value of the output, typically the critical Ri number\n\"\"\"\nfunction ∇Richardson_number(\n ∂b∂z::FT,\n Shear²::FT,\n minval::FT,\n Ri_c::FT,\n) where {FT}\n return min(∂b∂z / max(Shear², minval), Ri_c)\nend;\n\n\"\"\"\n turbulent_Prandtl_number(\n Pr_n::FT,\n Grad_Ri::FT,\n ω_pr::FT\n ) where {FT}\n\nReturns the turbulent Prandtl number, given:\n - `Pr_n`, the turbulent Prandtl number under neutral conditions\n - `Grad_Ri`, the gradient Richardson number\n\"\"\"\nfunction turbulent_Prandtl_number(Pr_n::FT, Grad_Ri::FT, ω_pr::FT) where {FT}\n if Grad_Ri > FT(0)\n factor =\n 2 * Grad_Ri /\n (1 + ω_pr * Grad_Ri - sqrt((1 + ω_pr * Grad_Ri)^2 - 4 * Grad_Ri))\n else\n factor = FT(1)\n end\n return Pr_n * factor\nend;\n" }, { "path": "test/Atmos/EDMF/compute_mse.jl", "content": "using ClimateMachine\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n using Plots\nend\n\nusing OrderedCollections\nusing Test\nusing NCDatasets\nusing Dierckx\nusing PrettyTables\nusing Printf\nusing ArtifactWrappers\n\n# Get PyCLES_output dataset folder:\n#! format: off\nPyCLES_output_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"PyCLES_output\",\n ArtifactFile[\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/johlutwhohvr66wn38cdo7a6rluvz708.nc\", filename = \"Rico.nc\",),\n ArtifactFile(url = \"https://caltech.box.com/shared/static/zraeiftuzlgmykzhppqwrym2upqsiwyb.nc\", filename = \"Gabls.nc\",),\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/toyvhbwmow3nz5bfa145m5fmcb2qbfuz.nc\", filename = \"DYCOMS_RF01.nc\",),\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/ivo4751camlph6u3k68ftmb1dl4z7uox.nc\", filename = \"TRMM_LBA.nc\",),\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/4osqp0jpt4cny8fq2ukimgfnyi787vsy.nc\", filename = \"ARM_SGP.nc\",),\n ArtifactFile(url = \"https://caltech.box.com/shared/static/jci8l11qetlioab4cxf5myr1r492prk6.nc\", filename = \"Bomex.nc\",),\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/pzuu6ii99by2s356ij69v5cb615200jq.nc\", filename = \"Soares.nc\",),\n # ArtifactFile(url = \"https://caltech.box.com/shared/static/7upt639siyc2umon8gs6qsjiqavof5cq.nc\", filename = \"Nieuwstadt.nc\",),\n ],\n)\nPyCLES_output_dataset_path = get_data_folder(PyCLES_output_dataset)\n#! format: on\n\ninclude(\"variable_map.jl\")\n\nfunction compute_mse(\n grid,\n bl,\n time_cm,\n dons_arr,\n ds,\n experiment,\n best_mse,\n t_compare,\n plot_dir = nothing,\n)\n mse = Dict()\n\n # Ensure domain matches:\n z_les = ds[\"z_half\"][:]\n z_cm = get_z(grid; rm_dupes = true)\n @info \"Z extent for LES vs CLIMA:\"\n @show extrema(z_cm)\n @show extrema(z_les)\n\n time_les = ds[\"t\"][:]\n\n # Find the nearest matching final time:\n t_cmp = min(time_cm[end], time_les[end])\n\n # Accidentally running a short simulation\n # could improve MSE. So, let's test that\n # we run for at least t_compare. We should\n # increase this as we can reach higher CFL.\n @test t_cmp >= t_compare\n\n # Ensure z_cm and dons_arr fields are consistent lengths:\n @test length(z_cm) == length(dons_arr[1][first(keys(dons_arr[1]))])\n\n data_cm = Dict()\n dons_cont = Dict()\n cm_variables = []\n computed_mse = []\n table_best_mse = []\n mse_reductions = []\n pycles_variables = []\n data_scales = []\n pycles_weight = []\n for (ftc) in keys(best_mse)\n # Only compare fields defined for var_map\n tup = var_map(ftc)\n tup == nothing && continue\n\n # Unpack the data\n LES_var = tup[1]\n facts = tup[2]\n push!(cm_variables, ftc)\n push!(pycles_variables, LES_var)\n data_ds = ds.group[\"profiles\"]\n data_les = data_ds[LES_var][:]\n\n # Scale the data for comparison\n ρ = data_ds[\"rho\"][:]\n a_up = data_ds[\"updraft_fraction\"][:]\n a_en = 1 .- data_ds[\"updraft_fraction\"][:]\n ρa_up = ρ .* a_up\n ρa_en = ρ .* a_en\n ρa = occursin(\"updraft\", ftc) in ftc ? ρa_up : ρa_en\n if :a in facts && :ρ in facts\n data_les .*= ρa\n push!(pycles_weight, \"ρa\")\n elseif :ρ in facts\n data_les .*= ρ\n push!(pycles_weight, \"ρ\")\n else\n push!(pycles_weight, \"1\")\n end\n\n # Interpolate data\n steady_data = length(size(data_les)) == 1\n if steady_data\n data_les_cont = Spline1D(z_les, data_les)\n else # unsteady data\n data_les_cont = Spline2D(time_les, z_les, data_les')\n end\n data_cm_arr_ = [\n dons_arr[i][ftc][i_z]\n for i in 1:length(time_cm), i_z in 1:length(z_cm)\n ]\n data_cm_arr = reshape(data_cm_arr_, (length(time_cm), length(z_cm)))\n data_cm[ftc] = Spline2D(time_cm, z_cm, data_cm_arr)\n\n # Compute data scale\n data_scale = sum(abs.(data_les)) / length(data_les)\n push!(data_scales, data_scale)\n\n # Plot comparison\n if plot_dir ≠ nothing\n p = plot()\n if steady_data\n data_les_cont_mapped = map(z -> data_les_cont(z), z_cm)\n else\n data_les_cont_mapped = map(z -> data_les_cont(t_cmp, z), z_cm)\n end\n plot!(\n data_les_cont_mapped,\n z_cm ./ 10^3,\n xlabel = ftc,\n ylabel = \"z [km]\",\n label = \"PyCLES\",\n )\n plot!(\n map(z -> data_cm[ftc](t_cmp, z), z_cm),\n z_cm ./ 10^3,\n xlabel = ftc,\n ylabel = \"z [km]\",\n label = \"CM\",\n )\n mkpath(plot_dir)\n ftc_name = replace(ftc, \".\" => \"_\")\n savefig(joinpath(plot_dir, \"$ftc_name.png\"))\n end\n\n # Compute mean squared error (mse)\n if steady_data\n mse_single_var = sum(map(z_cm) do z\n (data_les_cont(z) - data_cm[ftc](t_cmp, z))^2\n end)\n else\n mse_single_var = sum(map(z_cm) do z\n (data_les_cont(t_cmp, z) - data_cm[ftc](t_cmp, z))^2\n end)\n end\n # Normalize by data scale\n mse[ftc] = mse_single_var / data_scale^2\n\n push!(mse_reductions, (best_mse[ftc] - mse[ftc]) / best_mse[ftc] * 100)\n push!(computed_mse, mse[ftc])\n push!(table_best_mse, best_mse[ftc])\n end\n\n # Tabulate output\n header = [\n \"Variable\" \"Variable\" \"Weight\" \"Data scale\" \"MSE\" \"MSE\" \"MSE\"\n \"ClimateMachine (EDMF)\" \"PyCLES\" \"PyCLES\" \"\" \"Computed\" \"Best\" \"Reduction (%)\"\n ]\n table_data = hcat(\n cm_variables,\n pycles_variables,\n pycles_weight,\n data_scales,\n computed_mse,\n table_best_mse,\n mse_reductions,\n )\n\n @info @sprintf(\n \"Experiment comparison: %s at time t=%s\\n\",\n experiment,\n t_cmp\n )\n hl_worsened_mse = Highlighter(\n (data, i, j) -> !sufficient_mse(data[i, 5], data[i, 6]) && j == 5,\n crayon\"red bold\",\n )\n hl_worsened_mse_reduction = Highlighter(\n (data, i, j) -> !sufficient_mse(data[i, 5], data[i, 6]) && j == 7,\n crayon\"red bold\",\n )\n hl_improved_mse = Highlighter(\n (data, i, j) -> sufficient_mse(data[i, 5], data[i, 6]) && j == 7,\n crayon\"green bold\",\n )\n pretty_table(\n table_data,\n header,\n formatters = ft_printf(\"%.16e\", 5:6),\n header_crayon = crayon\"yellow bold\",\n subheader_crayon = crayon\"green bold\",\n highlighters = (\n hl_worsened_mse,\n hl_improved_mse,\n hl_worsened_mse_reduction,\n ),\n crop = :none,\n )\n\n return mse\nend\n\nsufficient_mse(computed_mse, best_mse) = computed_mse <= best_mse + sqrt(eps())\n\nfunction test_mse(computed_mse, best_mse, key)\n mse_not_regressed = sufficient_mse(computed_mse[key], best_mse[key])\n @test mse_not_regressed\n mse_not_regressed || @show key\nend\n\nfunction dons(diag_vs_z)\n return Dict(map(keys(first(diag_vs_z))) do k\n string(k) => [getproperty(ca, k) for ca in diag_vs_z]\n end)\nend\n\nget_dons_arr(diag_arr) = [dons(diag_vs_z) for diag_vs_z in diag_arr]\n\ndons_arr = get_dons_arr(diag_arr)\n" }, { "path": "test/Atmos/EDMF/edmf_kernels.jl", "content": "#### EDMF model kernels\n\nusing CLIMAParameters.Planet: e_int_v0, grav, day, R_d, R_v, molmass_ratio\nusing Printf\nusing ClimateMachine.Atmos: nodal_update_auxiliary_state!, Advect\n\nusing ClimateMachine.BalanceLaws\n\nusing ClimateMachine.MPIStateArrays: MPIStateArray\nusing ClimateMachine.DGMethods: LocalGeometry, DGModel\n\nimport ClimateMachine.Mesh.Filters: vars_state_filtered\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n prognostic_vars,\n prognostic_to_primitive!,\n primitive_to_prognostic!,\n get_prog_state,\n get_specific_state,\n flux,\n precompute,\n source,\n eq_tends,\n update_auxiliary_state!,\n init_state_prognostic!,\n compute_gradient_argument!,\n compute_gradient_flux!\n\nimport ClimateMachine.TurbulenceConvection:\n init_aux_turbconv!,\n turbconv_nodal_update_auxiliary_state!,\n turbconv_boundary_state!,\n turbconv_normal_boundary_flux_second_order!\n\nusing Thermodynamics: air_pressure, air_density\n\n\ninclude(joinpath(\"helper_funcs\", \"nondimensional_exchange_functions.jl\"))\ninclude(joinpath(\"helper_funcs\", \"lamb_smooth_minimum.jl\"))\ninclude(joinpath(\"helper_funcs\", \"utility_funcs.jl\"))\ninclude(joinpath(\"helper_funcs\", \"subdomain_statistics.jl\"))\ninclude(joinpath(\"helper_funcs\", \"diagnose_environment.jl\"))\ninclude(joinpath(\"helper_funcs\", \"subdomain_thermo_states.jl\"))\ninclude(joinpath(\"helper_funcs\", \"save_subdomain_temperature.jl\"))\ninclude(joinpath(\"closures\", \"entr_detr.jl\"))\ninclude(joinpath(\"closures\", \"pressure.jl\"))\ninclude(joinpath(\"closures\", \"mixing_length.jl\"))\ninclude(joinpath(\"closures\", \"turbulence_functions.jl\"))\ninclude(joinpath(\"closures\", \"surface_functions.jl\"))\n\n\nfunction vars_state(m::NTuple{N, Updraft}, st::Auxiliary, FT) where {N}\n return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}\nend\n\nvars_state(::Updraft, ::Auxiliary, FT) = @vars(T::FT)\nvars_state(::Environment, ::Auxiliary, FT) = @vars(T::FT)\n\nfunction vars_state(m::EDMF, st::Auxiliary, FT)\n @vars(\n environment::vars_state(m.environment, st, FT),\n updraft::vars_state(m.updraft, st, FT)\n )\nend\n\nfunction vars_state(::Updraft, ::Prognostic, FT)\n @vars(ρa::FT, ρaw::FT, ρaθ_liq::FT, ρaq_tot::FT,)\nend\n\nfunction vars_state(::Environment, ::Prognostic, FT)\n @vars(ρatke::FT, ρaθ_liq_cv::FT, ρaq_tot_cv::FT, ρaθ_liq_q_tot_cv::FT,)\nend\n\nfunction vars_state(::Updraft, ::Primitive, FT)\n @vars(a::FT, aw::FT, aθ_liq::FT, aq_tot::FT,)\nend\n\nfunction vars_state(::Environment, ::Primitive, FT)\n @vars(atke::FT, aθ_liq_cv::FT, aq_tot_cv::FT, aθ_liq_q_tot_cv::FT,)\nend\n\nfunction vars_state(\n m::NTuple{N, Updraft},\n st::Union{Prognostic, Primitive},\n FT,\n) where {N}\n return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}\nend\n\nfunction vars_state(m::EDMF, st::Union{Prognostic, Primitive}, FT)\n @vars(\n environment::vars_state(m.environment, st, FT),\n updraft::vars_state(m.updraft, st, FT)\n )\nend\n\nfunction vars_state(::Updraft, ::Gradient, FT)\n @vars(w::FT,)\nend\n\nfunction vars_state(::Environment, ::Gradient, FT)\n @vars(\n θ_liq::FT,\n q_tot::FT,\n w::FT,\n tke::FT,\n θ_liq_cv::FT,\n q_tot_cv::FT,\n θ_liq_q_tot_cv::FT,\n θv::FT,\n h_tot::FT,\n )\nend\n\nfunction vars_state(m::NTuple{N, Updraft}, st::Gradient, FT) where {N}\n return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}\nend\n\nfunction vars_state(m::EDMF, st::Gradient, FT)\n @vars(\n environment::vars_state(m.environment, st, FT),\n updraft::vars_state(m.updraft, st, FT),\n u::FT,\n v::FT\n )\nend\n\nfunction vars_state(m::NTuple{N, Updraft}, st::GradientFlux, FT) where {N}\n return Tuple{ntuple(i -> vars_state(m[i], st, FT), N)...}\nend\n\nfunction vars_state(::Updraft, st::GradientFlux, FT)\n @vars(∇w::SVector{3, FT},)\nend\n\nfunction vars_state(::Environment, ::GradientFlux, FT)\n @vars(\n ∇θ_liq::SVector{3, FT},\n ∇q_tot::SVector{3, FT},\n ∇w::SVector{3, FT},\n ∇tke::SVector{3, FT},\n ∇θ_liq_cv::SVector{3, FT},\n ∇q_tot_cv::SVector{3, FT},\n ∇θ_liq_q_tot_cv::SVector{3, FT},\n ∇θv::SVector{3, FT},\n ∇h_tot::SVector{3, FT},\n )\n\nend\n\nfunction vars_state(m::EDMF, st::GradientFlux, FT)\n @vars(\n environment::vars_state(m.environment, st, FT),\n updraft::vars_state(m.updraft, st, FT),\n ∇u::SVector{3, FT},\n ∇v::SVector{3, FT}\n )\nend\n\nfunction vars_state_filtered(::Updraft, FT)\n @vars(a::FT, aw::FT, aθ_liq::FT, aq_tot::FT,)\nend\n\nfunction vars_state_filtered(m::NTuple{N, Updraft}, FT) where {N}\n return Tuple{ntuple(i -> vars_state_filtered(m[i], FT), N)...}\nend\n\nfunction vars_state_filtered(::Environment, FT)\n @vars(atke::FT, aθ_liq_cv::FT, aq_tot_cv::FT, aθ_liq_q_tot_cv::FT,)\nend\n\nfunction vars_state_filtered(m::EDMF, FT)\n @vars(\n environment::vars_state_filtered(m.environment, FT),\n updraft::vars_state_filtered(m.updraft, FT)\n )\nend\n\nabstract type EDMFPrognosticVariable <: AbstractPrognosticVariable end\n\nabstract type EnvironmentPrognosticVariable <: EDMFPrognosticVariable end\nstruct en_ρatke <: EnvironmentPrognosticVariable end\nstruct en_ρaθ_liq_cv <: EnvironmentPrognosticVariable end\nstruct en_ρaq_tot_cv <: EnvironmentPrognosticVariable end\nstruct en_ρaθ_liq_q_tot_cv <: EnvironmentPrognosticVariable end\n\nabstract type UpdraftPrognosticVariable{i} <: EDMFPrognosticVariable end\nstruct up_ρa{i} <: UpdraftPrognosticVariable{i} end\nstruct up_ρaw{i} <: UpdraftPrognosticVariable{i} end\nstruct up_ρaθ_liq{i} <: UpdraftPrognosticVariable{i} end\nstruct up_ρaq_tot{i} <: UpdraftPrognosticVariable{i} end\n\nprognostic_vars(m::EDMF) =\n (prognostic_vars(m.environment)..., prognostic_vars(m.updraft)...)\nprognostic_vars(m::Environment) =\n (en_ρatke(), en_ρaθ_liq_cv(), en_ρaq_tot_cv(), en_ρaθ_liq_q_tot_cv())\n\nfunction prognostic_vars(m::NTuple{N, Updraft}) where {N}\n t_ρa = vuntuple(i -> up_ρa{i}(), N)\n t_ρaw = vuntuple(i -> up_ρaw{i}(), N)\n t_ρaθ_liq = vuntuple(i -> up_ρaθ_liq{i}(), N)\n t_ρaq_tot = vuntuple(i -> up_ρaq_tot{i}(), N)\n t = (t_ρa..., t_ρaw..., t_ρaθ_liq..., t_ρaq_tot...)\n return t\nend\n\nget_prog_state(state, ::en_ρatke) = (state.turbconv.environment, :ρatke)\nget_prog_state(state, ::en_ρaθ_liq_cv) =\n (state.turbconv.environment, :ρaθ_liq_cv)\nget_prog_state(state, ::en_ρaq_tot_cv) =\n (state.turbconv.environment, :ρaq_tot_cv)\nget_prog_state(state, ::en_ρaθ_liq_q_tot_cv) =\n (state.turbconv.environment, :ρaθ_liq_q_tot_cv)\n\nget_prog_state(state, ::up_ρa{i}) where {i} = (state.turbconv.updraft[i], :ρa)\nget_prog_state(state, ::up_ρaw{i}) where {i} = (state.turbconv.updraft[i], :ρaw)\nget_prog_state(state, ::up_ρaθ_liq{i}) where {i} =\n (state.turbconv.updraft[i], :ρaθ_liq)\nget_prog_state(state, ::up_ρaq_tot{i}) where {i} =\n (state.turbconv.updraft[i], :ρaq_tot)\n\nget_specific_state(state, ::en_ρatke) = (state.turbconv.environment, :atke)\nget_specific_state(state, ::en_ρaθ_liq_cv) =\n (state.turbconv.environment, :aθ_liq_cv)\nget_specific_state(state, ::en_ρaq_tot_cv) =\n (state.turbconv.environment, :aq_tot_cv)\nget_specific_state(state, ::en_ρaθ_liq_q_tot_cv) =\n (state.turbconv.environment, :aθ_liq_q_tot_cv)\n\nget_specific_state(state, ::up_ρa{i}) where {i} =\n (state.turbconv.updraft[i], :a)\nget_specific_state(state, ::up_ρaw{i}) where {i} =\n (state.turbconv.updraft[i], :aw)\nget_specific_state(state, ::up_ρaθ_liq{i}) where {i} =\n (state.turbconv.updraft[i], :aθ_liq)\nget_specific_state(state, ::up_ρaq_tot{i}) where {i} =\n (state.turbconv.updraft[i], :aq_tot)\n\nstruct EntrDetr{N_up} <: TendencyDef{Source} end\nstruct PressSource{N_up} <: TendencyDef{Source} end\nstruct BuoySource{N_up} <: TendencyDef{Source} end\nstruct ShearSource <: TendencyDef{Source} end\nstruct DissSource <: TendencyDef{Source} end\nstruct GradProdSource <: TendencyDef{Source} end\n\nprognostic_vars(::EntrDetr{N_up}) where {N_up} = (\n vuntuple(i -> up_ρa{i}, N_up)...,\n vuntuple(i -> up_ρaw{i}, N_up)...,\n vuntuple(i -> up_ρaθ_liq{i}, N_up)...,\n vuntuple(i -> up_ρaq_tot{i}, N_up)...,\n en_ρatke(),\n en_ρaθ_liq_cv(),\n en_ρaq_tot_cv(),\n en_ρaθ_liq_q_tot_cv(),\n)\nprognostic_vars(::PressSource{N_up}) where {N_up} =\n vuntuple(i -> up_ρaw{i}(), N_up)\n\nprognostic_vars(::BuoySource{N_up}) where {N_up} =\n vuntuple(i -> up_ρaw{i}(), N_up)\n\nEntrDetr(m::EDMF) = EntrDetr{n_updrafts(m)}()\nBuoySource(m::EDMF) = BuoySource{n_updrafts(m)}()\nPressSource(m::EDMF) = PressSource{n_updrafts(m)}()\n\n# Dycore tendencies\neq_tends(\n pv::Union{Momentum, Energy, TotalMoisture},\n m::EDMF,\n flux::Flux{SecondOrder},\n) = eq_tends(pv, m.coupling, flux)\n\neq_tends(\n pv::Union{Momentum, Energy, TotalMoisture},\n ::Decoupled,\n ::Flux{SecondOrder},\n) = ()\n\neq_tends(\n pv::Union{Momentum, Energy, TotalMoisture},\n ::Coupled,\n ::Flux{SecondOrder},\n) = (SGSFlux(),)\n\n# Turbconv tendencies\neq_tends(pv::EDMFPrognosticVariable, m::AtmosModel, tt::Flux{O}) where {O} =\n eq_tends(pv, turbconv_model(m), tt)\n\neq_tends(::EDMFPrognosticVariable, m::EDMF, ::Flux{O}) where {O} = ()\n\neq_tends(::EnvironmentPrognosticVariable, m::EDMF, ::Flux{SecondOrder}) =\n (Diffusion(),)\n\neq_tends(pv::EDMFPrognosticVariable, m::EDMF, ::Flux{FirstOrder}) = (Advect(),)\n\neq_tends(pv::PV, m::EDMF, ::Source) where {PV} = ()\n\neq_tends(::EDMFPrognosticVariable, m::EDMF, ::Source) = (EntrDetr(m),)\n\neq_tends(pv::en_ρatke, m::EDMF, ::Source) =\n (EntrDetr(m), PressSource(m), BuoySource(m), ShearSource(), DissSource())\n\neq_tends(\n ::Union{en_ρaθ_liq_cv, en_ρaq_tot_cv, en_ρaθ_liq_q_tot_cv},\n m::EDMF,\n ::Source,\n) = (EntrDetr(m), DissSource(), GradProdSource())\n\neq_tends(::up_ρaw, m::EDMF, ::Source) =\n (EntrDetr(m), PressSource(m), BuoySource(m))\n\nstruct SGSFlux <: TendencyDef{Flux{SecondOrder}} end\n\n\"\"\"\n init_aux_turbconv!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n aux::Vars,\n geom::LocalGeometry,\n ) where {FT}\n\nInitialize EDMF auxiliary variables.\n\"\"\"\nfunction init_aux_turbconv!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n aux::Vars,\n geom::LocalGeometry,\n) where {FT} end;\n\nfunction turbconv_nodal_update_auxiliary_state!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n save_subdomain_temperature!(m, state, aux)\nend;\n\nfunction prognostic_to_primitive!(\n turbconv::EDMF,\n atmos,\n moist::DryModel,\n prim::Vars,\n prog::Vars,\n)\n N_up = n_updrafts(turbconv)\n prim_en = prim.turbconv.environment\n prog_en = prog.turbconv.environment\n prim_up = prim.turbconv.updraft\n prog_up = prog.turbconv.updraft\n\n ρ_inv = 1 / prog.ρ\n prim_en.atke = prog_en.ρatke * ρ_inv\n prim_en.aθ_liq_cv = prog_en.ρaθ_liq_cv * ρ_inv\n prim_en.aθ_liq_q_tot_cv = prog_en.ρaθ_liq_q_tot_cv * ρ_inv\n\n if moist isa DryModel\n prim_en.aq_tot_cv = 0\n else\n prim_en.aq_tot_cv = prog_en.ρaq_tot_cv * ρ_inv\n end\n @unroll_map(N_up) do i\n prim_up[i].a = prog_up[i].ρa * ρ_inv\n prim_up[i].aw = prog_up[i].ρaw * ρ_inv\n prim_up[i].aθ_liq = prog_up[i].ρaθ_liq * ρ_inv\n if moist isa DryModel\n prim_up[i].aq_tot = 0\n else\n prim_up[i].aq_tot = prog_up[i].ρaq_tot * ρ_inv\n end\n end\nend\n\nfunction primitive_to_prognostic!(\n turbconv::EDMF,\n atmos,\n moist::DryModel,\n prog::Vars,\n prim::Vars,\n)\n\n N_up = n_updrafts(turbconv)\n prim_en = prim.turbconv.environment\n prog_en = prog.turbconv.environment\n prim_up = prim.turbconv.updraft\n prog_up = prog.turbconv.updraft\n\n ρ_gm = prog.ρ\n prog_en.ρatke = prim_en.atke * ρ_gm\n prog_en.ρaθ_liq_cv = prim_en.aθ_liq_cv * ρ_gm\n prog_en.ρaθ_liq_q_tot_cv = prim_en.aθ_liq_q_tot_cv * ρ_gm\n\n if moist isa DryModel\n prog_en.ρaq_tot_cv = 0\n else\n prog_en.ρaq_tot_cv = prim_en.aq_tot_cv * ρ_gm\n end\n @unroll_map(N_up) do i\n prog_up[i].ρa = prim_up[i].a * ρ_gm\n prog_up[i].ρaw = prim_up[i].aw * ρ_gm\n prog_up[i].ρaθ_liq = prim_up[i].aθ_liq * ρ_gm\n if moist isa DryModel\n prog_up[i].ρaq_tot = 0\n else\n prog_up[i].ρaq_tot = prim_up[i].aq_tot * ρ_gm\n end\n end\n\nend\n\nfunction compute_gradient_argument!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n N_up = n_updrafts(turbconv)\n z = altitude(m, aux)\n\n # Aliases:\n gm_tf = transform.turbconv\n up_tf = transform.turbconv.updraft\n en_tf = transform.turbconv.environment\n gm = state\n up = state.turbconv.updraft\n en = state.turbconv.environment\n\n # Recover thermo states\n ts = recover_thermo_state_all(m, state, aux)\n\n # Get environment variables\n env = environment_vars(state, N_up)\n param_set = parameter_set(m)\n @unroll_map(N_up) do i\n up_tf[i].w = fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)\n end\n _grav::FT = grav(param_set)\n\n ρ_inv = 1 / gm.ρ\n θ_liq_en = liquid_ice_pottemp(ts.en)\n\n if moisture_model(m) isa DryModel\n q_tot_en = FT(0)\n else\n q_tot_en = total_specific_humidity(ts.en)\n end\n\n # populate gradient arguments\n en_tf.θ_liq = θ_liq_en\n en_tf.q_tot = q_tot_en\n en_tf.w = env.w\n\n en_tf.tke = en.ρatke / (env.a * gm.ρ)\n en_tf.θ_liq_cv = en.ρaθ_liq_cv / (env.a * gm.ρ)\n\n if moisture_model(m) isa DryModel\n en_tf.q_tot_cv = FT(0)\n en_tf.θ_liq_q_tot_cv = FT(0)\n else\n en_tf.q_tot_cv = en.ρaq_tot_cv / (env.a * gm.ρ)\n en_tf.θ_liq_q_tot_cv = en.ρaθ_liq_q_tot_cv / (env.a * gm.ρ)\n end\n\n en_tf.θv = virtual_pottemp(ts.en)\n en_e_kin =\n FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w^2) # TBD: Check\n en_e_tot = total_energy(en_e_kin, _grav * z, ts.en)\n en_tf.h_tot = total_specific_enthalpy(ts.en, en_e_tot)\n\n gm_tf.u = gm.ρu[1] * ρ_inv\n gm_tf.v = gm.ρu[2] * ρ_inv\nend;\n\nfunction compute_gradient_flux!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n args = (; diffusive, state, aux, t)\n N_up = n_updrafts(turbconv)\n\n # Aliases:\n gm = state\n gm_dif = diffusive.turbconv\n gm_∇tf = ∇transform.turbconv\n up_dif = diffusive.turbconv.updraft\n up_∇tf = ∇transform.turbconv.updraft\n en = state.turbconv.environment\n en_dif = diffusive.turbconv.environment\n en_∇tf = ∇transform.turbconv.environment\n\n @unroll_map(N_up) do i\n up_dif[i].∇w = up_∇tf[i].w\n end\n\n env = environment_vars(state, N_up)\n ρa₀ = gm.ρ * env.a\n # first moment grid mean coming from environment gradients only\n en_dif.∇θ_liq = en_∇tf.θ_liq\n en_dif.∇q_tot = en_∇tf.q_tot\n en_dif.∇w = en_∇tf.w\n # second moment env cov\n en_dif.∇tke = en_∇tf.tke\n en_dif.∇θ_liq_cv = en_∇tf.θ_liq_cv\n en_dif.∇q_tot_cv = en_∇tf.q_tot_cv\n en_dif.∇θ_liq_q_tot_cv = en_∇tf.θ_liq_q_tot_cv\n\n en_dif.∇θv = en_∇tf.θv\n en_dif.∇h_tot = en_∇tf.h_tot\n\n gm_dif.∇u = gm_∇tf.u\n gm_dif.∇v = gm_∇tf.v\nend;\n\nfunction source(::up_ρa{i}, ::EntrDetr, atmos, args) where {i}\n @unpack E_dyn, Δ_dyn, ρa_up = args.precomputed.turbconv\n return fix_void_up(ρa_up[i], E_dyn[i] - Δ_dyn[i])\nend\n\nfunction source(::up_ρaw{i}, ::EntrDetr, atmos, args) where {i}\n @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, w_up = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * env.w)\n detr = fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * w_up[i])\n\n return entr - detr\nend\n\nfunction source(::up_ρaθ_liq{i}, ::EntrDetr, atmos, args) where {i}\n @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, ts_en = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n θ_liq_en = liquid_ice_pottemp(ts_en)\n entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * θ_liq_en)\n detr =\n fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * up[i].ρaθ_liq / ρa_up[i])\n\n return entr - detr\nend\n\nfunction source(::up_ρaq_tot{i}, ::EntrDetr, atmos, args) where {i}\n @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n q_tot_en = total_specific_humidity(ts_en)\n entr = fix_void_up(ρa_up[i], (E_dyn[i] + E_trb[i]) * q_tot_en)\n detr =\n fix_void_up(ρa_up[i], (Δ_dyn[i] + E_trb[i]) * up[i].ρaq_tot / ρa_up[i])\n\n return entr - detr\nend\n\nfunction source(::en_ρatke, ::EntrDetr, atmos, args)\n @unpack E_dyn, Δ_dyn, E_trb, env, ρa_up, w_up = args.precomputed.turbconv\n @unpack state = args\n up = state.turbconv.updraft\n en = state.turbconv.environment\n gm = state\n N_up = n_updrafts(turbconv_model(atmos))\n ρ_inv = 1 / gm.ρ\n tke_en = enforce_positivity(en.ρatke) * ρ_inv / env.a\n\n entr_detr = vuntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n E_trb[i] * (env.w - gm.ρu[3] * ρ_inv) * (env.w - w_up[i]) -\n (E_dyn[i] + E_trb[i]) * tke_en +\n Δ_dyn[i] * (w_up[i] - env.w) * (w_up[i] - env.w) / 2,\n )\n end\n return sum(entr_detr)\nend\n\nfunction source(::en_ρaθ_liq_cv, ::EntrDetr, atmos, args)\n @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv\n @unpack state = args\n ts_gm = args.precomputed.ts\n up = state.turbconv.updraft\n en = state.turbconv.environment\n N_up = n_updrafts(turbconv_model(atmos))\n θ_liq = liquid_ice_pottemp(ts_gm)\n θ_liq_en = liquid_ice_pottemp(ts_en)\n\n entr_detr = vuntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n Δ_dyn[i] *\n (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) *\n (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) +\n E_trb[i] *\n (θ_liq_en - θ_liq) *\n (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) +\n E_trb[i] *\n (θ_liq_en - θ_liq) *\n (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) -\n (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_cv,\n )\n end\n return sum(entr_detr)\nend\n\nfunction source(::en_ρaq_tot_cv, ::EntrDetr, atmos, args)\n @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv\n @unpack state = args\n FT = eltype(state)\n up = state.turbconv.updraft\n en = state.turbconv.environment\n gm = state\n N_up = n_updrafts(turbconv_model(atmos))\n q_tot_en = total_specific_humidity(ts_en)\n ρ_inv = 1 / gm.ρ\n ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot\n\n entr_detr = vuntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n Δ_dyn[i] *\n (up[i].ρaq_tot / ρa_up[i] - q_tot_en) *\n (up[i].ρaq_tot / ρa_up[i] - q_tot_en) +\n E_trb[i] *\n (q_tot_en - ρq_tot * ρ_inv) *\n (q_tot_en - up[i].ρaq_tot / ρa_up[i]) +\n E_trb[i] *\n (q_tot_en - ρq_tot * ρ_inv) *\n (q_tot_en - up[i].ρaq_tot / ρa_up[i]) -\n (E_dyn[i] + E_trb[i]) * en.ρaq_tot_cv,\n )\n end\n return sum(entr_detr)\nend\n\nfunction source(::en_ρaθ_liq_q_tot_cv, ::EntrDetr, atmos, args)\n @unpack E_dyn, Δ_dyn, E_trb, ρa_up, ts_en = args.precomputed.turbconv\n @unpack state = args\n FT = eltype(state)\n ts_gm = args.precomputed.ts\n up = state.turbconv.updraft\n en = state.turbconv.environment\n gm = state\n N_up = n_updrafts(turbconv_model(atmos))\n q_tot_en = total_specific_humidity(ts_en)\n θ_liq = liquid_ice_pottemp(ts_gm)\n θ_liq_en = liquid_ice_pottemp(ts_en)\n ρ_inv = 1 / gm.ρ\n ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot\n\n entr_detr = vuntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n Δ_dyn[i] *\n (up[i].ρaθ_liq / ρa_up[i] - θ_liq_en) *\n (up[i].ρaq_tot / ρa_up[i] - q_tot_en) +\n E_trb[i] *\n (θ_liq_en - θ_liq) *\n (q_tot_en - up[i].ρaq_tot / ρa_up[i]) +\n E_trb[i] *\n (q_tot_en - ρq_tot * ρ_inv) *\n (θ_liq_en - up[i].ρaθ_liq / ρa_up[i]) -\n (E_dyn[i] + E_trb[i]) * en.ρaθ_liq_q_tot_cv,\n )\n end\n return sum(entr_detr)\nend\n\nfunction source(::en_ρatke, ::PressSource, atmos, args)\n @unpack env, ρa_up, dpdz, w_up = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n N_up = n_updrafts(turbconv_model(atmos))\n press_tke = vuntuple(N_up) do i\n fix_void_up(ρa_up[i], ρa_up[i] * (w_up[i] - env.w) * dpdz[i])\n end\n return sum(press_tke)\nend\n\nfunction source(::en_ρatke, ::ShearSource, atmos, args)\n @unpack env, K_m, Shear² = args.precomputed.turbconv\n gm = args.state\n ρa₀ = gm.ρ * env.a\n # production from mean gradient and Dissipation\n return ρa₀ * K_m * Shear² # tke Shear source\nend\n\nfunction source(::en_ρatke, ::BuoySource, atmos, args)\n @unpack env, K_h, ∂b∂z_env = args.precomputed.turbconv\n gm = args.state\n ρa₀ = gm.ρ * env.a\n return -ρa₀ * K_h * ∂b∂z_env # tke Buoyancy source\nend\n\nfunction source(::en_ρatke, ::DissSource, atmos, args)\n @unpack Diss₀ = args.precomputed.turbconv\n en = args.state.turbconv.environment\n return -Diss₀ * en.ρatke # tke Dissipation\nend\n\nfunction source(::en_ρaθ_liq_cv, ::DissSource, atmos, args)\n @unpack Diss₀ = args.precomputed.turbconv\n en = args.state.turbconv.environment\n return -Diss₀ * en.ρaθ_liq_cv\nend\n\nfunction source(::en_ρaq_tot_cv, ::DissSource, atmos, args)\n @unpack Diss₀ = args.precomputed.turbconv\n en = args.state.turbconv.environment\n return -Diss₀ * en.ρaq_tot_cv\nend\n\nfunction source(::en_ρaθ_liq_q_tot_cv, ::DissSource, atmos, args)\n @unpack Diss₀ = args.precomputed.turbconv\n en = args.state.turbconv.environment\n return -Diss₀ * en.ρaθ_liq_q_tot_cv\nend\n\nfunction source(::en_ρaθ_liq_cv, ::GradProdSource, atmos, args)\n @unpack env, K_h = args.precomputed.turbconv\n gm = args.state\n en_dif = args.diffusive.turbconv.environment\n ρa₀ = gm.ρ * env.a\n return ρa₀ * (2 * K_h * en_dif.∇θ_liq[3] * en_dif.∇θ_liq[3])\nend\n\nfunction source(::en_ρaq_tot_cv, ::GradProdSource, atmos, args)\n @unpack env, K_h = args.precomputed.turbconv\n gm = args.state\n en_dif = args.diffusive.turbconv.environment\n ρa₀ = gm.ρ * env.a\n return ρa₀ * (2 * K_h * en_dif.∇q_tot[3] * en_dif.∇q_tot[3])\nend\n\nfunction source(::en_ρaθ_liq_q_tot_cv, ::GradProdSource, atmos, args)\n @unpack env, K_h = args.precomputed.turbconv\n gm = args.state\n en_dif = args.diffusive.turbconv.environment\n ρa₀ = gm.ρ * env.a\n return ρa₀ * (2 * K_h * en_dif.∇θ_liq[3] * en_dif.∇q_tot[3])\nend\n\nfunction source(::up_ρaw{i}, ::BuoySource, atmos, args) where {i}\n @unpack buoy = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n return up[i].ρa * buoy.up[i]\nend\n\nfunction source(::up_ρaw{i}, ::PressSource, atmos, args) where {i}\n @unpack dpdz = args.precomputed.turbconv\n up = args.state.turbconv.updraft\n return -up[i].ρa * dpdz[i]\nend\n\nfunction compute_ρa_up(atmos, state, aux)\n # Aliases:\n turbconv = turbconv_model(atmos)\n gm = state\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n a_min = turbconv.subdomains.a_min\n a_max = turbconv.subdomains.a_max\n # in future GCM implementations we need to think about grid mean advection\n ρa_up = vuntuple(N_up) do i\n gm.ρ * enforce_unit_bounds(up[i].ρa / gm.ρ, a_min, a_max)\n end\n return ρa_up\nend\n\nfunction flux(::up_ρa{i}, ::Advect, atmos, args) where {i}\n @unpack state, aux = args\n @unpack ρa_up = args.precomputed.turbconv\n up = state.turbconv.updraft\n ẑ = vertical_unit_vector(atmos, aux)\n return fix_void_up(ρa_up[i], up[i].ρaw) * ẑ\nend\nfunction flux(::up_ρaw{i}, ::Advect, atmos, args) where {i}\n @unpack state, aux = args\n @unpack ρa_up, w_up = args.precomputed.turbconv\n up = state.turbconv.updraft\n ẑ = vertical_unit_vector(atmos, aux)\n return fix_void_up(ρa_up[i], up[i].ρaw * w_up[i]) * ẑ\n\nend\nfunction flux(::up_ρaθ_liq{i}, ::Advect, atmos, args) where {i}\n @unpack state, aux = args\n @unpack ρa_up, w_up = args.precomputed.turbconv\n up = state.turbconv.updraft\n ẑ = vertical_unit_vector(atmos, aux)\n return fix_void_up(ρa_up[i], w_up[i] * up[i].ρaθ_liq) * ẑ\n\nend\nfunction flux(::up_ρaq_tot{i}, ::Advect, atmos, args) where {i}\n @unpack state, aux = args\n @unpack ρa_up, w_up = args.precomputed.turbconv\n up = state.turbconv.updraft\n ẑ = vertical_unit_vector(atmos, aux)\n return fix_void_up(ρa_up[i], w_up[i] * up[i].ρaq_tot) * ẑ\n\nend\n\nfunction flux(::en_ρatke, ::Advect, atmos, args)\n @unpack state, aux = args\n @unpack env = args.precomputed.turbconv\n en = state.turbconv.environment\n ẑ = vertical_unit_vector(atmos, aux)\n return en.ρatke * env.w * ẑ\nend\nfunction flux(::en_ρaθ_liq_cv, ::Advect, atmos, args)\n @unpack state, aux = args\n @unpack env = args.precomputed.turbconv\n en = state.turbconv.environment\n ẑ = vertical_unit_vector(atmos, aux)\n return en.ρaθ_liq_cv * env.w * ẑ\nend\nfunction flux(::en_ρaq_tot_cv, ::Advect, atmos, args)\n @unpack state, aux = args\n @unpack env = args.precomputed.turbconv\n en = state.turbconv.environment\n ẑ = vertical_unit_vector(atmos, aux)\n return en.ρaq_tot_cv * env.w * ẑ\nend\nfunction flux(::en_ρaθ_liq_q_tot_cv, ::Advect, atmos, args)\n @unpack state, aux = args\n @unpack env = args.precomputed.turbconv\n en = state.turbconv.environment\n ẑ = vertical_unit_vector(atmos, aux)\n return en.ρaθ_liq_q_tot_cv * env.w * ẑ\nend\n\nfunction precompute(::EDMF, bl, args, ts, ::Flux{FirstOrder})\n @unpack state, aux = args\n FT = eltype(state)\n turbconv = turbconv_model(bl)\n env = environment_vars(state, n_updrafts(turbconv))\n ρa_up = compute_ρa_up(bl, state, aux)\n up = state.turbconv.updraft\n gm = state\n N_up = n_updrafts(turbconv)\n ρ_inv = 1 / gm.ρ\n w_up = vuntuple(N_up) do i\n fix_void_up(ρa_up[i], up[i].ρaw / ρa_up[i])\n end\n\n θ_liq_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρaθ_liq / up[i].ρa, liquid_ice_pottemp(ts))\n end\n a_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρa * ρ_inv)\n end\n if !(moisture_model(bl) isa DryModel)\n q_tot_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρaq_tot / up[i].ρa, gm.moisture.ρq_tot)\n end\n else\n q_tot_up = vuntuple(i -> FT(0), N_up)\n end\n\n return (; env, a_up, q_tot_up, ρa_up, θ_liq_up, w_up)\nend\n\nfunction precompute(::EDMF, bl, args, ts, ::Flux{SecondOrder})\n @unpack state, aux, diffusive, t = args\n en_dif = diffusive.turbconv.environment\n up = state.turbconv.updraft\n gm = state\n ts_gm = ts\n FT = eltype(state)\n z = altitude(bl, aux)\n param_set = parameter_set(bl)\n turbconv = turbconv_model(bl)\n N_up = n_updrafts(turbconv)\n _grav::FT = grav(param_set)\n ρ_inv = 1 / gm.ρ\n\n env = environment_vars(state, N_up)\n ts_en = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts_gm)\n ts_up = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts_gm)\n\n buoy = compute_buoyancy(bl, state, env, ts_en, ts_up, aux.ref_state)\n\n E_dyn, Δ_dyn, E_trb = entr_detr(bl, state, aux, ts_up, ts_en, env, buoy)\n\n Shear² =\n diffusive.turbconv.∇u[3]^2 +\n diffusive.turbconv.∇v[3]^2 +\n diffusive.turbconv.environment.∇w[3]^2\n l_mix, ∂b∂z_env, Pr_t = mixing_length(\n bl,\n turbconv.mix_len,\n args,\n Δ_dyn,\n E_trb,\n Shear²,\n ts_gm,\n ts_en,\n env,\n )\n ρa_up = compute_ρa_up(bl, state, aux)\n\n en = state.turbconv.environment\n tke_en = enforce_positivity(en.ρatke) / env.a / state.ρ\n K_m = turbconv.mix_len.c_m * l_mix * sqrt(tke_en)\n K_h = K_m / Pr_t\n ρaw_up = vuntuple(i -> up[i].ρaw, N_up)\n\n w_up = vuntuple(N_up) do i\n fix_void_up(ρa_up[i], up[i].ρaw / ρa_up[i])\n end\n\n ρu_gm_tup = Tuple(gm.ρu)\n ρa_en = gm.ρ * env.a\n # TODO: Consider turbulent contribution:\n\n e_kin_up =\n FT(1 / 2) .* ntuple(N_up) do i\n (ρu_gm_tup[1] * ρ_inv)^2 + (ρu_gm_tup[2] * ρ_inv)^2 + w_up[i]^2\n end\n e_kin_en =\n FT(1 // 2) * ((gm.ρu[1] * ρ_inv)^2 + (gm.ρu[2] * ρ_inv)^2 + env.w)^2\n\n e_tot_up = ntuple(i -> total_energy(e_kin_up[i], _grav * z, ts_up[i]), N_up)\n h_tot_up = ntuple(i -> total_specific_enthalpy(ts_up[i], e_tot_up[i]), N_up)\n\n e_tot_en = total_energy(e_kin_en, _grav * z, ts_en)\n h_tot_en = total_specific_enthalpy(ts_en, e_tot_en)\n h_tot_gm = total_specific_enthalpy(ts, gm.energy.ρe * ρ_inv)\n\n massflux_h_tot = sum(\n ntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n ρa_up[i] *\n (h_tot_gm - h_tot_up[i]) *\n (gm.ρu[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),\n )\n end,\n )\n massflux_h_tot +=\n (ρa_en * (h_tot_gm - h_tot_en) * (ρu_gm_tup[3] * ρ_inv - env.w))\n\n ρh_sgs_flux = SVector{3, FT}(\n 0,\n 0,\n -gm.ρ * env.a * K_h * en_dif.∇h_tot[3] + massflux_h_tot,\n )\n\n return (;\n env,\n ρa_up,\n ρaw_up,\n ts_en,\n ts_up,\n E_dyn,\n Δ_dyn,\n E_trb,\n l_mix,\n ∂b∂z_env,\n K_h,\n K_m,\n Pr_t,\n ρh_sgs_flux,\n )\nend\n\n\"\"\"\n compute_buoyancy(\n bl::BalanceLaw,\n state::Vars,\n env::NamedTuple,\n ts_en::ThermodynamicState,\n ts_up,\n ref_state::Vars\n )\n\nCompute buoyancies of subdomains\n\"\"\"\nfunction compute_buoyancy(\n bl::BalanceLaw,\n state::Vars,\n env::NamedTuple,\n ts_en::ThermodynamicState,\n ts_up,\n ref_state::Vars,\n)\n FT = eltype(state)\n param_set = parameter_set(bl)\n N_up = n_updrafts(turbconv_model(bl))\n _grav::FT = grav(param_set)\n gm = state\n ρ_inv = 1 / gm.ρ\n buoyancy_en = -_grav * (air_density(ts_en) - ref_state.ρ) * ρ_inv\n up = state.turbconv.updraft\n\n a_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρa * ρ_inv)\n end\n\n abs_buoyancy_up = vuntuple(N_up) do i\n -_grav * (air_density(ts_up[i]) - ref_state.ρ) * ρ_inv\n end\n b_gm = grid_mean_b(env, a_up, N_up, abs_buoyancy_up, buoyancy_en)\n\n # remove the gm_b from all subdomains\n buoyancy_up = vuntuple(N_up) do i\n abs_buoyancy_up[i] - b_gm\n end\n\n buoyancy_en -= b_gm\n return (; up = buoyancy_up, en = buoyancy_en)\nend\n\nfunction precompute(::EDMF, bl, args, ts, ::Source)\n @unpack state, aux, diffusive, t = args\n ts_gm = ts\n gm = state\n up = state.turbconv.updraft\n turbconv = turbconv_model(bl)\n N_up = n_updrafts(turbconv)\n # Get environment variables\n env = environment_vars(state, N_up)\n # Recover thermo states\n ts_en = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts_gm)\n ts_up = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts_gm)\n ρa_up = compute_ρa_up(bl, state, aux)\n\n Shear² =\n diffusive.turbconv.∇u[3]^2 +\n diffusive.turbconv.∇v[3]^2 +\n diffusive.turbconv.environment.∇w[3]^2\n\n buoy = compute_buoyancy(bl, state, env, ts_en, ts_up, aux.ref_state)\n E_dyn, Δ_dyn, E_trb = entr_detr(bl, state, aux, ts_up, ts_en, env, buoy)\n\n dpdz = perturbation_pressure(bl, args, env, buoy)\n\n l_mix, ∂b∂z_env, Pr_t = mixing_length(\n bl,\n turbconv.mix_len,\n args,\n Δ_dyn,\n E_trb,\n Shear²,\n ts_gm,\n ts_en,\n env,\n )\n\n w_up = vuntuple(N_up) do i\n fix_void_up(ρa_up[i], up[i].ρaw / ρa_up[i])\n end\n\n en = state.turbconv.environment\n tke_en = enforce_positivity(en.ρatke) / env.a / state.ρ\n K_m = turbconv.mix_len.c_m * l_mix * sqrt(tke_en)\n K_h = K_m / Pr_t\n Diss₀ = turbconv.mix_len.c_d * sqrt(tke_en) / l_mix\n tke_buoy_prod = -gm.ρ * env.a * K_h * ∂b∂z_env # tke Buoyancy source\n\n return (;\n env,\n Diss₀,\n buoy,\n K_m,\n K_h,\n ρa_up,\n w_up,\n ts_en,\n ts_up,\n E_dyn,\n Δ_dyn,\n E_trb,\n dpdz,\n l_mix,\n ∂b∂z_env,\n Pr_t,\n Shear²,\n tke_buoy_prod,\n )\nend\n\nfunction flux(::Energy, ::SGSFlux, atmos, args)\n @unpack state, aux, diffusive = args\n @unpack ts = args.precomputed\n @unpack ρh_sgs_flux = args.precomputed.turbconv\n return ρh_sgs_flux\nend\n\nfunction flux(::TotalMoisture, ::SGSFlux, atmos, args)\n @unpack state, diffusive = args\n @unpack env, K_h, ρa_up, ρaw_up, ts_en = args.precomputed.turbconv\n FT = eltype(state)\n en_dif = diffusive.turbconv.environment\n up = state.turbconv.updraft\n gm = state\n ρ_inv = 1 / gm.ρ\n N_up = n_updrafts(turbconv_model(atmos))\n ρq_tot = moisture_model(atmos) isa DryModel ? FT(0) : gm.moisture.ρq_tot\n ρaq_tot_up = vuntuple(i -> up[i].ρaq_tot, N_up)\n ρa_en = gm.ρ * env.a\n q_tot_en = total_specific_humidity(ts_en)\n\n ρu_gm_tup = Tuple(gm.ρu)\n\n massflux_q_tot = sum(\n ntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n ρa_up[i] *\n (ρq_tot * ρ_inv - ρaq_tot_up[i] / ρa_up[i]) *\n (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),\n )\n end,\n )\n massflux_q_tot +=\n (ρa_en * (ρq_tot * ρ_inv - q_tot_en) * (ρu_gm_tup[3] * ρ_inv - env.w))\n ρq_tot_sgs_flux = -gm.ρ * env.a * K_h * en_dif.∇q_tot[3] + massflux_q_tot\n return SVector{3, FT}(0, 0, ρq_tot_sgs_flux)\nend\n\nfunction flux(::Momentum, ::SGSFlux, atmos, args)\n @unpack state, diffusive = args\n @unpack env, K_m, ρa_up, ρaw_up = args.precomputed.turbconv\n FT = eltype(state)\n en_dif = diffusive.turbconv.environment\n gm_dif = diffusive.turbconv\n up = state.turbconv.updraft\n gm = state\n ρ_inv = 1 / gm.ρ\n N_up = n_updrafts(turbconv_model(atmos))\n ρa_en = gm.ρ * env.a\n\n ρu_gm_tup = Tuple(gm.ρu)\n\n massflux_w = sum(\n ntuple(N_up) do i\n fix_void_up(\n ρa_up[i],\n ρa_up[i] *\n (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]) *\n (ρu_gm_tup[3] * ρ_inv - ρaw_up[i] / ρa_up[i]),\n )\n end,\n )\n massflux_w += (\n ρa_en * (ρu_gm_tup[3] * ρ_inv - env.w) * (ρu_gm_tup[3] * ρ_inv - env.w)\n )\n ρw_sgs_flux = -gm.ρ * env.a * K_m * en_dif.∇w[3] + massflux_w\n ρu_sgs_flux = -gm.ρ * env.a * K_m * gm_dif.∇u[3]\n ρv_sgs_flux = -gm.ρ * env.a * K_m * gm_dif.∇v[3]\n return SMatrix{3, 3, FT, 9}(\n 0,\n 0,\n ρu_sgs_flux,\n 0,\n 0,\n ρv_sgs_flux,\n 0,\n 0,\n ρw_sgs_flux,\n )\nend\n\nfunction flux(::en_ρaθ_liq_cv, ::Diffusion, atmos, args)\n @unpack state, aux, diffusive = args\n @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv\n en_dif = diffusive.turbconv.environment\n gm = state\n ẑ = vertical_unit_vector(atmos, aux)\n return -gm.ρ * env.a * K_h * en_dif.∇θ_liq_cv[3] * ẑ\nend\nfunction flux(::en_ρaq_tot_cv, ::Diffusion, atmos, args)\n @unpack state, aux, diffusive = args\n @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv\n en_dif = diffusive.turbconv.environment\n gm = state\n ẑ = vertical_unit_vector(atmos, aux)\n return -gm.ρ * env.a * K_h * en_dif.∇q_tot_cv[3] * ẑ\nend\nfunction flux(::en_ρaθ_liq_q_tot_cv, ::Diffusion, atmos, args)\n @unpack state, aux, diffusive = args\n @unpack env, l_mix, Pr_t, K_h = args.precomputed.turbconv\n en_dif = diffusive.turbconv.environment\n gm = state\n ẑ = vertical_unit_vector(atmos, aux)\n return -gm.ρ * env.a * K_h * en_dif.∇θ_liq_q_tot_cv[3] * ẑ\nend\nfunction flux(::en_ρatke, ::Diffusion, atmos, args)\n @unpack state, aux, diffusive = args\n @unpack env, K_m = args.precomputed.turbconv\n gm = state\n en_dif = diffusive.turbconv.environment\n ẑ = vertical_unit_vector(atmos, aux)\n return -gm.ρ * env.a * K_m * en_dif.∇tke[3] * ẑ\nend\n\n# First order boundary conditions\nfunction turbconv_boundary_state!(\n nf,\n bc::EDMFBottomBC,\n atmos::AtmosModel{FT},\n state⁺::Vars,\n args,\n) where {FT}\n @unpack state⁻, aux⁻, aux_int⁻ = args\n turbconv = turbconv_model(atmos)\n N_up = n_updrafts(turbconv)\n up⁺ = state⁺.turbconv.updraft\n en⁺ = state⁺.turbconv.environment\n gm⁻ = state⁻\n gm_a⁻ = aux⁻\n\n zLL = altitude(atmos, aux_int⁻)\n surf_vals = subdomain_surface_values(atmos, gm⁻, gm_a⁻, zLL)\n a_up_surf = surf_vals.a_up_surf\n\n @unroll_map(N_up) do i\n up⁺[i].ρaw = FT(0)\n up⁺[i].ρa = gm⁻.ρ * a_up_surf[i]\n up⁺[i].ρaθ_liq = gm⁻.ρ * a_up_surf[i] * surf_vals.θ_liq_up_surf[i]\n if !(moisture_model(atmos) isa DryModel)\n up⁺[i].ρaq_tot = gm⁻.ρ * a_up_surf[i] * surf_vals.q_tot_up_surf[i]\n else\n up⁺[i].ρaq_tot = FT(0)\n end\n end\n\n a_en = environment_area(gm⁻, N_up)\n en⁺.ρatke = gm⁻.ρ * a_en * surf_vals.tke\n en⁺.ρaθ_liq_cv = gm⁻.ρ * a_en * surf_vals.θ_liq_cv\n if !(moisture_model(atmos) isa DryModel)\n en⁺.ρaq_tot_cv = gm⁻.ρ * a_en * surf_vals.q_tot_cv\n en⁺.ρaθ_liq_q_tot_cv = gm⁻.ρ * a_en * surf_vals.θ_liq_q_tot_cv\n else\n en⁺.ρaq_tot_cv = FT(0)\n en⁺.ρaθ_liq_q_tot_cv = FT(0)\n end\nend;\nfunction turbconv_boundary_state!(\n nf,\n bc::EDMFTopBC,\n atmos::AtmosModel{FT},\n state⁺::Vars,\n args,\n) where {FT}\n N_up = n_updrafts(turbconv_model(atmos))\n up⁺ = state⁺.turbconv.updraft\n @unroll_map(N_up) do i\n up⁺[i].ρaw = FT(0)\n end\nend;\n\n\n# The boundary conditions for second-order unknowns\n# (here we prescribe a flux at state⁺ to match that at state⁻ so that the flux divergence is zero)\nfunction turbconv_normal_boundary_flux_second_order!(\n nf,\n bc::EDMFBottomBC,\n atmos::AtmosModel,\n fluxᵀn::Vars,\n args,\n)\n @unpack state⁻, aux⁻, diffusive⁻, hyperdiff⁻, t, n⁻ = args\n en_flx = fluxᵀn.turbconv.environment\n tend_type = Flux{SecondOrder}()\n _args⁻ = (;\n state = state⁻,\n aux = aux⁻,\n t,\n diffusive = diffusive⁻,\n hyperdiffusive = hyperdiff⁻,\n )\n pargs = merge(_args⁻, (precomputed = precompute(atmos, _args⁻, tend_type),))\n\n total_flux = Σfluxes(\n en_ρatke(),\n eq_tends(en_ρatke(), atmos, tend_type),\n atmos,\n pargs,\n )\n nd_ρatke = dot(n⁻, total_flux)\n en_flx.ρatke = nd_ρatke\n\n total_flux = Σfluxes(\n en_ρaθ_liq_cv(),\n eq_tends(en_ρaθ_liq_cv(), atmos, tend_type),\n atmos,\n pargs,\n )\n nd_ρaθ_liq_cv = dot(n⁻, total_flux)\n en_flx.ρaθ_liq_cv = nd_ρaθ_liq_cv\n\n if !(moisture_model(atmos) isa DryModel)\n total_flux = Σfluxes(\n en_ρaq_tot_cv(),\n eq_tends(en_ρaq_tot_cv(), atmos, tend_type),\n atmos,\n pargs,\n )\n nd_ρaq_tot_cv = dot(n⁻, total_flux)\n en_flx.ρaq_tot_cv = nd_ρaq_tot_cv\n\n total_flux = Σfluxes(\n en_ρaθ_liq_q_tot_cv(),\n eq_tends(en_ρaθ_liq_q_tot_cv(), atmos, tend_type),\n atmos,\n pargs,\n )\n nd_ρaθ_liq_q_tot_cv = dot(n⁻, total_flux)\n en_flx.ρaθ_liq_q_tot_cv = nd_ρaθ_liq_q_tot_cv\n end\nend;\n\nfunction turbconv_normal_boundary_flux_second_order!(\n nf,\n bc::EDMFTopBC,\n atmos::AtmosModel{FT},\n fluxᵀn::Vars,\n args,\n) where {FT}\n\n turbconv = turbconv_model(atmos)\n N_up = n_updrafts(turbconv)\n up_flx = fluxᵀn.turbconv.updraft\n en_flx = fluxᵀn.turbconv.environment\n @unroll_map(N_up) do i\n up_flx[i].ρa = FT(0)\n up_flx[i].ρaθ_liq = FT(0)\n up_flx[i].ρaq_tot = FT(0)\n end\n en_flx.ρatke = FT(0)\n en_flx.ρaθ_liq_cv = FT(0)\n en_flx.ρaq_tot_cv = FT(0)\n en_flx.ρaθ_liq_q_tot_cv = FT(0)\n\nend;\n" }, { "path": "test/Atmos/EDMF/edmf_model.jl", "content": "#### EDMF model\nusing DocStringExtensions\nusing CLIMAParameters: AbstractEarthParameterSet\nusing CLIMAParameters.Atmos.EDMF\nusing CLIMAParameters.SubgridScale\n\n\"\"\"\n EntrainmentDetrainment\n\nAn Entrainment-Detrainment model for EDMF, containing\nall related model and free parameters.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct EntrainmentDetrainment{FT <: AbstractFloat}\n \"Entrainment TKE scale\"\n c_λ::FT\n \"Entrainment factor\"\n c_ε::FT\n \"Detrainment factor\"\n c_δ::FT\n \"Turbulent Entrainment factor\"\n c_t::FT\n \"Detrainment RH power\"\n β::FT\n \"Logistic function scale ‵[1/s]‵\"\n μ_0::FT\n \"Updraft mixing fraction\"\n χ::FT\n \"Minimum updraft velocity\"\n w_min::FT\n \"Exponential area limiter scale\"\n lim_ϵ::FT\n \"Exponential area limiter amplitude\"\n lim_amp::FT\nend\n\n\"\"\"\n EntrainmentDetrainment{FT}(param_set) where {FT}\n\nConstructor for `EntrainmentDetrainment` for EDMF, given:\n - `param_set`, an AbstractEarthParameterSet\n\"\"\"\nfunction EntrainmentDetrainment{FT}(\n param_set::AbstractEarthParameterSet,\n) where {FT}\n c_λ_ = c_λ(param_set)\n c_ε_ = c_ε(param_set)\n c_δ_ = c_δ(param_set)\n c_t_ = c_t(param_set)\n β_ = β(param_set)\n μ_0_ = μ_0(param_set)\n χ_ = χ(param_set)\n w_min_ = w_min(param_set)\n lim_ϵ_ = lim_ϵ(param_set)\n lim_amp_ = lim_amp(param_set)\n\n args = (c_λ_, c_ε_, c_δ_, c_t_, β_, μ_0_, χ_, w_min_, lim_ϵ_, lim_amp_)\n\n return EntrainmentDetrainment{FT}(args...)\nend\n\n\"\"\"\n SubdomainModel\n\nA subdomain model for EDMF, containing\nall related model and free parameters.\n\nTODO: `a_max` is valid for all subdomains,\n but it is insufficient to ensure `a_en`\n is not negative. Limits can be imposed\n for updrafts, but this is a limit\n is dictated by 1 - Σᵢ aᵢ, which must be\n somehow satisfied by regularizing prognostic\n source terms.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct SubdomainModel{FT <: AbstractFloat}\n \"Minimum area fraction for any subdomain\"\n a_min::FT\n \"Maximum area fraction for any subdomain\"\n a_max::FT\nend\n\nfunction SubdomainModel(\n ::Type{FT},\n N_up;\n a_min::FT = FT(0),\n a_max::FT = 1 - N_up * a_min,\n) where {FT}\n return SubdomainModel(; a_min = a_min, a_max = a_max)\nend\n\n\"\"\"\n SurfaceModel\n\nA surface model for EDMF, containing all boundary\nvalues and parameters needed by the model.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct SurfaceModel{FT <: AbstractFloat, SV}\n \"Area\"\n a::FT\n \"Surface covariance stability coefficient\"\n ψϕ_stab::FT\n \"Square ratio of rms turbulent velocity to friction velocity\"\n κ_star²::FT\n \"Updraft normalized standard deviation at the surface\"\n upd_surface_std::SV\n # The following will be deleted after SurfaceFlux coupling\n \"Liquid water potential temperature ‵[k]‵\"\n θ_liq::FT = 299.1\n \"Specific humidity ‵[kg/kg]‵\"\n q_tot::FT = 22.45e-3\n \"Sensible heat flux ‵[w/m^2]‵\"\n shf::FT = 9.5\n \"Latent heat flux ‵[w/m^2]‵\"\n lhf::FT = 147.2\n \"Friction velocity\"\n ustar::FT = 0.28\n \"Monin - Obukhov length\"\n obukhov_length::FT = 0\n \"Height of the lowest level\"\n zLL::FT = 60\nend\n\n\"\"\"\n SurfaceModel{FT}(N_up, param_set) where {FT}\n\nConstructor for `SurfaceModel` for EDMF, given:\n - `N_up`, the number of updrafts\n - `param_set`, an AbstractEarthParameterSet\n\"\"\"\nfunction SurfaceModel{FT}(N_up, param_set::AbstractEarthParameterSet) where {FT}\n a_surf_ = a_surf(param_set)\n κ_star²_ = κ_star²(param_set)\n ψϕ_stab_ = ψϕ_stab(param_set)\n\n if a_surf_ > FT(0)\n upd_surface_std = SVector(\n ntuple(N_up) do i\n percentile_bounds_mean_norm(\n 1 - a_surf_ + (i - 1) * FT(a_surf_ / N_up),\n 1 - a_surf_ + i * FT(a_surf_ / N_up),\n 1000,\n )\n end,\n )\n else\n upd_surface_std = SVector(ntuple(i -> FT(0), N_up))\n end\n SV = typeof(upd_surface_std)\n return SurfaceModel{FT, SV}(;\n upd_surface_std = upd_surface_std,\n a = a_surf_,\n κ_star² = κ_star²_,\n ψϕ_stab = ψϕ_stab_,\n )\nend\n\n\"\"\"\n NeutralDrySurfaceModel\n\nA surface model for EDMF simulations in a\ndry, neutral environment, containing all boundary\nvalues and parameters needed by the model.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct NeutralDrySurfaceModel{FT <: AbstractFloat}\n \"Area\"\n a::FT\n \"Square ratio of rms turbulent velocity to friction velocity\"\n κ_star²::FT\n \"Friction velocity\"\n ustar::FT = 0.3\n \"Height of the lowest level\"\n zLL::FT = 60\n \"Monin - Obukhov length\"\n obukhov_length::FT = 0\nend\n\n\"\"\"\n NeutralDrySurfaceModel{FT}(N_up, param_set) where {FT}\n\nConstructor for `NeutralDrySurfaceModel` for EDMF, given:\n - `param_set`, an AbstractEarthParameterSet\n\"\"\"\nfunction NeutralDrySurfaceModel{FT}(\n param_set::AbstractEarthParameterSet,\n) where {FT}\n a_surf_ = a_surf(param_set)\n κ_star²_ = κ_star²(param_set)\n return NeutralDrySurfaceModel{FT}(; a = a_surf_, κ_star² = κ_star²_)\nend\n\n\"\"\"\n PressureModel\n\nA pressure model for EDMF, containing\nall related model and free parameters.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct PressureModel{FT <: AbstractFloat}\n \"Pressure drag\"\n α_d::FT\n \"Pressure advection\"\n α_a::FT\n \"Pressure buoyancy\"\n α_b::FT\n \"Minimum diagnostic updraft height for closures\"\n H_up_min::FT\nend\n\n\"\"\"\n PressureModel{FT}(param_set) where {FT}\n\nConstructor for `PressureModel` for EDMF, given:\n - `param_set`, an AbstractEarthParameterSet\n\"\"\"\nfunction PressureModel{FT}(param_set::AbstractEarthParameterSet) where {FT}\n α_d_ = α_d(param_set)\n α_a_ = α_a(param_set)\n α_b_ = α_b(param_set)\n H_up_min_ = H_up_min(param_set)\n\n args = (α_d_, α_a_, α_b_, H_up_min_)\n\n return PressureModel{FT}(args...)\nend\n\n\"\"\"\n MixingLengthModel\n\nA mixing length model for EDMF, containing\nall related model and free parameters.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct MixingLengthModel{FT <: AbstractFloat}\n \"dissipation coefficient\"\n c_d::FT\n \"Eddy Viscosity\"\n c_m::FT\n \"Static Stability coefficient\"\n c_b::FT\n \"Empirical stability function coefficient\"\n a1::FT\n \"Empirical stability function coefficient\"\n a2::FT\n \"Von Karmen constant\"\n κ::FT\n \"Prandtl number empirical coefficient\"\n ω_pr::FT\n \"Prandtl number scale\"\n Pr_n::FT\n \"Critical Richardson number\"\n Ri_c::FT\n \"smooth minimum's fractional upper bound\"\n smin_ub::FT\n \"smooth minimum's regularization minimum\"\n smin_rm::FT\n \"Maximum mixing length\"\n max_length::FT\n \"Random small number variable that should be addressed\"\n random_minval::FT\nend\n\n\"\"\"\n MixingLengthModel{FT}(param_set) where {FT}\n\nConstructor for `MixingLengthModel` for EDMF, given:\n - `param_set`, an AbstractEarthParameterSet\n\"\"\"\nfunction MixingLengthModel{FT}(param_set::AbstractEarthParameterSet) where {FT}\n c_d_ = c_d(param_set)\n c_m_ = c_m(param_set)\n c_b_ = c_b(param_set)\n a1_ = a1(param_set)\n a2_ = a2(param_set)\n κ = von_karman_const(param_set)\n ω_pr_ = ω_pr(param_set)\n Pr_n_ = Pr_n(param_set)\n Ri_c_ = Ri_c(param_set)\n smin_ub_ = smin_ub(param_set)\n smin_rm_ = smin_rm(param_set)\n max_length = 1e6\n random_minval = 1e-9\n\n args = (\n c_d_,\n c_m_,\n c_b_,\n a1_,\n a2_,\n κ,\n ω_pr_,\n Pr_n_,\n Ri_c_,\n smin_ub_,\n smin_rm_,\n max_length,\n random_minval,\n )\n\n return MixingLengthModel{FT}(args...)\nend\n\nabstract type AbstractStatisticalModel end\nstruct SubdomainMean <: AbstractStatisticalModel end\nstruct GaussQuad <: AbstractStatisticalModel end\nstruct LogNormalQuad <: AbstractStatisticalModel end\n\n\"\"\"\n MicrophysicsModel\n\nA microphysics model for EDMF, containing\nall related model and free parameters and\nassumed subdomain distributions.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct MicrophysicsModel{FT <: AbstractFloat, SM}\n \"Subdomain statistical model\"\n statistical_model::SM\nend\n\n\"\"\"\n MicrophysicsModel(\n FT;\n statistical_model = SubdomainMean()\n )\n\nConstructor for `MicrophysicsModel` for EDMF, given:\n - `FT`, the float type used\n - `statistical_model`, the assumed environmental distribution\n of thermodynamic variables.\n\"\"\"\nfunction MicrophysicsModel(FT; statistical_model = SubdomainMean())\n args = (statistical_model,)\n return MicrophysicsModel{FT, typeof(statistical_model)}(args...)\nend\n\n\"\"\"\n Environment <: BalanceLaw\nA `BalanceLaw` for the environment subdomain arising in EDMF.\n\"\"\"\nBase.@kwdef struct Environment{FT <: AbstractFloat, N_quad} <: BalanceLaw end\n\n\"\"\"\n Updraft <: BalanceLaw\nA `BalanceLaw` for the updraft subdomains arising in EDMF.\n\"\"\"\nBase.@kwdef struct Updraft{FT <: AbstractFloat} <: BalanceLaw end\n\nabstract type Coupling end\n\n\"\"\"\n Decoupled <: Coupling\n\nDispatch on decoupled model (default)\n\n - The EDMF SGS tendencies do not modify the grid-mean equations.\n\"\"\"\nstruct Decoupled <: Coupling end\n\n\"\"\"\n Coupled <: Coupling\n\nDispatch on coupled model\n\n - The EDMF SGS tendencies modify the grid-mean equations.\n\"\"\"\nstruct Coupled <: Coupling end\n\n\"\"\"\n EDMF <: TurbulenceConvectionModel\n\nA turbulence convection model for the EDMF\nscheme, containing all closure models and\nfree parameters.\n\n# Fields\n$(DocStringExtensions.FIELDS)\n\"\"\"\nBase.@kwdef struct EDMF{\n FT <: AbstractFloat,\n N,\n UP,\n EN,\n ED,\n P,\n S,\n MP,\n ML,\n SD,\n C,\n} <: TurbulenceConvectionModel\n \"Updrafts\"\n updraft::UP\n \"Environment\"\n environment::EN\n \"Entrainment-Detrainment model\"\n entr_detr::ED\n \"Pressure model\"\n pressure::P\n \"Surface model\"\n surface::S\n \"Microphysics model\"\n micro_phys::MP\n \"Mixing length model\"\n mix_len::ML\n \"Subdomain model\"\n subdomains::SD\n \"Coupling mode\"\n coupling::C\nend\n\n\"\"\"\n EDMF(\n FT, N_up, N_quad, param_set;\n updraft = ntuple(i -> Updraft{FT}(), N_up),\n environment = Environment{FT, N_quad}(),\n entr_detr = EntrainmentDetrainment{FT}(param_set),\n pressure = PressureModel{FT}(param_set),\n surface = SurfaceModel{FT}(N_up, param_set),\n micro_phys = MicrophysicsModel(FT),\n mix_len = MixingLengthModel{FT}(param_set),\n subdomain = SubdomainModel(FT, N_up),\n )\nConstructor for `EDMF` subgrid-scale scheme, given:\n - `FT`, the AbstractFloat type used\n - `N_up`, the number of updrafts\n - `N_quad`, the quadrature order. `N_quad^2` is\n the total number of quadrature points\n used for environmental distributions.\n - `updraft`, a tuple containing N_up updraft BalanceLaws\n - `environment`, the environment BalanceLaw\n - `entr_detr`, an `EntrainmentDetrainment` model\n - `pressure`, a `PressureModel`\n - `surface`, a `SurfaceModel`\n - `micro_phys`, a `MicrophysicsModel`\n - `mix_len`, a `MixingLengthModel`\n - `subdomain`, a `SubdomainModel`\n - `coupling`, a coupling type\n\"\"\"\nfunction EDMF(\n FT,\n N_up,\n N_quad,\n param_set;\n updraft = ntuple(i -> Updraft{FT}(), N_up),\n environment = Environment{FT, N_quad}(),\n entr_detr = EntrainmentDetrainment{FT}(param_set),\n pressure = PressureModel{FT}(param_set),\n surface = SurfaceModel{FT}(N_up, param_set),\n micro_phys = MicrophysicsModel(FT),\n mix_len = MixingLengthModel{FT}(param_set),\n subdomain = SubdomainModel(FT, N_up),\n coupling = Decoupled(),\n)\n args = (\n updraft,\n environment,\n entr_detr,\n pressure,\n surface,\n micro_phys,\n mix_len,\n subdomain,\n coupling,\n )\n return EDMF{FT, N_up, typeof.(args)...}(args...)\nend\n\n\nimport ClimateMachine.TurbulenceConvection: turbconv_sources, turbconv_bcs\n\n\"\"\"\n EDMFTopBC <: TurbConvBC\n\nBoundary conditions for the top of the EDMF.\n\"\"\"\nstruct EDMFTopBC <: TurbConvBC end\n\n\"\"\"\n EDMFBottomBC <: TurbConvBC\n\nBoundary conditions for the bottom of the EDMF.\n\"\"\"\nstruct EDMFBottomBC <: TurbConvBC end\n\n\nn_updrafts(m::EDMF{FT, N_up}) where {FT, N_up} = N_up\nn_updrafts(m::TurbulenceConvectionModel) = 0\nturbconv_filters(m::TurbulenceConvectionModel) = ()\nturbconv_filters(m::EDMF) = (\n \"turbconv.environment.ρatke\",\n \"turbconv.environment.ρaθ_liq_cv\",\n \"turbconv.environment.ρaq_tot_cv\",\n \"turbconv.updraft\",\n)\nn_quad_points(m::Environment{FT, N_quad}) where {FT, N_quad} = N_quad\nturbconv_sources(m::EDMF) = ()\nturbconv_bcs(::EDMF) = (EDMFBottomBC(), EDMFTopBC())\n" }, { "path": "test/Atmos/EDMF/ekman_layer.jl", "content": "#!/usr/bin/env julia --project\n#=\n# This driver file simulates the ekman_layer_model.jl in a single column setting.\n# \n# The user may select in main() the following configurations:\n# - DG or FV vertical discretization by changing the boolean `finite_volume`,\n# - Compressible() or Anelastic1D() changing the compressibility,\n# - Constant kinematic viscosity, Smagorinsky-Lilly or EDMF SGS fluxes.\n#\n# The default is DG, Anelastic1D(), constant kinematic viscosity of 0.1. \n#\n=#\n\nusing JLD2, FileIO\nusing ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nusing ClimateMachine.Atmos\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\nimport CLIMAParameters\nimport ClimateMachine.DGMethods.FVReconstructions: FVLinear\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"ekman_layer_model.jl\"))\ninclude(joinpath(\"helper_funcs\", \"diagnostics_configuration.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\nCLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 290.0\nCLIMAParameters.Atmos.EDMF.a_surf(::EarthParameterSet) = 0.0\nfunction set_clima_parameters(filename)\n eval(:(include($filename)))\nend\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables if turbconv=EDMF(...) is selected.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n z = altitude(m, aux)\n\n param_set = parameter_set(m)\n ts = new_thermo_state(m, state, aux)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = FT(0)\n end\n\n en.ρatke =\n z > FT(250) ? FT(0) :\n gm.ρ *\n FT(0.3375) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n en.ρaθ_liq_cv = FT(0)\n en.ρaq_tot_cv = FT(0)\n en.ρaθ_liq_q_tot_cv = FT(0)\n return nothing\nend;\n\nfunction main(::Type{FT}, cl_args) where {FT}\n # Change boolean to control vertical discretization type\n finite_volume = false\n\n # Choice of compressibility and CFL\n # compressibility = Compressible()\n compressibility = Anelastic1D()\n str_comp = compressibility == Compressible() ? \"COMPRESS\" : \"ANELASTIC\"\n\n # Choice of SGS model\n # turbconv = NoTurbConv()\n N_updrafts = 1\n N_quad = 3\n turbconv = EDMF(\n FT,\n N_updrafts,\n N_quad,\n param_set,\n surface = NeutralDrySurfaceModel{FT}(param_set),\n )\n\n C_smag_ = C_smag(param_set)\n turbulence = ConstantKinematicViscosity(FT(0.1))\n # turbulence = SmagorinskyLilly{FT}(C_smag_)\n\n # Prescribe domain parameters\n zmax = FT(400)\n # Simulation time\n t0 = FT(0)\n timeend = FT(3600 * 2) # Change to 7h for low-level jet\n CFLmax = compressibility == Compressible() ? FT(1) : FT(100)\n\n config_type = SingleStackConfigType\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n if finite_volume\n N = (1, 0)\n nelem_vert = 80\n ref_state = HydrostaticState(\n DryAdiabaticProfile{FT}(param_set),\n ;\n subtract_off = false,\n )\n output_prefix = string(\"EL_\", str_comp, \"_FVM\")\n fv_reconstruction = FVLinear()\n else\n N = 4\n nelem_vert = 20\n ref_state = HydrostaticState(DryAdiabaticProfile{FT}(param_set),)\n output_prefix = string(\"EL_\", str_comp, \"_DG\")\n fv_reconstruction = nothing\n end\n\n surface_flux = cl_args[\"surface_flux\"]\n model = ekman_layer_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbulence = turbulence,\n turbconv = turbconv,\n ref_state = ref_state,\n compressibility = compressibility,\n )\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n output_prefix,\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n fv_reconstruction = fv_reconstruction,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config = config_diagnostics(driver_config)\n\n # boyd vandeven filter\n cb_boyd = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n AtmosSpecificFilterPerturbations(driver_config.bl),\n solver_config.dg.grid,\n BoydVandevenFilter(solver_config.dg.grid, 1, 4);\n state_auxiliary = solver_config.dg.state_auxiliary,\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n # Mass tendencies = 0 for Anelastic1D model,\n # so mass should be completely conserved:\n Δρ_lim = compressibility == Compressible() ? FT(0.001) : FT(0.00000001)\n check_cons = (ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", Δρ_lim),)\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n if !isnothing(cl_args[\"cparam_file\"])\n ClimateMachine.Settings.output_dir = cl_args[\"cparam_file\"] * \".output\"\n end\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cb_boyd, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n# ArgParse in global scope to modify Clima Parameters\nsl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sl_args, \"EkmanLayer\")\n@add_arg_table! sl_args begin\n \"--cparam-file\"\n help = \"specify CLIMAParameters file\"\n arg_type = Union{String, Nothing}\n default = nothing\n\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sl\"\n arg_type = String\n default = \"prescribed\"\nend\n\ncl_args = ClimateMachine.init(parse_clargs = true, custom_clargs = sl_args)\nif !isnothing(cl_args[\"cparam_file\"])\n filename = cl_args[\"cparam_file\"]\n set_clima_parameters(filename)\nend\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\n## Uncomment lines to save output using JLD2\n# output_dir = @__DIR__;\n# mkpath(output_dir);\n# function dons(diag_vs_z)\n# return Dict(map(keys(first(diag_vs_z))) do k\n# string(k) => [getproperty(ca, k) for ca in diag_vs_z]\n# end)\n# end\n# get_dons_arr(diag_arr) = [dons(diag_vs_z) for diag_vs_z in diag_arr]\n# dons_arr = get_dons_arr(diag_arr)\n# println(dons_arr[1].keys)\n# z = get_z(solver_config.dg.grid; rm_dupes = true);\n# save(\n# string(output_dir, \"/ekman.jld2\"),\n# \"dons_arr\",\n# dons_arr,\n# \"time_data\",\n# time_data,\n# \"z\",\n# z,\n# )\n\nnothing\n" }, { "path": "test/Atmos/EDMF/helper_funcs/diagnose_environment.jl", "content": "#### Diagnose environment variables\n\n\"\"\"\n environment_vars(state::Vars, N_up::Int)\n\nA NamedTuple of environment variables\n\"\"\"\nfunction environment_vars(state::Vars, N_up::Int)\n return (a = environment_area(state, N_up), w = environment_w(state, N_up))\nend\n\n\"\"\"\n environment_area(\n state::Vars,\n N_up::Int,\n )\n\nReturns the environmental area fraction, given:\n - `state`, state variables\n - `N_up`, number of updrafts\n\"\"\"\nfunction environment_area(state::Vars, N_up::Int)\n up = state.turbconv.updraft\n return 1 - sum(vuntuple(i -> up[i].ρa, N_up)) / state.ρ\nend\n\n\"\"\"\n environment_w(state::Vars, N_up::Int)\n\nReturns the environmental vertical velocity, given:\n - `state`, state variables\n - `N_up`, number of updrafts\n\"\"\"\nfunction environment_w(state::Vars, N_up::Int)\n ρ_inv = 1 / state.ρ\n a_en = environment_area(state, N_up)\n up = state.turbconv.updraft\n return (state.ρu[3] - sum(vuntuple(i -> up[i].ρaw, N_up))) / a_en * ρ_inv\nend\n\n\"\"\"\n grid_mean_b(env, a_up, N_up::Int, buoyancy_up, buoyancy_en)\n\nReturns the grid-mean buoyancy with respect to the\nreference state, given:\n - `env`, environment variables\n - `a_up`, updraft area fractions\n - `N_up`, number of updrafts\n - `buoyancy_up`, updraft buoyancies\n - `buoyancy_en`, environment buoyancy\n\"\"\"\nfunction grid_mean_b(env, a_up, N_up::Int, buoyancy_up, buoyancy_en)\n ∑abuoyancy_up = sum(vuntuple(i -> buoyancy_up[i] * a_up[i], N_up))\n return env.a * buoyancy_en + ∑abuoyancy_up\nend\n" }, { "path": "test/Atmos/EDMF/helper_funcs/diagnostics_configuration.jl", "content": "\"\"\"\n config_diagnostics(driver_config, timeend; interval=nothing)\n\nReturns the state and tendency diagnostic groups\n\"\"\"\nfunction config_diagnostics(driver_config, timeend; interval = nothing)\n FT = eltype(driver_config.grid)\n info = driver_config.config_info\n if interval == nothing\n interval = \"$(cld(timeend, 2) + 10)ssecs\"\n #interval = \"10steps\"\n end\n\n boundaries = [\n FT(0) FT(0) FT(0)\n FT(info.hmax) FT(info.hmax) FT(info.zmax)\n ]\n axes = (\n [FT(1)],\n [FT(1)],\n collect(range(boundaries[1, 3], boundaries[2, 3], step = FT(50)),),\n )\n interpol = ClimateMachine.InterpolationConfiguration(\n driver_config,\n boundaries;\n axes = axes,\n )\n ds_dgngrp = setup_dump_state_diagnostics(\n SingleStackConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n )\n dt_dgngrp = setup_dump_tendencies_diagnostics(\n SingleStackConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n )\n return ClimateMachine.DiagnosticsConfiguration([ds_dgngrp, dt_dgngrp])\nend\n" }, { "path": "test/Atmos/EDMF/helper_funcs/lamb_smooth_minimum.jl", "content": "using LambertW\n\n\"\"\"\n lambertw_gpu(N)\n\nReturns `real(LambertW.lambertw(Float64(N - 1) / MathConstants.e))`,\nvalid for `N = 2` and `N = 3`.\n\nTODO: add `LambertW.lambertw` support to KernelAbstractions.\n\"\"\"\nfunction lambertw_gpu(N)\n if !(N == 2 || N == 3)\n error(\"Bad N in lambertw_gpu\")\n end\n return (0.2784645427610738, 0.46305551336554884)[N - 1]\nend\n\n\"\"\"\n lamb_smooth_minimum(\n l::AbstractArray{FT};\n frac_upper_bound::FT,\n reg_min::FT,\n ) where {FT}\n\nReturns the smooth minimum of the elements of an array\nfollowing the formulation of\nLopez-Gomez et al. (JAMES, 2020), Appendix A, given:\n - `l`, an array of candidate elements\n - `frac_upper_bound`, defines the upper bound of the\n smooth minimum as `smin(x) = min(x)*(1+frac_upper_bound)`\n - `reg_min`, defines the minimum value of the regularizer Λ\n\"\"\"\nfunction lamb_smooth_minimum(\n l::AbstractArray{FT},\n frac_upper_bound::FT,\n reg_min::FT,\n) where {FT}\n\n xmin = minimum(l)\n\n # Get regularizer for exponential weights\n N_l = length(l)\n denom = FT(lambertw_gpu(N_l))\n Λ = max(FT(xmin) * frac_upper_bound / denom, reg_min)\n\n num = sum(i -> l[i] * exp(-(l[i] - xmin) / Λ), 1:N_l)\n den = sum(i -> exp(-(l[i] - xmin) / Λ), 1:N_l)\n smin = num / den\n\n return smin\nend\n" }, { "path": "test/Atmos/EDMF/helper_funcs/nondimensional_exchange_functions.jl", "content": "\"\"\"\n nondimensional_exchange_functions(\n m::AtmosModel{FT},\n entr::EntrainmentDetrainment,\n state::Vars,\n aux::Vars,\n t::Real,\n ts_up,\n ts_en,\n env,\n buoy,\n i,\n ) where {FT}\n\nReturns the nondimensional entrainment and detrainment\nfunctions following Cohen et al. (JAMES, 2020), given:\n - `m`, an `AtmosModel`\n - `entr`, an `EntrainmentDetrainment` model\n - `state`, state variables\n - `aux`, auxiliary variables\n - `ts_up`, updraft thermodynamic states\n - `ts_en`, environment thermodynamic states\n - `env`, NamedTuple of environment variables\n - `buoy`, NamedTuple of environment and updraft buoyancies\n - `i`, the updraft index\n\"\"\"\nfunction nondimensional_exchange_functions(\n m::AtmosModel{FT},\n entr::EntrainmentDetrainment,\n state::Vars,\n aux::Vars,\n ts_up,\n ts_en,\n env,\n buoy,\n i,\n) where {FT}\n\n # Alias convention:\n gm = state\n up = state.turbconv.updraft\n up_aux = aux.turbconv.updraft\n en_aux = aux.turbconv.environment\n\n # precompute vars\n w_min = entr.w_min\n N_up = n_updrafts(turbconv_model(m))\n ρinv = 1 / gm.ρ\n a_up_i = up[i].ρa * ρinv\n w_up_i = fix_void_up(up[i].ρa, up[i].ρaw / up[i].ρa)\n\n # thermodynamic variables\n RH_up = relative_humidity(ts_up[i])\n RH_en = relative_humidity(ts_en)\n\n Δw = filter_w(w_up_i - env.w, w_min)\n Δb = buoy.up[i] - buoy.en\n\n c_δ = sign(condensate(ts_en) + condensate(ts_up[i])) * entr.c_δ\n\n # compute dry and moist aux functions\n μ_ij = (entr.χ - a_up_i / (a_up_i + env.a)) * Δb / Δw\n D_ε = entr.c_ε / (1 + exp(-μ_ij / entr.μ_0))\n M_ε = c_δ * (max((RH_en^entr.β - RH_up^entr.β), 0))^(1 / entr.β)\n D_δ = entr.c_ε / (1 + exp(μ_ij / entr.μ_0))\n M_δ = c_δ * (max((RH_up^entr.β - RH_en^entr.β), 0))^(1 / entr.β)\n return D_ε, D_δ, M_δ, M_ε\nend;\n" }, { "path": "test/Atmos/EDMF/helper_funcs/save_subdomain_temperature.jl", "content": "# Convenience wrapper\nsave_subdomain_temperature!(m, state, aux) =\n save_subdomain_temperature!(m, moisture_model(m), state, aux)\n\nusing KernelAbstractions: @print\n\n\"\"\"\n save_subdomain_temperature!(\n m::AtmosModel,\n moist::EquilMoist,\n state::Vars,\n aux::Vars,\n )\n\nUpdates the subdomain sensible temperature, given:\n - `m`, an `AtmosModel`\n - `moist`, an `EquilMoist` model\n - `state`, state variables\n - `aux`, auxiliary variables\n\"\"\"\nfunction save_subdomain_temperature!(\n m::AtmosModel,\n moist::EquilMoist,\n state::Vars,\n aux::Vars,\n)\n N_up = n_updrafts(turbconv_model(m))\n ts = recover_thermo_state(m, state, aux)\n ts_up = new_thermo_state_up(m, state, aux, ts)\n ts_en = new_thermo_state_en(m, state, aux, ts)\n\n @unroll_map(N_up) do i\n aux.turbconv.updraft[i].T = air_temperature(ts_up[i])\n end\n aux.turbconv.environment.T = air_temperature(ts_en)\n return nothing\nend\n\n# No need to save temperature for DryModel.\nfunction save_subdomain_temperature!(\n m::AtmosModel,\n moist::DryModel,\n state::Vars,\n aux::Vars,\n) end\n" }, { "path": "test/Atmos/EDMF/helper_funcs/subdomain_statistics.jl", "content": "#### Subdomain statistics\n\nfunction compute_subdomain_statistics(m::AtmosModel, args, ts_gm, ts_en)\n turbconv = turbconv_model(m)\n return compute_subdomain_statistics(\n turbconv.micro_phys.statistical_model,\n m,\n args,\n ts_gm,\n ts_en,\n )\nend\n\n\"\"\"\n compute_subdomain_statistics(\n statistical_model::SubdomainMean,\n m::AtmosModel{FT},\n args,\n ts_gm,\n ts_en,\n ) where {FT}\n\nReturns a cloud fraction and cloudy and dry thermodynamic\nstates in the subdomain.\n\"\"\"\nfunction compute_subdomain_statistics(\n statistical_model::SubdomainMean,\n m::AtmosModel{FT},\n args,\n ts_gm,\n ts_en,\n) where {FT}\n cloud_frac = has_condensate(ts_en) ? FT(1) : FT(0)\n dry = ts_en\n cloudy = ts_en\n return (dry = dry, cloudy = cloudy, cloud_frac = cloud_frac)\nend\n" }, { "path": "test/Atmos/EDMF/helper_funcs/subdomain_thermo_states.jl", "content": "#### thermo states for subdomains\n\nusing KernelAbstractions: @print\n\nexport new_thermo_state_up,\n new_thermo_state_en,\n recover_thermo_state_all,\n recover_thermo_state_up,\n recover_thermo_state_en\n\n####\n#### Interface\n####\n\n\"\"\"\n new_thermo_state_up(bl, state, aux)\n\nUpdraft thermodynamic states given:\n - `bl`, parent `BalanceLaw`\n - `state`, state variables\n - `aux`, auxiliary variables\n\n!!! note\n This method calls saturation adjustment for\n EquilMoist models.\n\"\"\"\nnew_thermo_state_up(\n bl::AtmosModel,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState = recover_thermo_state(bl, state, aux),\n) = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts)\n\n\"\"\"\n new_thermo_state_en(bl, state, aux)\n\nEnvironment thermodynamic state given:\n - `bl`, parent `BalanceLaw`\n - `state`, state variables\n - `aux`, auxiliary variables\n\n!!! note\n This method calls saturation adjustment for\n EquilMoist models.\n\"\"\"\nnew_thermo_state_en(\n bl::AtmosModel,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState = recover_thermo_state(bl, state, aux),\n) = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts)\n\n\"\"\"\n recover_thermo_state_all(bl, state, aux)\n\nRecover NamedTuple of all thermo states\n\n# TODO: Define/call `recover_thermo_state` when it's safely implemented\n (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)\n\"\"\"\nfunction recover_thermo_state_all(bl, state, aux)\n ts = new_thermo_state(bl, state, aux)\n return (\n gm = ts,\n en = new_thermo_state_en(bl, moisture_model(bl), state, aux, ts),\n up = new_thermo_state_up(bl, moisture_model(bl), state, aux, ts),\n )\nend\n\n\"\"\"\n recover_thermo_state_up(bl, state, aux, ts = new_thermo_state(bl, state, aux))\n\nRecover the updraft thermodynamic states given:\n - `bl`, parent `BalanceLaw`\n - `state`, state variables\n - `aux`, auxiliary variables\n\n!!! warn\n Right now we are directly calling new_thermo_state_up to avoid\n inconsistent aux states in kernels where the aux states are\n out of sync with the boundary state.\n\n# TODO: Define/call `recover_thermo_state` when it's safely implemented\n (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)\n\"\"\"\nfunction recover_thermo_state_up(\n bl,\n state,\n aux,\n ts = new_thermo_state(bl, state, aux),\n)\n return new_thermo_state_up(bl, moisture_model(bl), state, aux, ts)\nend\n\n\"\"\"\n recover_thermo_state_en(bl, state, aux, ts = recover_thermo_state(bl, state, aux))\n\nRecover the environment thermodynamic state given:\n - `bl`, parent `BalanceLaw`\n - `state`, state variables\n - `aux`, auxiliary variables\n\n!!! warn\n Right now we are directly calling new_thermo_state_up to avoid\n inconsistent aux states in kernels where the aux states are\n out of sync with the boundary state.\n\n# TODO: Define/call `recover_thermo_state` when it's safely implemented\n (see https://github.com/CliMA/ClimateMachine.jl/issues/1648)\n\"\"\"\nfunction recover_thermo_state_en(\n bl,\n state,\n aux,\n ts = new_thermo_state(bl, state, aux),\n)\n return new_thermo_state_en(bl, moisture_model(bl), state, aux, ts)\nend\n\n####\n#### Implementation\n####\n\nfunction new_thermo_state_up(\n m::AtmosModel{FT},\n moist::DryModel,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState,\n) where {FT}\n N_up = n_updrafts(turbconv_model(m))\n up = state.turbconv.updraft\n p = air_pressure(ts)\n param_set = parameter_set(m)\n # compute thermo state for updrafts\n ts_up = vuntuple(N_up) do i\n ρa_up = up[i].ρa\n ρaθ_liq_up = up[i].ρaθ_liq\n θ_liq_up =\n fix_void_up(ρa_up, ρaθ_liq_up / ρa_up, liquid_ice_pottemp(ts))\n PhaseDry_pθ(param_set, p, θ_liq_up)\n end\n return ts_up\nend\n\nfunction new_thermo_state_up(\n m::AtmosModel{FT},\n moist::EquilMoist,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState,\n) where {FT}\n\n N_up = n_updrafts(turbconv_model(m))\n up = state.turbconv.updraft\n p = air_pressure(ts)\n param_set = parameter_set(m)\n # compute thermo state for updrafts\n ts_up = vuntuple(N_up) do i\n ρa_up = up[i].ρa\n ρaθ_liq_up = up[i].ρaθ_liq\n ρaq_tot_up = up[i].ρaq_tot\n θ_liq_up =\n fix_void_up(ρa_up, ρaθ_liq_up / ρa_up, liquid_ice_pottemp(ts))\n q_tot_up = fix_void_up(\n ρa_up,\n ρaq_tot_up / ρa_up,\n total_specific_humidity(ts),\n )\n PhaseEquil_pθq(param_set, p, θ_liq_up, q_tot_up)\n end\n return ts_up\nend\n\nfunction new_thermo_state_en(\n m::AtmosModel,\n moist::DryModel,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState,\n)\n turbconv = turbconv_model(m)\n N_up = n_updrafts(turbconv)\n up = state.turbconv.updraft\n\n # diagnose environment thermo state\n ρ_inv = 1 / state.ρ\n p = air_pressure(ts)\n θ_liq = liquid_ice_pottemp(ts)\n a_en = environment_area(state, N_up)\n ρaθ_liq_up = vuntuple(N_up) do i\n fix_void_up(up[i].ρa, up[i].ρaθ_liq)\n end\n θ_liq_en = (θ_liq - sum(vuntuple(j -> ρaθ_liq_up[j] * ρ_inv, N_up))) / a_en\n a_min = turbconv.subdomains.a_min\n a_max = turbconv.subdomains.a_max\n param_set = parameter_set(m)\n if !(0 <= θ_liq_en)\n @print(\"ρaθ_liq_up = \", ρaθ_liq_up[Val(1)], \"\\n\")\n @print(\"θ_liq = \", θ_liq, \"\\n\")\n @print(\"θ_liq_en = \", θ_liq_en, \"\\n\")\n error(\"Environment θ_liq_en out-of-bounds in new_thermo_state_en\")\n end\n ts_en = PhaseDry_pθ(param_set, p, θ_liq_en)\n return ts_en\nend\n\nfunction new_thermo_state_en(\n m::AtmosModel,\n moist::EquilMoist,\n state::Vars,\n aux::Vars,\n ts::ThermodynamicState,\n)\n turbconv = turbconv_model(m)\n N_up = n_updrafts(turbconv)\n up = state.turbconv.updraft\n\n # diagnose environment thermo state\n ρ_inv = 1 / state.ρ\n p = air_pressure(ts)\n θ_liq = liquid_ice_pottemp(ts)\n q_tot = total_specific_humidity(ts)\n a_en = environment_area(state, N_up)\n θ_liq_en = (θ_liq - sum(vuntuple(j -> up[j].ρaθ_liq * ρ_inv, N_up))) / a_en\n q_tot_en = (q_tot - sum(vuntuple(j -> up[j].ρaq_tot * ρ_inv, N_up))) / a_en\n a_min = turbconv.subdomains.a_min\n a_max = turbconv.subdomains.a_max\n param_set = parameter_set(m)\n if !(0 <= θ_liq_en)\n @print(\"θ_liq_en = \", θ_liq_en, \"\\n\")\n error(\"Environment θ_liq_en out-of-bounds in new_thermo_state_en\")\n end\n if !(0 <= q_tot_en <= 1)\n @print(\"q_tot_en = \", q_tot_en, \"\\n\")\n error(\"Environment q_tot_en out-of-bounds in new_thermo_state_en\")\n end\n ts_en = PhaseEquil_pθq(param_set, p, θ_liq_en, q_tot_en)\n return ts_en\nend\n" }, { "path": "test/Atmos/EDMF/helper_funcs/utility_funcs.jl", "content": "\"\"\"\n filter_w(w::FT, w_min::FT) where {FT}\n\nReturn velocity such that\n`abs(filter_w(w, w_min)) >= abs(w_min)`\nwhile preserving the sign of `w`.\n\"\"\"\nfilter_w(w::FT, w_min::FT) where {FT} =\n max(abs(w), abs(w_min)) * (w < 0 ? sign(w) : 1)\n\n\"\"\"\n enforce_unit_bounds(a_up_i::FT, a_min::FT, a_max::FT) where {FT}\n\nEnforce variable to be positive.\n\nIdeally, this safety net will be removed once we\nhave robust positivity preserving methods. For now,\nwe need this to avoid domain error in certain\ncircumstances.\n\"\"\"\nenforce_unit_bounds(a_up_i::FT, a_min::FT, a_max::FT) where {FT} =\n clamp(a_up_i, a_min, a_max)\n\n\"\"\"\n enforce_positivity(x::FT) where {FT}\n\nEnforce variable to be positive.\n\nIdeally, this safety net will be removed once we\nhave robust positivity preserving methods. For now,\nwe need this to avoid domain error in certain\ncircumstances.\n\"\"\"\nenforce_positivity(x::FT) where {FT} = max(x, FT(0))\n\n\"\"\"\n fix_void_up(ρa_up_i::FT, val::FT, fallback = FT(0)) where {FT}\n\nSubstitute value by a consistent fallback in case of\nnegligible area fraction (void updraft).\n\"\"\"\nfunction fix_void_up(ρa_up_i::FT, val::FT, fallback = FT(0)) where {FT}\n tol = sqrt(eps(FT))\n return ρa_up_i > tol ? val : fallback\nend\n" }, { "path": "test/Atmos/EDMF/report_mse_bomex.jl", "content": "using ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n plot_dir = joinpath(clima_dir, \"output\", \"bomex_edmf\", \"pycles_comparison\")\nelse\n plot_dir = nothing\nend\n\ninclude(joinpath(@__DIR__, \"compute_mse.jl\"))\n\ndata_file = Dataset(joinpath(PyCLES_output_dataset_path, \"Bomex.nc\"), \"r\")\n\n#! format: off\nbest_mse = OrderedDict()\nbest_mse[\"prog_ρ\"] = 3.4936203419421122e-02\nbest_mse[\"prog_ρu_1\"] = 3.0715584660655654e+03\nbest_mse[\"prog_ρu_2\"] = 1.2897432819935302e-03\nbest_mse[\"prog_moisture_ρq_tot\"] = 4.1572900459802775e-02\nbest_mse[\"prog_turbconv_environment_ρatke\"] = 8.5590220998989253e+02\nbest_mse[\"prog_turbconv_environment_ρaθ_liq_cv\"] = 8.5667227621106434e+01\nbest_mse[\"prog_turbconv_environment_ρaq_tot_cv\"] = 1.6437582130429374e+02\nbest_mse[\"prog_turbconv_updraft_1_ρa\"] = 8.0846471185616252e+01\nbest_mse[\"prog_turbconv_updraft_1_ρaw\"] = 8.5071186636845333e-02\nbest_mse[\"prog_turbconv_updraft_1_ρaθ_liq\"] = 9.0386637425608303e+00\nbest_mse[\"prog_turbconv_updraft_1_ρaq_tot\"] = 1.0803377515313473e+01\n#! format: on\n\ncomputed_mse = compute_mse(\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n time_data,\n dons_arr,\n data_file,\n \"Bomex\",\n best_mse,\n 400,\n plot_dir,\n)\n\n@testset \"BOMEX EDMF Solution Quality Assurance (QA) tests\" begin\n #! format: off\n test_mse(computed_mse, best_mse, \"prog_ρ\")\n test_mse(computed_mse, best_mse, \"prog_ρu_1\")\n test_mse(computed_mse, best_mse, \"prog_ρu_2\")\n test_mse(computed_mse, best_mse, \"prog_moisture_ρq_tot\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρatke\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρaθ_liq_cv\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρaq_tot_cv\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρa\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaw\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaθ_liq\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaq_tot\")\n #! format: on\nend\n" }, { "path": "test/Atmos/EDMF/report_mse_sbl_anelastic.jl", "content": "using ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n plot_dir = joinpath(clima_dir, \"output\", \"sbl_edmf\", \"pycles_comparison\")\nelse\n plot_dir = nothing\nend\n\ninclude(joinpath(@__DIR__, \"compute_mse.jl\"))\n\ndata_file = Dataset(joinpath(PyCLES_output_dataset_path, \"Gabls.nc\"), \"r\")\n\n#! format: off\nbest_mse = OrderedDict()\nbest_mse[\"prog_ρ\"] = 9.3809207150296822e-03\nbest_mse[\"prog_ρu_1\"] = 6.7269975116623837e+03\nbest_mse[\"prog_ρu_2\"] = 6.8630628605220889e-01\n#! format: on\n\ncomputed_mse = compute_mse(\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n time_data,\n dons_arr,\n data_file,\n \"Gabls\",\n best_mse,\n 1800,\n plot_dir,\n)\n\n@testset \"SBL Anelastic Solution Quality Assurance (QA) tests\" begin\n #! format: off\n test_mse(computed_mse, best_mse, \"prog_ρ\")\n test_mse(computed_mse, best_mse, \"prog_ρu_1\")\n test_mse(computed_mse, best_mse, \"prog_ρu_2\")\n #! format: on\nend\n" }, { "path": "test/Atmos/EDMF/report_mse_sbl_coupled_edmf_an1d.jl", "content": "using ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n plot_dir = joinpath(clima_dir, \"output\", \"sbl_edmf\", \"pycles_comparison\")\nelse\n plot_dir = nothing\nend\n\ninclude(joinpath(@__DIR__, \"compute_mse.jl\"))\n\ndata_file = Dataset(joinpath(PyCLES_output_dataset_path, \"Gabls.nc\"), \"r\")\n\n#! format: off\nbest_mse = OrderedDict()\nbest_mse[\"prog_ρ\"] = 9.3808142287632006e-03\nbest_mse[\"prog_ρu_1\"] = 6.7427140307178524e+03\nbest_mse[\"prog_ρu_2\"] = 7.2306841961112966e-01\nbest_mse[\"prog_turbconv_environment_ρatke\"] = 2.9810806322235749e+02\nbest_mse[\"prog_turbconv_environment_ρaθ_liq_cv\"] = 8.1270487249851797e+01\nbest_mse[\"prog_turbconv_updraft_1_ρa\"] = 2.7223017351724246e+02\nbest_mse[\"prog_turbconv_updraft_1_ρaw\"] = 5.5909371368198686e+02\nbest_mse[\"prog_turbconv_updraft_1_ρaθ_liq\"] = 2.7933498788454813e+02\n#! format: on\n\ncomputed_mse = compute_mse(\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n time_data,\n dons_arr,\n data_file,\n \"Gabls\",\n best_mse,\n 1800,\n plot_dir,\n)\n\n@testset \"SBL Coupled EDMF Solution Quality Assurance (QA) tests\" begin\n #! format: off\n test_mse(computed_mse, best_mse, \"prog_ρ\")\n test_mse(computed_mse, best_mse, \"prog_ρu_1\")\n test_mse(computed_mse, best_mse, \"prog_ρu_2\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρatke\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρaθ_liq_cv\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρa\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaw\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaθ_liq\")\n #! format: on\nend\n" }, { "path": "test/Atmos/EDMF/report_mse_sbl_edmf.jl", "content": "using ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n plot_dir = joinpath(clima_dir, \"output\", \"sbl_edmf\", \"pycles_comparison\")\nelse\n plot_dir = nothing\nend\n\ninclude(joinpath(@__DIR__, \"compute_mse.jl\"))\n\ndata_file = Dataset(joinpath(PyCLES_output_dataset_path, \"Gabls.nc\"), \"r\")\n\n#! format: off\nbest_mse = OrderedDict()\nbest_mse[\"prog_ρ\"] = 6.8041561724036673e-03\nbest_mse[\"prog_ρu_1\"] = 6.2586216904022822e+03\nbest_mse[\"prog_ρu_2\"] = 1.2965826326450741e-04\nbest_mse[\"prog_turbconv_environment_ρatke\"] = 4.9035225536000081e+02\nbest_mse[\"prog_turbconv_environment_ρaθ_liq_cv\"] = 8.7727377301948991e+01\nbest_mse[\"prog_turbconv_updraft_1_ρa\"] = 1.8213581913998944e+01\nbest_mse[\"prog_turbconv_updraft_1_ρaw\"] = 1.7800899665237452e-01\nbest_mse[\"prog_turbconv_updraft_1_ρaθ_liq\"] = 1.3358308964295857e+01\n#! format: on\n\ncomputed_mse = compute_mse(\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n time_data,\n dons_arr,\n data_file,\n \"Gabls\",\n best_mse,\n 60,\n plot_dir,\n)\n\n@testset \"SBL EDMF Solution Quality Assurance (QA) tests\" begin\n #! format: off\n test_mse(computed_mse, best_mse, \"prog_ρ\")\n test_mse(computed_mse, best_mse, \"prog_ρu_1\")\n test_mse(computed_mse, best_mse, \"prog_ρu_2\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρatke\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_environment_ρaθ_liq_cv\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρa\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaw\")\n test_mse(computed_mse, best_mse, \"prog_turbconv_updraft_1_ρaθ_liq\")\n #! format: on\nend\n" }, { "path": "test/Atmos/EDMF/report_mse_sbl_ss_implicit.jl", "content": "using ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\nif parse(Bool, get(ENV, \"CLIMATEMACHINE_PLOT_EDMF_COMPARISON\", \"false\"))\n plot_dir = joinpath(clima_dir, \"output\", \"sbl_edmf\", \"pycles_comparison\")\nelse\n plot_dir = nothing\nend\n\ninclude(joinpath(@__DIR__, \"compute_mse.jl\"))\n\ndata_file = Dataset(joinpath(PyCLES_output_dataset_path, \"Gabls.nc\"), \"r\")\n\n#! format: off\nbest_mse = OrderedDict()\nbest_mse[\"prog_ρ\"] = 7.9878085788692155e-03\nbest_mse[\"prog_ρu_1\"] = 3.0923908085978383e+03\nbest_mse[\"prog_ρu_2\"] = 4.4228530597867703e+01\n#! format: on\n\ncomputed_mse = compute_mse(\n solver_config.dg.grid,\n solver_config.dg.balance_law,\n time_data,\n dons_arr,\n data_file,\n \"Gabls\",\n best_mse,\n 3600 * 6,\n plot_dir,\n)\n\n@testset \"SBL Implicit Solution Quality Assurance (QA) tests\" begin\n #! format: off\n test_mse(computed_mse, best_mse, \"prog_ρ\")\n test_mse(computed_mse, best_mse, \"prog_ρu_1\")\n test_mse(computed_mse, best_mse, \"prog_ρu_2\")\n #! format: on\nend\n" }, { "path": "test/Atmos/EDMF/stable_bl_anelastic1d.jl", "content": "using JLD2, FileIO\nusing ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nimport ClimateMachine.BalanceLaws: projection\nimport ClimateMachine.DGMethods\nusing ClimateMachine.Atmos\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\nimport CLIMAParameters\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"stable_bl_model.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\n# CLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 290.0 # default\nCLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 265\nCLIMAParameters.Atmos.EDMF.a_surf(::EarthParameterSet) = 0.0\nfunction set_clima_parameters(filename)\n eval(:(include($filename)))\nend\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n # GCM setting - Initialize the grid mean profiles of prognostic variables (ρ,e_int,q_tot,u,v,w)\n z = altitude(m, aux)\n\n # SCM setting - need to have separate cases coded and called from a folder - see what LES does\n # a thermo state is used here to convert the input θ to e_int profile\n e_int = internal_energy(m, state, aux)\n param_set = parameter_set(m)\n ts = PhaseDry(param_set, e_int, state.ρ)\n T = air_temperature(ts)\n p = air_pressure(ts)\n q = PhasePartition(ts)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = FT(0)\n end\n\n # initialize environment covariance with zero for now\n if z <= FT(250)\n en.ρatke =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n en.ρaθ_liq_cv =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n else\n en.ρatke = FT(0)\n en.ρaθ_liq_cv = FT(0)\n end\n en.ρaq_tot_cv = FT(0)\n en.ρaθ_liq_q_tot_cv = FT(0)\n return nothing\nend;\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # Choice of compressibility and CFL\n # compressibility = Compressible()\n compressibility = Anelastic1D()\n str_comp = compressibility == Compressible() ? \"COMPRESS\" : \"ANELASTIC\"\n\n # DG polynomial order\n N = 4\n nelem_vert = 20\n\n # Prescribe domain parameters\n zmax = FT(400)\n # Simulation time\n t0 = FT(0)\n timeend = FT(1800 * 1)\n CFLmax = compressibility == Compressible() ? FT(1) : FT(100)\n\n config_type = SingleStackConfigType\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n # Choice of SGS model\n N_updrafts = 1\n N_quad = 3\n turbconv = NoTurbConv()\n # turbconv = EDMF(\n # FT,\n # N_updrafts,\n # N_quad,\n # param_set,\n # surface = NeutralDrySurfaceModel{FT}(param_set),\n # )\n\n C_smag_ = C_smag(param_set)\n # turbulence = ConstantKinematicViscosity(FT(0.1))\n turbulence = SmagorinskyLilly{FT}(C_smag_)\n model = stable_bl_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbulence = turbulence,\n turbconv = turbconv,\n compressibility = compressibility,\n )\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n string(\"SBL_\", str_comp, \"_1D\"),\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config = config_diagnostics(driver_config)\n\n # boyd vandeven filter\n cb_boyd = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n (\"energy.ρe\",),\n solver_config.dg.grid,\n BoydVandevenFilter(\n solver_config.dg.grid,\n 1, #default=0\n 4, #default=32\n ),\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n # Mass tendencies = 0 for Anelastic1D model,\n # so mass should be completely conserved:\n check_cons =\n (ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.00000001)),)\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n if !isnothing(cl_args[\"cparam_file\"])\n ClimateMachine.Settings.output_dir = cl_args[\"cparam_file\"] * \".output\"\n end\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cb_boyd, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n# ArgParse in global scope to modify Clima Parameters\nsbl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sbl_args, \"StableBoundaryLayer\")\n@add_arg_table! sbl_args begin\n \"--cparam-file\"\n help = \"specify CLIMAParameters file\"\n arg_type = Union{String, Nothing}\n default = nothing\n\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sbl\"\n arg_type = String\n default = \"custom_sbl\"\nend\n\ncl_args = ClimateMachine.init(parse_clargs = true, custom_clargs = sbl_args)\nif !isnothing(cl_args[\"cparam_file\"])\n filename = cl_args[\"cparam_file\"]\n set_clima_parameters(filename)\nend\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\ninclude(joinpath(@__DIR__, \"report_mse_sbl_anelastic.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/stable_bl_coupled_edmf_an1d.jl", "content": "using JLD2, FileIO\nusing ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nimport ClimateMachine.BalanceLaws: projection\nimport ClimateMachine.DGMethods\nusing ClimateMachine.Atmos\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\nimport CLIMAParameters\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"stable_bl_model.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\nCLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 265\nCLIMAParameters.Atmos.EDMF.a_surf(::EarthParameterSet) = 0.0\nfunction set_clima_parameters(filename)\n eval(:(include($filename)))\nend\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n # GCM setting - Initialize the grid mean profiles of prognostic variables (ρ,e_int,q_tot,u,v,w)\n z = altitude(m, aux)\n\n # SCM setting - need to have separate cases coded and called from a folder - see what LES does\n # a thermo state is used here to convert the input θ to e_int profile\n e_int = internal_energy(m, state, aux)\n param_set = parameter_set(m)\n ts = PhaseDry(param_set, e_int, state.ρ)\n T = air_temperature(ts)\n p = air_pressure(ts)\n q = PhasePartition(ts)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = FT(0)\n end\n\n # initialize environment covariance with zero for now\n if z <= FT(250)\n en.ρatke =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n en.ρaθ_liq_cv =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n else\n en.ρatke = FT(0)\n en.ρaθ_liq_cv = FT(0)\n end\n en.ρaq_tot_cv = FT(0)\n en.ρaθ_liq_q_tot_cv = FT(0)\n return nothing\nend;\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # Choice of compressibility and CFL\n # compressibility = Compressible()\n compressibility = Anelastic1D()\n str_comp = compressibility == Compressible() ? \"COMPRESS\" : \"ANELASTIC\"\n\n # DG polynomial order\n N = 4\n nelem_vert = 20\n\n # Prescribe domain parameters\n zmax = FT(400)\n # Simulation time\n t0 = FT(0)\n timeend = FT(1800 * 1)\n CFLmax = compressibility == Compressible() ? FT(1) : FT(100)\n\n config_type = SingleStackConfigType\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n # Choice of SGS model\n N_updrafts = 1\n N_quad = 3\n turbconv = EDMF(\n FT,\n N_updrafts,\n N_quad,\n param_set,\n surface = NeutralDrySurfaceModel{FT}(param_set),\n coupling = Coupled(),\n )\n\n turbulence = ConstantKinematicViscosity(FT(0.0))\n model = stable_bl_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbulence = turbulence,\n turbconv = turbconv,\n compressibility = compressibility,\n )\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n string(\"SBL_COUPLED_\", str_comp, \"_1D\"),\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config = config_diagnostics(driver_config)\n\n # boyd vandeven filter\n num_state_prognostic = number_states(driver_config.bl, Prognostic())\n cb_boyd = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n 1:num_state_prognostic,\n solver_config.dg.grid,\n BoydVandevenFilter(\n solver_config.dg.grid,\n 1, #default=0\n 4, #default=32\n ),\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n # Mass tendencies = 0 for Anelastic1D model,\n # so mass should be completely conserved:\n check_cons =\n (ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.00000001)),)\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n if !isnothing(cl_args[\"cparam_file\"])\n ClimateMachine.Settings.output_dir = cl_args[\"cparam_file\"] * \".output\"\n end\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cb_boyd, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n# ArgParse in global scope to modify Clima Parameters\nsbl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sbl_args, \"StableBoundaryLayer\")\n@add_arg_table! sbl_args begin\n \"--cparam-file\"\n help = \"specify CLIMAParameters file\"\n arg_type = Union{String, Nothing}\n default = nothing\n\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sbl\"\n arg_type = String\n default = \"custom_sbl\"\nend\n\ncl_args = ClimateMachine.init(parse_clargs = true, custom_clargs = sbl_args)\nif !isnothing(cl_args[\"cparam_file\"])\n filename = cl_args[\"cparam_file\"]\n set_clima_parameters(filename)\nend\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\n# Uncomment lines to save output using JLD2\noutput_dir = @__DIR__;\nmkpath(output_dir);\nfunction dons(diag_vs_z)\n return Dict(map(keys(first(diag_vs_z))) do k\n string(k) => [getproperty(ca, k) for ca in diag_vs_z]\n end)\nend\nget_dons_arr(diag_arr) = [dons(diag_vs_z) for diag_vs_z in diag_arr]\ndons_arr = get_dons_arr(diag_arr)\nprintln(dons_arr[1].keys)\nz = get_z(solver_config.dg.grid; rm_dupes = true);\nsave(\n string(output_dir, \"/sbl_coupled.jld2\"),\n \"dons_arr\",\n dons_arr,\n \"time_data\",\n time_data,\n \"z\",\n z,\n)\n\ninclude(joinpath(@__DIR__, \"report_mse_sbl_coupled_edmf_an1d.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/stable_bl_edmf.jl", "content": "using JLD2, FileIO\nusing ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\nimport CLIMAParameters\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"stable_bl_model.jl\"))\ninclude(joinpath(\"helper_funcs\", \"diagnostics_configuration.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\nCLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 265\nCLIMAParameters.Atmos.EDMF.a_surf(::EarthParameterSet) = 0.0\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n # GCM setting - Initialize the grid mean profiles of prognostic variables (ρ,e_int,q_tot,u,v,w)\n z = altitude(m, aux)\n\n # SCM setting - need to have separate cases coded and called from a folder - see what LES does\n # a thermo state is used here to convert the input θ to e_int profile\n e_int = internal_energy(m, state, aux)\n param_set = parameter_set(m)\n ts = PhaseDry(param_set, e_int, state.ρ)\n T = air_temperature(ts)\n p = air_pressure(ts)\n q = PhasePartition(ts)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = FT(0)\n end\n\n # initialize environment covariance with zero for now\n if z <= FT(250)\n en.ρatke =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n en.ρaθ_liq_cv =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n else\n en.ρatke = FT(0)\n en.ρaθ_liq_cv = FT(0)\n end\n en.ρaq_tot_cv = FT(0)\n en.ρaθ_liq_q_tot_cv = FT(0)\n return nothing\nend;\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # DG polynomial order\n N = 4\n nelem_vert = 15\n\n # Prescribe domain parameters\n zmax = FT(400)\n\n t0 = FT(0)\n\n # Simulation time\n timeend = FT(60)\n CFLmax = FT(0.50)\n\n config_type = SingleStackConfigType\n\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n N_updrafts = 1\n N_quad = 3 # Using N_quad = 1 leads to norm(Q) = NaN at init.\n turbconv = EDMF(\n FT,\n N_updrafts,\n N_quad,\n param_set,\n surface = NeutralDrySurfaceModel{FT}(param_set),\n )\n\n model = stable_bl_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbconv = turbconv,\n )\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"SBL_EDMF\",\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config =\n config_diagnostics(driver_config, timeend; interval = \"10ssecs\")\n\n cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n (turbconv_filters(turbconv)...,),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n # Save to disc (for debugging):\n # @save \"bomex_edmf_nstep=$nstep.jld2\" diag_vs_z\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n check_cons = (\n ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.001)),\n ClimateMachine.ConservationCheck(\"energy.ρe\", \"3000steps\", FT(0.1)),\n )\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cbtmarfilter, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n\n# add a command line argument to specify the kind of surface flux\n# TODO: this will move to the future namelist functionality\nsbl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sbl_args, \"StableBoundaryLayer\")\n@add_arg_table! sbl_args begin\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sbl\"\n arg_type = String\n default = \"custom_sbl\"\nend\n\ncl_args = ClimateMachine.init(\n parse_clargs = true,\n custom_clargs = sbl_args,\n output_dir = get(ENV, \"CLIMATEMACHINE_SETTINGS_OUTPUT_DIR\", \"output\"),\n fix_rng_seed = true,\n)\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\n## Uncomment lines to save output using JLD2\n# output_dir = @__DIR__;\n# mkpath(output_dir);\n# function dons(diag_vs_z)\n# return Dict(map(keys(first(diag_vs_z))) do k\n# string(k) => [getproperty(ca, k) for ca in diag_vs_z]\n# end)\n# end\n# get_dons_arr(diag_arr) = [dons(diag_vs_z) for diag_vs_z in diag_arr]\n# dons_arr = get_dons_arr(diag_arr)\n# println(dons_arr[1].keys)\n# z = get_z(solver_config.dg.grid; rm_dupes = true);\n# save(\n# string(output_dir, \"/sbl_edmf.jld2\"),\n# \"dons_arr\",\n# dons_arr,\n# \"time_data\",\n# time_data,\n# \"z\",\n# z,\n# )\n\ninclude(joinpath(@__DIR__, \"report_mse_sbl_edmf.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/stable_bl_edmf_fvm.jl", "content": "using JLD2, FileIO\nusing ClimateMachine\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nimport ClimateMachine.DGMethods.FVReconstructions: FVLinear\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\nimport CLIMAParameters\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"stable_bl_model.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\nCLIMAParameters.Planet.T_surf_ref(::EarthParameterSet) = 265\nCLIMAParameters.Atmos.EDMF.a_surf(::EarthParameterSet) = 0.0\n\n\"\"\"\n init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n ) where {FT}\n\nInitialize EDMF state variables.\nThis method is only called at `t=0`.\n\"\"\"\nfunction init_state_prognostic!(\n turbconv::EDMF{FT},\n m::AtmosModel{FT},\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n) where {FT}\n # Aliases:\n gm = state\n en = state.turbconv.environment\n up = state.turbconv.updraft\n N_up = n_updrafts(turbconv)\n # GCM setting - Initialize the grid mean profiles of prognostic variables (ρ,e_int,q_tot,u,v,w)\n z = altitude(m, aux)\n\n # SCM setting - need to have separate cases coded and called from a folder - see what LES does\n # a thermo state is used here to convert the input θ to e_int profile\n e_int = internal_energy(m, state, aux)\n param_set = parameter_set(m)\n ts = PhaseDry(param_set, e_int, state.ρ)\n T = air_temperature(ts)\n p = air_pressure(ts)\n q = PhasePartition(ts)\n θ_liq = liquid_ice_pottemp(ts)\n\n a_min = turbconv.subdomains.a_min\n @unroll_map(N_up) do i\n up[i].ρa = gm.ρ * a_min\n up[i].ρaw = gm.ρu[3] * a_min\n up[i].ρaθ_liq = gm.ρ * a_min * θ_liq\n up[i].ρaq_tot = FT(0)\n end\n\n # initialize environment covariance with zero for now\n if z <= FT(250)\n en.ρatke =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n en.ρaθ_liq_cv =\n gm.ρ *\n FT(0.4) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0) *\n FT(1 - z / 250.0)\n else\n en.ρatke = FT(0)\n en.ρaθ_liq_cv = FT(0)\n end\n en.ρaq_tot_cv = FT(0)\n en.ρaθ_liq_q_tot_cv = FT(0)\n return nothing\nend;\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # DG polynomial order\n N = (1, 0)\n nelem_vert = 80\n\n # Prescribe domain parameters\n zmax = FT(400)\n\n t0 = FT(0)\n\n # Simulation time\n timeend = FT(60)\n CFLmax = FT(0.50)\n\n config_type = SingleStackConfigType\n\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n # solver_method = LSRK144NiegemannDiehlBusch,\n solver_method = LSRK54CarpenterKennedy,\n )\n\n N_updrafts = 1\n N_quad = 3 # Using N_quad = 1 leads to norm(Q) = NaN at init.\n turbconv = EDMF(\n FT,\n N_updrafts,\n N_quad,\n param_set,\n surface = NeutralDrySurfaceModel{FT}(param_set),\n )\n\n model = stable_bl_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbconv = turbconv,\n ref_state = HydrostaticState(\n DecayingTemperatureProfile{FT}(param_set);\n subtract_off = false,\n ),\n )\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"SBL_EDMF\",\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n numerical_flux_first_order = RoeNumericalFlux(),\n fv_reconstruction = HBFVReconstruction(model, FVLinear()),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config = config_diagnostics(driver_config)\n\n cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(100) do\n nstep = getsteps(solver_config.solver)\n Filters.apply!(\n solver_config.Q,\n (turbconv_filters(turbconv)...,),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n # Save to disc (for debugging):\n # @save \"bomex_edmf_nstep=$nstep.jld2\" diag_vs_z\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n check_cons = (\n ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.01)),\n ClimateMachine.ConservationCheck(\"energy.ρe\", \"3000steps\", FT(0.1)),\n )\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(500) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cbtmarfilter, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n# add a command line argument to specify the kind of surface flux\n# TODO: this will move to the future namelist functionality\nsbl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sbl_args, \"StableBoundaryLayer\")\n@add_arg_table! sbl_args begin\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sbl\"\n arg_type = String\n default = \"custom_sbl\"\nend\n\ncl_args = ClimateMachine.init(\n parse_clargs = true,\n custom_clargs = sbl_args,\n output_dir = get(ENV, \"CLIMATEMACHINE_SETTINGS_OUTPUT_DIR\", \"output\"),\n fix_rng_seed = true,\n)\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\ninclude(joinpath(@__DIR__, \"report_mse_sbl_edmf.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/stable_bl_single_stack_implicit.jl", "content": "using ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.BalanceLaws: vars_state\nusing JLD2, FileIO\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\ninclude(joinpath(clima_dir, \"experiments\", \"AtmosLES\", \"stable_bl_model.jl\"))\ninclude(\"edmf_model.jl\")\ninclude(\"edmf_kernels.jl\")\n\nfunction main(::Type{FT}, cl_args) where {FT}\n\n surface_flux = cl_args[\"surface_flux\"]\n\n # DG polynomial order\n N = 1\n nelem_vert = 50\n\n # Prescribe domain parameters\n zmax = FT(400)\n\n t0 = FT(0)\n\n # Simulation time\n timeend = FT(3600 * 6)\n CFLmax = FT(40.0)\n\n config_type = SingleStackConfigType\n\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n N_updrafts = 1\n N_quad = 3 # Using N_quad = 1 leads to norm(Q) = NaN at init.\n turbconv = NoTurbConv()\n\n C_smag = FT(0.23)\n\n model = stable_bl_model(\n FT,\n config_type,\n zmax,\n surface_flux;\n turbulence = SmagorinskyLilly{FT}(C_smag),\n turbconv = turbconv,\n )\n\n # Assemble configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"SBL_EDMF\",\n N,\n nelem_vert,\n zmax,\n param_set,\n model;\n hmax = FT(40),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n\n #################### Change the ode_solver to implicit solver\n\n dg = solver_config.dg\n Q = solver_config.Q\n\n\n vdg = DGModel(\n driver_config;\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n\n # linear solver relative tolerance rtol which should be slightly smaller than the nonlinear solver tol\n linearsolver = BatchedGeneralizedMinimalResidual(\n dg,\n Q;\n max_subspace_size = 30,\n atol = -1.0,\n rtol = 5e-5,\n )\n\n \"\"\"\n N(q)(Q) = Qhat => F(Q) = N(q)(Q) - Qhat\n\n F(Q) == 0\n ||F(Q^i) || / ||F(Q^0) || < tol\n\n \"\"\"\n # ϵ is a sensity parameter for this problem, it determines the finite difference Jacobian dF = (F(Q + ϵdQ) - F(Q))/ϵ\n # I have also try larger tol, but tol = 1e-3 does not work\n nonlinearsolver =\n JacobianFreeNewtonKrylovSolver(Q, linearsolver; tol = 1e-4, ϵ = 1.e-10)\n\n # this is a second order time integrator, to change it to a first order time integrator\n # change it ARK1ForwardBackwardEuler, which can reduce the cost by half at the cost of accuracy\n # and stability\n # preconditioner_update_freq = 50 means updating the preconditioner every 50 Newton solves,\n # update it more freqent will accelerate the convergence of linear solves, but updating it\n # is very expensive\n ode_solver = ARK2ImplicitExplicitMidpoint(\n dg,\n vdg,\n NonLinearBackwardEulerSolver(\n nonlinearsolver;\n isadjustable = true,\n preconditioner_update_freq = 50,\n ),\n Q;\n dt = solver_config.dt,\n t0 = 0,\n split_explicit_implicit = false,\n variant = NaiveVariant(),\n )\n\n solver_config.solver = ode_solver\n\n #######################################\n\n # --- Zero-out horizontal variations:\n vsp = vars_state(model, Prognostic(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :turbconv),\n )\n horizontally_average!(\n driver_config.grid,\n solver_config.Q,\n varsindex(vsp, :energy, :ρe),\n )\n vsa = vars_state(model, Auxiliary(), FT)\n horizontally_average!(\n driver_config.grid,\n solver_config.dg.state_auxiliary,\n varsindex(vsa, :turbconv),\n )\n # ---\n\n dgn_config = config_diagnostics(driver_config)\n\n cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n (turbconv_filters(turbconv)...,),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n\n diag_arr = [single_stack_diagnostics(solver_config)]\n time_data = FT[0]\n\n # Define the number of outputs from `t0` to `timeend`\n n_outputs = 5\n # This equates to exports every ceil(Int, timeend/n_outputs) time-step:\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n cb_data_vs_time =\n GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n diag_vs_z = single_stack_diagnostics(solver_config)\n\n nstep = getsteps(solver_config.solver)\n # Save to disc (for debugging):\n # @save \"bomex_edmf_nstep=$nstep.jld2\" diag_vs_z\n\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n check_cons =\n (ClimateMachine.ConservationCheck(\"ρ\", \"3000steps\", FT(0.001)),)\n\n cb_print_step = GenericCallbacks.EveryXSimulationSteps(100) do\n @show getsteps(solver_config.solver)\n nothing\n end\n\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n check_cons = check_cons,\n user_callbacks = (cbtmarfilter, cb_data_vs_time, cb_print_step),\n check_euclidean_distance = true,\n )\n\n diag_vs_z = single_stack_diagnostics(solver_config)\n push!(diag_arr, diag_vs_z)\n push!(time_data, gettime(solver_config.solver))\n\n return solver_config, diag_arr, time_data\nend\n\n# add a command line argument to specify the kind of surface flux\n# TODO: this will move to the future namelist functionality\nsbl_args = ArgParseSettings(autofix_names = true)\nadd_arg_group!(sbl_args, \"StableBoundaryLayer\")\n@add_arg_table! sbl_args begin\n \"--surface-flux\"\n help = \"specify surface flux for energy and moisture\"\n metavar = \"prescribed|bulk|custom_sbl\"\n arg_type = String\n default = \"custom_sbl\"\nend\n\ncl_args = ClimateMachine.init(\n parse_clargs = true,\n custom_clargs = sbl_args,\n output_dir = get(ENV, \"CLIMATEMACHINE_SETTINGS_OUTPUT_DIR\", \"output\"),\n fix_rng_seed = true,\n)\n\nsolver_config, diag_arr, time_data = main(Float64, cl_args)\n\ninclude(joinpath(@__DIR__, \"report_mse_sbl_ss_implicit.jl\"))\n\nnothing\n" }, { "path": "test/Atmos/EDMF/variable_map.jl", "content": "#! format: off\nvar_map(s::String) = var_map(Val(Symbol(s)))\nvar_map(::Val{T}) where {T} = nothing\n\nvar_map(::Val{:prog_ρ}) = (\"rho\", ())\nvar_map(::Val{:prog_ρu_1}) = (\"u_mean\", (:ρ,))\nvar_map(::Val{:prog_ρu_2}) = (\"v_mean\", (:ρ,))\nvar_map(::Val{:prog_moisture_ρq_tot}) = (\"qt_mean\", (:ρ,))\nvar_map(::Val{:prog_turbconv_updraft_1_ρa}) = (\"updraft_fraction\", (:ρ,))\nvar_map(::Val{:prog_turbconv_updraft_1_ρaw}) = (\"updraft_w\", (:ρ, :a))\nvar_map(::Val{:prog_turbconv_updraft_1_ρaq_tot}) = (\"updraft_qt\", (:ρ, :a))\nvar_map(::Val{:prog_turbconv_updraft_1_ρaθ_liq}) = (\"updraft_thetali\", (:ρ, :a))\nvar_map(::Val{:prog_turbconv_environment_ρatke}) = (\"tke_mean\", (:ρ, :a))\nvar_map(::Val{:prog_turbconv_environment_ρaθ_liq_cv}) = (\"env_thetali2\", (:ρ, :a))\nvar_map(::Val{:prog_turbconv_environment_ρaq_tot_cv}) = (\"env_qt2\", (:ρ, :a))\n#! format: on\n" }, { "path": "test/Atmos/Model/Artifacts.toml", "content": "[ref_state]\ngit-tree-sha1 = \"f6b0e460e5660732eb9e755b1e06d395ea3fb518\"\n" }, { "path": "test/Atmos/Model/discrete_hydrostatic_balance.jl", "content": "using ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.Atmos\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers: ManyColumnLU\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Mesh.Grids\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nusing LinearAlgebra\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nfunction init_to_ref_state!(problem, bl, state, aux, localgeo, t)\n FT = eltype(state)\n state.ρ = aux.ref_state.ρ\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = aux.ref_state.ρe\nend\n\nfunction config_balanced(\n FT,\n poly_order,\n temp_profile,\n numflux,\n (config_type, config_fun, config_args),\n)\n ref_state = HydrostaticState(temp_profile; subtract_off = false)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = ConstantDynamicViscosity(FT(0)),\n hyperdiffusion = NoHyperDiffusion(),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n config_type,\n physics;\n source = (Gravity(),),\n init_state_prognostic = init_to_ref_state!,\n )\n\n config = config_fun(\n \"balanced state\",\n poly_order,\n config_args...,\n param_set,\n nothing;\n model = model,\n numerical_flux_first_order = numflux,\n )\n\n return config\nend\n\nfunction main()\n FT = Float64\n poly_order = 4\n\n timestart = FT(0)\n timeend = FT(100)\n domain_height = FT(50e3)\n\n LES_params = let\n LES_resolution = ntuple(_ -> domain_height / 3poly_order, 3)\n LES_domain = ntuple(_ -> domain_height, 3)\n (LES_resolution, LES_domain...)\n end\n\n GCM_params = let\n GCM_resolution = (3, 3)\n (GCM_resolution, domain_height)\n end\n\n GCM = (AtmosGCMConfigType, ClimateMachine.AtmosGCMConfiguration, GCM_params)\n LES = (AtmosLESConfigType, ClimateMachine.AtmosLESConfiguration, LES_params)\n\n imex_solver_type = ClimateMachine.IMEXSolverType(\n splitting_type = HEVISplitting(),\n implicit_model = AtmosAcousticGravityLinearModel,\n implicit_solver = ManyColumnLU,\n solver_method = ARK2GiraldoKellyConstantinescu,\n )\n explicit_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK54CarpenterKennedy,\n )\n\n @testset for config in (LES, GCM)\n @testset for ode_solver_type in (explicit_solver_type, imex_solver_type)\n @testset for numflux in (\n CentralNumericalFluxFirstOrder(),\n RoeNumericalFlux(),\n HLLCNumericalFlux(),\n )\n @testset for temp_profile in (\n IsothermalProfile(param_set, FT),\n DecayingTemperatureProfile{FT}(param_set),\n )\n driver_config = config_balanced(\n FT,\n poly_order,\n temp_profile,\n numflux,\n config,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n timestart,\n timeend,\n driver_config,\n Courant_number = FT(0.1),\n init_on_cpu = true,\n ode_solver_type = ode_solver_type,\n CFL_direction = EveryDirection(),\n diffdir = HorizontalDirection(),\n )\n\n Qinit = similar(solver_config.Q)\n Qinit .= solver_config.Q\n\n ClimateMachine.invoke!(solver_config)\n\n error =\n euclidean_distance(solver_config.Q, Qinit) / norm(Qinit)\n @test error <= 100 * eps(FT)\n end\n end\n end\n end\nend\n\nmain()\n" }, { "path": "test/Atmos/Model/get_atmos_ref_states.jl", "content": "using ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.Mesh.Grids\nusing Thermodynamics\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing Test\nClimateMachine.init()\n\nfunction get_atmos_ref_states(nelem_vert, N_poly, RH)\n\n FT = Float64\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(\n DecayingTemperatureProfile{FT}(param_set),\n RH,\n ),\n )\n model = AtmosModel{FT}(\n SingleStackConfigType,\n physics;\n init_state_prognostic = (_...) -> nothing,\n )\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"ref_state\",\n N_poly,\n nelem_vert,\n FT(25e3),\n param_set,\n model,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n FT(0),\n FT(10),\n driver_config;\n skip_update_aux = true,\n ode_dt = FT(1),\n )\n\n return solver_config\nend\n" }, { "path": "test/Atmos/Model/ref_state.jl", "content": "include(\"get_atmos_ref_states.jl\")\nusing JLD2\nusing Pkg.Artifacts\nusing ArtifactWrappers\nusing Thermodynamics\nconst TD = Thermodynamics\n\n@testset \"Hydrostatic reference states - regression test\" begin\n ref_state_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"ref_state\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/gyq292ns79wm9xpmy1sse3qtnpcxw54q.jld2\",\n filename = \"ref_state.jld2\",\n ),],\n )\n ref_state_dataset_path = get_data_folder(ref_state_dataset)\n data_file = joinpath(ref_state_dataset_path, \"ref_state.jld2\")\n\n RH = 0.5\n (nelem_vert, N_poly) = (20, 4)\n solver_config = get_atmos_ref_states(nelem_vert, N_poly, RH)\n dons_arr = dict_of_nodal_states(solver_config, (Auxiliary(),))\n T = dons_arr[\"ref_state.T\"]\n p = dons_arr[\"ref_state.p\"]\n ρ = dons_arr[\"ref_state.ρ\"]\n\n @load \"$data_file\" T_ref p_ref ρ_ref\n @test all(isapprox.(T, T_ref; rtol = 1e-6))\n @test all(p .≈ p_ref)\n @test all(ρ .≈ ρ_ref)\nend\n\n@testset \"Hydrostatic reference states - correctness\" begin\n\n RH = 0.5\n # Fails on (80, 1)\n for (nelem_vert, N_poly) in [(40, 2), (20, 4)]\n solver_config = get_atmos_ref_states(nelem_vert, N_poly, RH)\n dons_arr = dict_of_nodal_states(solver_config)\n phase_type = PhaseEquil\n T = dons_arr[\"ref_state.T\"]\n p = dons_arr[\"ref_state.p\"]\n ρ = dons_arr[\"ref_state.ρ\"]\n q_tot = dons_arr[\"ref_state.ρq_tot\"] ./ ρ\n q_pt = PhasePartition.(q_tot)\n\n # TODO: test that ρ and p are in discrete hydrostatic balance\n\n # Test state for thermodynamic consistency (with ideal gas law)\n T_igl =\n TD.air_temperature_from_ideal_gas_law.(Ref(param_set), p, ρ, q_pt)\n @test all(T .≈ T_igl)\n\n # Test that relative humidity in reference state is approximately\n # input relative humidity\n RH_ref = relative_humidity.(Ref(param_set), T, p, Ref(phase_type), q_pt)\n @show max(abs.(RH .- RH_ref)...)\n @test all(isapprox.(RH, RH_ref, atol = 0.05))\n end\n\nend\n" }, { "path": "test/Atmos/Model/runtests.jl", "content": "using Test\n\n@testset \"Atmos Model\" begin\n include(\"ref_state.jl\")\nend\n" }, { "path": "test/Atmos/Parameterizations/Microphysics/KM_ice.jl", "content": "using Dierckx\n\ninclude(\"KinematicModel.jl\")\n\n# speed up the relaxation timescales for cloud water and cloud ice\nCLIMAParameters.Atmos.Microphysics.τ_cond_evap(::AbstractParameterSet) = 0.5\nCLIMAParameters.Atmos.Microphysics.τ_sub_dep(::AbstractParameterSet) = 0.5\n\nfunction vars_state(m::KinematicModel, ::Prognostic, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n ρq_tot::FT\n ρq_liq::FT\n ρq_ice::FT\n ρq_rai::FT\n ρq_sno::FT\n end\nend\n\nfunction vars_state(m::KinematicModel, ::Auxiliary, FT)\n @vars begin\n # defined in init_state_auxiliary\n p::FT\n z_coord::FT\n x_coord::FT\n # defined in update_aux\n u::FT\n w::FT\n q_tot::FT\n q_vap::FT\n q_liq::FT\n q_ice::FT\n q_rai::FT\n q_sno::FT\n e_tot::FT\n e_kin::FT\n e_pot::FT\n e_int::FT\n T::FT\n S_liq::FT\n S_ice::FT\n RH::FT\n rain_w::FT\n snow_w::FT\n # more diagnostics\n src_cloud_liq::FT\n src_cloud_ice::FT\n src_rain_acnv::FT\n src_snow_acnv::FT\n src_liq_rain_accr::FT\n src_liq_snow_accr::FT\n src_ice_snow_accr::FT\n src_ice_rain_accr::FT\n src_snow_rain_accr::FT\n src_rain_accr_sink::FT\n src_rain_evap::FT\n src_snow_subl::FT\n src_snow_melt::FT\n flag_cloud_liq::FT\n flag_cloud_ice::FT\n flag_rain::FT\n flag_snow::FT\n # helpers for bc\n ρe_init::FT\n ρq_tot_init::FT\n end\nend\n\nfunction init_kinematic_eddy!(eddy_model, state, aux, localgeo, t, spline_fun)\n\n FT = eltype(state)\n _grav::FT = grav(param_set)\n\n dc = eddy_model.data_config\n\n (x, y, z) = localgeo.coord\n (xc, yc, zc) = localgeo.center_coord\n\n @inbounds begin\n\n init_T, init_qt, init_p, init_ρ, init_dρ = spline_fun\n\n # density\n q_pt_0 = PhasePartition(init_qt(z))\n R_m, cp_m, cv_m, γ = gas_constants(param_set, q_pt_0)\n T::FT = init_T(z)\n ρ::FT = init_ρ(z)\n state.ρ = ρ\n aux.p = init_p(z)\n\n # moisture\n state.ρq_tot = ρ * init_qt(z)\n state.ρq_liq = ρ * q_pt_0.liq\n state.ρq_ice = ρ * q_pt_0.ice\n state.ρq_rai = ρ * FT(0)\n state.ρq_sno = ρ * FT(0)\n\n # [Grabowski1998](@cite)\n # velocity (derivative of streamfunction)\n # This is actually different than what comes out from taking a\n # derivative of Ψ from the paper. I have sin(π/2/X(x-xc)).\n # This setup makes more sense to me though.\n _Z::FT = FT(15000)\n _X::FT = FT(10000)\n _xc::FT = FT(30000)\n _A::FT = FT(4.8 * 1e4)\n _S::FT = FT(2.5 * 1e-2) * FT(0.01) #TODO\n _ρ_00::FT = FT(1)\n ρu::FT = FT(0)\n ρw::FT = FT(0)\n fact =\n _A / _ρ_00 * (\n init_ρ(z) * FT(π) / _Z * cos(FT(π) / _Z * z) +\n init_dρ(z) * sin(FT(π) / _Z * z)\n )\n\n if zc < _Z\n if x >= (_xc + _X)\n ρu = _S * z - fact\n ρw = FT(0)\n elseif x <= (_xc - _X)\n ρu = _S * z + fact\n ρw = FT(0)\n else\n ρu = _S * z - fact * sin(FT(π / 2.0) / _X * (x - _xc))\n ρw =\n _A * init_ρ(z) / _ρ_00 * FT(π / 2.0) / _X *\n sin(FT(π) / _Z * z) *\n cos(FT(π / 2.0) / _X * (x - _xc))\n end\n else\n ρu = _S * z\n ρw = FT(0)\n end\n state.ρu = SVector(ρu, FT(0), ρw)\n u::FT = ρu / ρ\n w::FT = ρw / ρ\n\n # energy\n e_kin::FT = 1 // 2 * (u^2 + w^2)\n e_pot::FT = _grav * z\n e_int::FT = internal_energy(param_set, T, q_pt_0)\n e_tot::FT = e_kin + e_pot + e_int\n state.ρe = ρ * e_tot\n end\n return nothing\nend\n\nfunction nodal_update_auxiliary_state!(\n m::KinematicModel,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n FT = eltype(state)\n _grav::FT = grav(param_set)\n _T_freeze::FT = T_freeze(param_set)\n\n @inbounds begin\n\n if t == FT(0)\n aux.ρe_init = state.ρe\n aux.ρq_tot_init = state.ρq_tot\n end\n\n # velocity\n aux.u = state.ρu[1] / state.ρ\n aux.w = state.ρu[3] / state.ρ\n # water\n aux.q_tot = state.ρq_tot / state.ρ\n aux.q_liq = state.ρq_liq / state.ρ\n aux.q_ice = state.ρq_ice / state.ρ\n aux.q_rai = state.ρq_rai / state.ρ\n aux.q_sno = state.ρq_sno / state.ρ\n aux.q_vap = aux.q_tot - aux.q_liq - aux.q_ice\n # energy\n aux.e_tot = state.ρe / state.ρ\n aux.e_kin = 1 // 2 * (aux.u^2 + aux.w^2)\n aux.e_pot = _grav * aux.z_coord\n aux.e_int = aux.e_tot - aux.e_kin - aux.e_pot\n # supersaturation\n q = PhasePartition(aux.q_tot, aux.q_liq, aux.q_ice)\n aux.T = air_temperature(param_set, aux.e_int, q)\n ts_neq = PhaseNonEquil_ρTq(param_set, state.ρ, aux.T, q)\n aux.S_liq = max(0, supersaturation(ts_neq, Liquid()))\n aux.S_ice = max(0, supersaturation(ts_neq, Ice()))\n aux.RH = relative_humidity(ts_neq) * FT(100)\n\n aux.rain_w =\n terminal_velocity(param_set, CM1M.RainType(), state.ρ, aux.q_rai)\n aux.snow_w =\n terminal_velocity(param_set, CM1M.SnowType(), state.ρ, aux.q_sno)\n\n # more diagnostics\n ts_eq = PhaseEquil_ρTq(param_set, state.ρ, aux.T, aux.q_tot)\n q_eq = PhasePartition(ts_eq)\n\n aux.src_cloud_liq =\n conv_q_vap_to_q_liq_ice(param_set, CM1M.LiquidType(), q_eq, q)\n aux.src_cloud_ice =\n conv_q_vap_to_q_liq_ice(param_set, CM1M.IceType(), q_eq, q)\n\n aux.src_rain_acnv = conv_q_liq_to_q_rai(param_set, aux.q_liq)\n aux.src_snow_acnv = conv_q_ice_to_q_sno(param_set, q, state.ρ, aux.T)\n\n aux.src_liq_rain_accr = accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.RainType(),\n aux.q_liq,\n aux.q_rai,\n state.ρ,\n )\n aux.src_liq_snow_accr = accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.SnowType(),\n aux.q_liq,\n aux.q_sno,\n state.ρ,\n )\n aux.src_ice_snow_accr = accretion(\n param_set,\n CM1M.IceType(),\n CM1M.SnowType(),\n aux.q_ice,\n aux.q_sno,\n state.ρ,\n )\n aux.src_ice_rain_accr = accretion(\n param_set,\n CM1M.IceType(),\n CM1M.RainType(),\n aux.q_ice,\n aux.q_rai,\n state.ρ,\n )\n\n aux.src_rain_accr_sink =\n accretion_rain_sink(param_set, aux.q_ice, aux.q_rai, state.ρ)\n\n if aux.T < _T_freeze\n aux.src_snow_rain_accr = accretion_snow_rain(\n param_set,\n CM1M.SnowType(),\n CM1M.RainType(),\n aux.q_sno,\n aux.q_rai,\n state.ρ,\n )\n else\n aux.src_snow_rain_accr = accretion_snow_rain(\n param_set,\n CM1M.RainType(),\n CM1M.SnowType(),\n aux.q_rai,\n aux.q_sno,\n state.ρ,\n )\n end\n\n aux.src_rain_evap = evaporation_sublimation(\n param_set,\n CM1M.RainType(),\n q,\n aux.q_rai,\n state.ρ,\n aux.T,\n )\n aux.src_snow_subl = evaporation_sublimation(\n param_set,\n CM1M.SnowType(),\n q,\n aux.q_sno,\n state.ρ,\n aux.T,\n )\n\n aux.src_snow_melt = snow_melt(param_set, aux.q_sno, state.ρ, aux.T)\n\n aux.flag_cloud_liq = FT(0)\n aux.flag_cloud_ice = FT(0)\n aux.flag_rain = FT(0)\n aux.flag_snow = FT(0)\n if (aux.q_liq >= FT(0))\n aux.flag_cloud_liq = FT(1)\n end\n if (aux.q_ice >= FT(0))\n aux.flag_cloud_ice = FT(1)\n end\n if (aux.q_rai >= FT(0))\n aux.flag_rain = FT(1)\n end\n if (aux.q_sno >= FT(0))\n aux.flag_snow = FT(1)\n end\n end\nend\n\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n bctype,\n m::KinematicModel,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n # 1 - left (x = 0, z = ...)\n # 2 - right (x = -1, z = ...)\n # 3,4 - y boundary (periodic)\n # 5 - bottom (x = ..., z = 0)\n # 6 - top (x = ..., z = -1)\n\n FT = eltype(state⁻)\n @inbounds state⁺.ρ = state⁻.ρ\n @inbounds state⁺.ρe = aux⁻.ρe_init\n @inbounds state⁺.ρq_tot = aux⁻.ρq_tot_init\n @inbounds state⁺.ρq_liq = FT(0) #state⁻.ρq_liq\n @inbounds state⁺.ρq_ice = FT(0) #state⁻.ρq_ice\n @inbounds state⁺.ρq_rai = FT(0)\n @inbounds state⁺.ρq_sno = FT(0)\n\n if bctype == 1\n @inbounds state⁺.ρu = SVector(state⁻.ρu[1], FT(0), FT(0))\n end\n if bctype == 2\n @inbounds state⁺.ρu = SVector(state⁻.ρu[1], FT(0), FT(0))\n @inbounds state⁺.ρe = state⁻.ρe\n @inbounds state⁺.ρq_tot = state⁻.ρq_tot\n @inbounds state⁺.ρq_liq = state⁻.ρq_liq\n @inbounds state⁺.ρq_ice = state⁻.ρq_ice\n end\n if bctype == 5\n @inbounds state⁺.ρu -= 2 * dot(state⁻.ρu, n) .* SVector(n)\n end\n if bctype == 6\n @inbounds state⁺.ρu = SVector(state⁻.ρu[1], FT(0), state⁻.ρu[3])\n end\nend\n\n@inline function wavespeed(\n m::KinematicModel,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n FT = eltype(state)\n @inbounds begin\n u = state.ρu / state.ρ\n q_rai::FT = state.ρq_rai / state.ρ\n q_sno::FT = state.ρq_sno / state.ρ\n rain_w = terminal_velocity(param_set, CM1M.RainType(), state.ρ, q_rai)\n snow_w = terminal_velocity(param_set, CM1M.SnowType(), state.ρ, q_sno)\n\n nu =\n nM[1] * u[1] +\n nM[3] * max(u[3], rain_w, snow_w, u[3] - rain_w, u[3] - snow_w)\n end\n return abs(nu)\nend\n\n@inline function flux_first_order!(\n m::KinematicModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n FT = eltype(state)\n @inbounds begin\n q_rai::FT = state.ρq_rai / state.ρ\n q_sno::FT = state.ρq_sno / state.ρ\n rain_w = terminal_velocity(param_set, CM1M.RainType(), state.ρ, q_rai)\n snow_w = terminal_velocity(param_set, CM1M.SnowType(), state.ρ, q_sno)\n\n # advect moisture ...\n flux.ρ = SVector(state.ρu[1], FT(0), state.ρu[3])\n flux.ρq_tot = SVector(\n state.ρu[1] * state.ρq_tot / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_tot / state.ρ,\n )\n flux.ρq_liq = SVector(\n state.ρu[1] * state.ρq_liq / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_liq / state.ρ,\n )\n flux.ρq_ice = SVector(\n state.ρu[1] * state.ρq_ice / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_ice / state.ρ,\n )\n flux.ρq_rai = SVector(\n state.ρu[1] * state.ρq_rai / state.ρ,\n FT(0),\n (state.ρu[3] / state.ρ - rain_w) * state.ρq_rai,\n )\n flux.ρq_sno = SVector(\n state.ρu[1] * state.ρq_sno / state.ρ,\n FT(0),\n (state.ρu[3] / state.ρ - snow_w) * state.ρq_sno,\n )\n # ... energy ...\n flux.ρe = SVector(\n state.ρu[1] / state.ρ * (state.ρe + aux.p),\n FT(0),\n state.ρu[3] / state.ρ * (state.ρe + aux.p),\n )\n # ... and don't advect momentum (kinematic setup)\n end\nend\n\nfunction source!(\n m::KinematicModel,\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n FT = eltype(state)\n\n _grav::FT = grav(param_set)\n\n _e_int_v0::FT = e_int_v0(param_set)\n _e_int_i0::FT = e_int_i0(param_set)\n\n _cv_d::FT = cv_d(param_set)\n _cv_v::FT = cv_v(param_set)\n _cv_l::FT = cv_l(param_set)\n _cv_i::FT = cv_i(param_set)\n\n _T_0::FT = T_0(param_set)\n _T_freeze = T_freeze(param_set)\n\n @inbounds begin\n e_tot = state.ρe / state.ρ\n q_tot = state.ρq_tot / state.ρ\n q_liq = state.ρq_liq / state.ρ\n q_ice = state.ρq_ice / state.ρ\n q_rai = state.ρq_rai / state.ρ\n q_sno = state.ρq_sno / state.ρ\n u = state.ρu[1] / state.ρ\n w = state.ρu[3] / state.ρ\n ρ = state.ρ\n e_int = e_tot - 1 // 2 * (u^2 + w^2) - _grav * aux.z_coord\n\n q = PhasePartition(q_tot, q_liq, q_ice)\n T = air_temperature(param_set, e_int, q)\n _Lf = latent_heat_fusion(param_set, T)\n # equilibrium state at current T\n ts_eq = PhaseEquil_ρTq(param_set, state.ρ, T, q_tot)\n q_eq = PhasePartition(ts_eq)\n\n # zero out the source terms\n source.ρq_tot = FT(0)\n source.ρq_liq = FT(0)\n source.ρq_ice = FT(0)\n source.ρq_rai = FT(0)\n source.ρq_sno = FT(0)\n source.ρe = FT(0)\n\n # vapour -> cloud liquid water\n source.ρq_liq +=\n ρ * conv_q_vap_to_q_liq_ice(param_set, CM1M.LiquidType(), q_eq, q)\n # vapour -> cloud ice\n source.ρq_ice +=\n ρ * conv_q_vap_to_q_liq_ice(param_set, CM1M.IceType(), q_eq, q)\n\n ## cloud liquid water -> rain\n acnv = ρ * conv_q_liq_to_q_rai(param_set, q_liq)\n source.ρq_liq -= acnv\n source.ρq_tot -= acnv\n source.ρq_rai += acnv\n source.ρe -= acnv * (_cv_l - _cv_d) * (T - _T_0)\n\n ## cloud ice -> snow\n acnv = ρ * conv_q_ice_to_q_sno(param_set, q, state.ρ, T)\n source.ρq_ice -= acnv\n source.ρq_tot -= acnv\n source.ρq_sno += acnv\n source.ρe -= acnv * ((_cv_i - _cv_d) * (T - _T_0) - _e_int_i0)\n\n # cloud liquid water + rain -> rain\n accr =\n ρ * accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.RainType(),\n q_liq,\n q_rai,\n state.ρ,\n )\n source.ρq_liq -= accr\n source.ρq_tot -= accr\n source.ρq_rai += accr\n source.ρe -= accr * (_cv_l - _cv_d) * (T - _T_0)\n\n # cloud ice + snow -> snow\n accr =\n ρ * accretion(\n param_set,\n CM1M.IceType(),\n CM1M.SnowType(),\n q_ice,\n q_sno,\n state.ρ,\n )\n source.ρq_ice -= accr\n source.ρq_tot -= accr\n source.ρq_sno += accr\n source.ρe -= accr * ((_cv_i - _cv_d) * (T - _T_0) - _e_int_i0)\n\n # cloud liquid water + snow -> snow or rain\n accr =\n ρ * accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.SnowType(),\n q_liq,\n q_sno,\n state.ρ,\n )\n if T < _T_freeze\n source.ρq_liq -= accr\n source.ρq_tot -= accr\n source.ρq_sno += accr\n source.ρe -= accr * ((_cv_i - _cv_d) * (T - _T_0) - _e_int_i0)\n else\n source.ρq_liq -= accr\n source.ρq_tot -= accr\n source.ρq_sno -= accr * (_cv_l / _Lf * (T - _T_freeze))\n source.ρq_rai += accr * (FT(1) + _cv_l / _Lf * (T - _T_freeze))\n source.ρe +=\n -accr * ((_cv_l - _cv_d) * (T - _T_0) + _cv_l * (T - _T_freeze))\n end\n\n # cloud ice + rain -> snow\n accr =\n ρ * accretion(\n param_set,\n CM1M.IceType(),\n CM1M.RainType(),\n q_ice,\n q_rai,\n state.ρ,\n )\n accr_rain_sink =\n ρ * accretion_rain_sink(param_set, q_ice, q_rai, state.ρ)\n source.ρq_ice -= accr\n source.ρq_tot -= accr\n source.ρq_rai -= accr_rain_sink\n source.ρq_sno += accr + accr_rain_sink\n source.ρe +=\n accr_rain_sink * _Lf -\n accr * ((_cv_i - _cv_d) * (T - _T_0) - _e_int_i0)\n\n # rain + snow -> snow or rain\n if T < _T_freeze\n accr =\n ρ * accretion_snow_rain(\n param_set,\n CM1M.SnowType(),\n CM1M.RainType(),\n q_sno,\n q_rai,\n state.ρ,\n )\n source.ρq_sno += accr\n source.ρq_rai -= accr\n source.ρe += accr * _Lf\n else\n accr =\n ρ * accretion_snow_rain(\n param_set,\n CM1M.RainType(),\n CM1M.SnowType(),\n q_rai,\n q_sno,\n state.ρ,\n )\n source.ρq_rai += accr\n source.ρq_sno -= accr\n source.ρe -= accr * _Lf\n end\n\n # rain -> vapour\n evap =\n ρ * evaporation_sublimation(\n param_set,\n CM1M.RainType(),\n q,\n q_rai,\n state.ρ,\n T,\n )\n source.ρq_rai += evap\n source.ρq_tot -= evap\n source.ρe -= evap * (_cv_l - _cv_d) * (T - _T_0)\n\n # snow -> vapour\n subl =\n ρ * evaporation_sublimation(\n param_set,\n CM1M.SnowType(),\n q,\n q_sno,\n state.ρ,\n T,\n )\n source.ρq_sno += subl\n source.ρq_tot -= subl\n source.ρe -= subl * ((_cv_i - _cv_d) * (T - _T_0) - _e_int_i0)\n\n # snow -> rain\n melt = ρ * snow_melt(param_set, q_sno, state.ρ, T)\n\n source.ρq_sno -= melt\n source.ρq_rai += melt\n source.ρe -= melt * _Lf\n end\nend\n\nfunction main()\n # Working precision\n FT = Float64\n # DG polynomial order\n N = 4\n # Domain resolution and size\n Δx = FT(500)\n Δy = FT(1)\n Δz = FT(250)\n resolution = (Δx, Δy, Δz)\n # Domain extents\n xmax = 90000\n ymax = 10\n zmax = 16000\n # initial configuration\n wmax = FT(0.6) # max velocity of the eddy [m/s]\n θ_0 = FT(289) # init. theta value (const) [K]\n p_0 = FT(101500) # surface pressure [Pa]\n p_1000 = FT(100000) # reference pressure in theta definition [Pa]\n qt_0 = FT(7.5 * 1e-3) # init. total water specific humidity (const) [kg/kg]\n z_0 = FT(0) # surface height\n\n # time stepping\n t_ini = FT(0)\n t_end = FT(5 * 60) #FT(4 * 60 * 60) #TODO\n dt = FT(0.25)\n #CFL = FT(1.75)\n filter_freq = 1\n output_freq = 1200\n interval = \"1200steps\"\n\n # periodicity and boundary numbers\n periodicity_x = false\n periodicity_y = true\n periodicity_z = false\n idx_bc_left = 1\n idx_bc_right = 2\n idx_bc_front = 3\n idx_bc_back = 4\n idx_bc_bottom = 5\n idx_bc_top = 6\n\n\n #! format: off\n z_range = [\n 0, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900, 2000, 2100, 2200, 2300, 2400, 2500,\n 2600, 2700, 2800, 2900, 3000, 3100, 3200, 3300, 3400, 3500, 3600, 3700, 3800, 3900, 4000, 4100, 4200, 4300, 4400, 4500, 4600, 4700, 4800, 4900,\n 5000, 5100, 5200, 5300, 5400, 5500, 5600, 5700, 5800, 5900, 6000, 6100, 6200, 6300, 6400, 6500, 6600, 6700, 6800, 6900, 7000, 7100, 7200, 7300,\n 7400, 7500, 7600, 7700, 7800, 7900, 8000, 8100, 8200, 8300, 8400, 8500, 8600, 8700, 8800, 8900, 9000, 9100, 9200, 9300, 9400, 9500, 9600, 9700,\n 9800, 9900, 10000, 10100, 10200, 10300, 10400, 10500, 10600, 10700, 10800, 10900, 11000, 11100, 11200, 11300, 11400, 11500, 11600, 11700,\n 11800, 11900, 12000, 12100, 12200, 12300, 12400, 12500, 12600, 12700, 12800, 12900, 13000, 13100, 13200, 13300, 13400, 13500, 13600, 13700,\n 13800, 13900, 14000, 14100, 14200, 14300, 14400, 14500, 14600, 14700, 14800, 14900, 15000, 15100, 15200, 15300, 15400, 15500, 15600, 15700,\n 15800, 15900, 16000, 16100, 16200, 16300, 16400, 16500, 16600, 16700, 16800, 16900, 17000]\n\n T_range = [\n 299.184, 297.628892697526, 296.498576625316, 295.708541362848, 295.174276489603, 294.811271585061, 294.5350162287, 294.261,\n 293.920421043582, 293.507311764631, 293.031413133474, 292.502466120435, 291.930211695841, 291.324390830018, 290.694744493293, 290.051013655991,\n 289.402939288437, 288.760262346828, 288.130849318341, 287.515505258001, 286.913372692428, 286.323594148241, 285.74531215206, 285.177669230506,\n 284.619807910198, 284.070870717755, 283.530000179798, 282.996338822946, 282.469029173819, 281.947213759037, 281.43003510522, 280.916635738988,\n 280.406158186959, 279.897744975755, 279.390538631995, 278.883681682298, 278.376316653285, 277.867586071576, 277.356632463789, 276.842598356546,\n 276.324685598782, 275.802667484647, 275.276636664165, 274.746689276312, 274.212921460063, 273.675429354393, 273.134309098278, 272.589656830692,\n 272.041568690611, 271.49014081701, 270.935469348864, 270.377650425148, 269.816780184838, 269.252954766908, 268.686270310335, 268.116822954093,\n 267.544708837157, 266.970024098502, 266.392864877104, 265.813327311938, 265.231507541979, 264.64747745881, 264.060960130126, 263.47141707779,\n 262.878303486873, 262.281074542448, 261.679185429586, 261.072091333358, 260.459247438837, 259.840108931094, 259.214130995201, 258.580768816228,\n 257.939477579249, 257.289712469334, 256.630928671556, 255.962581370986, 255.284125752695, 254.595017001755, 253.894710327766, 253.182887991267,\n 252.46001426264, 251.726721646298, 250.983642646651, 250.231409768111, 249.470655515088, 248.702012391995, 247.926112903242, 247.14358955324,\n 246.355074846402, 245.561201287137, 244.762601379858, 243.959907628976, 243.153728314245, 242.344349049499, 241.531825963971, 240.716210279609,\n 239.897553218363, 239.075906002182, 238.251319853016, 237.423845992814, 236.593535643524, 235.760440027097, 234.924610365481, 234.086097880626,\n 233.244953794481, 232.401229328996, 231.554975706119, 230.7062441478, 229.855085875989, 229.001552112634, 228.145694079684, 227.287574051742,\n 226.42757986392, 225.566470305038, 224.705024343724, 223.84402094861, 222.984239088323, 222.126457731494, 221.271455846752, 220.420012402726,\n 219.572906368047, 218.730916711342, 217.894822401242, 217.065402406377, 216.243435695375, 215.429701236865, 214.624977999479, 213.830040741192,\n 213.045586818019, 212.272245655601, 211.51064437337, 210.761410090759, 210.0251699272, 209.302551002126, 208.594180434969, 207.90068534516,\n 207.222682432974, 206.560708873131, 205.915265185628, 205.286851694754, 204.675968724799, 204.083116600052, 203.508795644802, 202.95350618334,\n 202.417748539953, 201.902023577843, 201.407120974388, 200.934594512862, 200.486122895611, 200.063384824981, 199.668059003318, 199.301824132969,\n 198.96635891628, 198.663342055596, 198.394452253265, 198.161279286222, 197.963547643009, 197.799211358657, 197.666154525954, 197.562261237686,\n 197.485415586642, 197.433501665607, 197.40440356737, 197.396005384718, 197.406191210439, 197.432899562191, 197.47476703369, 197.530913717737,\n 197.600469393593, 197.682563840515, 197.776326837764]\n\n q_range = [\n 0.0162321692080669, 0.0167673003973785, 0.0169842024326598, 0.0169337519538153, 0.0166663497239751, 0.0162321692080669, 0.015681352854829,\n 0.0150641570577189, 0.0144310460364482, 0.0138327343680399, 0.013320177602368, 0.012928720727755, 0.0126305522364067, 0.0123819207241334,\n 0.0121389587958033, 0.0118577075098814, 0.0115066332638913, 0.0111042588941075, 0.010681699769686, 0.0102701709364724, 0.0099009900990099,\n 0.00959624936114117, 0.00934081059469477, 0.00911016542683696, 0.00887975233624587, 0.0086249628234361, 0.00832855696312967, 0.0080029519599734,\n 0.00766802230757524, 0.00734369421478042, 0.007049945387747, 0.00680137732291751, 0.00659092388806814, 0.00640607314371358, 0.00623428890344645,\n 0.00606301560481065, 0.00588287130354119, 0.005697230175774, 0.00551266483334836, 0.00533575758863248, 0.00517309988062077, 0.00502869447640062,\n 0.004896158939909, 0.00476650506495816, 0.00463073495294506, 0.00447984071677452, 0.00430803604057719, 0.00412246330674869, 0.0039335098478028,\n 0.00375157835183152, 0.00358708648864089, 0.00344726315880705, 0.00332653452987796, 0.00321611598159287, 0.00310721394018452, 0.00299102691924227,\n 0.00286116318208214, 0.00272089896736183, 0.00257593295768495, 0.00243197016394643, 0.00229472213908012, 0.00216860732664965, 0.00205285001385606,\n 0.00194537329026098, 0.00184409796137261, 0.00174694285001248, 0.0016522111336286, 0.00155974942036123, 0.00146979043510184, 0.0013825670343689,\n 0.00129831219414761, 0.00121701485317001, 0.001137687161242, 0.00105909638608903, 0.000980009009141718, 0.00089919072834449, 0.000815997019873583,\n 0.000732145660692171, 0.00064994635423271, 0.000571710312185518, 0.000499750124937531, 0.000435858502312957, 0.000379744409038409,\n 0.000330595173284745, 0.000287597487595375, 0.000249937515621095, 0.000216895300299521, 0.000188127832006747, 0.000163386230957203,\n 0.000142421466578296, 0.000124984376952881, 0.000110689772327441, 9.86086827423503e-05, 8.7676088112503e-05, 7.68268601082915e-05,\n 6.49957752746072e-05, 5.15108887771554e-05, 3.72737503739521e-05, 2.35794100864691e-05, 1.17230575137195e-05, 2.999991000027e-06,\n 1.69619792857815e-06, 3.07897588648539e-06, 2.26364535967942e-06, 3.65542843474416e-07, 1.49999775000338e-06, 2.45283190408077e-06,\n 2.55351934700394e-06, 2.09779045350945e-06, 1.38137398580681e-06, 6.99999510000343e-07, 2.86443287526343e-07, 1.21662921693633e-07,\n 1.23660803334821e-07, 2.10439089554136e-07, 2.99999910000027e-07, 3.27788493770303e-07, 2.99022297121994e-07, 2.36361815382337e-07,\n 1.62467544245733e-07, 9.9999990000001e-08, 6.64022211520834e-08, 5.82475139663064e-08, 6.68916919016523e-08, 8.36905765913686e-08,\n 9.9999990000001e-08, 1.09002484657468e-07, 1.11187531664659e-07, 1.0887133139317e-07, 1.04370084083517e-07, 9.9999990000001e-08,\n 9.75877787592422e-08, 9.70022971221656e-08, 9.76229211090621e-08, 9.88290267308115e-08, 9.9999990000001e-08, 1.00646340200832e-07,\n 1.00803219684934e-07, 1.0063692406888e-07, 1.00313748968563e-07, 9.9999990000001e-08, 9.98268004299175e-08, 9.97847641264901e-08,\n 9.9829322608123e-08, 9.99159173931719e-08, 9.9999990000001e-08, 1.00046398078965e-07, 1.00057663808278e-07, 1.0004572549811e-07,\n 1.00022521458629e-07, 9.9999990000001e-08, 9.99875472541912e-08, 9.99845206403455e-08, 9.99877153994048e-08, 9.99939367723098e-08,\n 9.9999990000001e-08, 1.00003352904274e-07, 1.00004193630342e-07, 1.00003352904274e-07, 1.00001671452137e-07, 9.9999990000001e-08,\n 9.99989811287191e-08, 9.99986448382919e-08, 9.99988129835055e-08, 9.99993174191465e-08, 9.9999990000001e-08, 1.00000662580856e-07,\n 1.00001167016497e-07, 1.0000133516171e-07, 1.00000998871283e-07, 9.9999990000001e-08]\n\n p_range = [\n 101500, 100315.025749516, 99139.9660271538, 97974.7710180798, 96819.3909074598, 95673.7758804597, 94537.876176083, 93411.6422486837,\n 92295.0246064536, 91187.9737719499, 90090.4404029628, 89002.3752464385, 87923.7291106931, 86854.4528968052, 85794.49764766, 84743.814421206,\n 83702.3543252787, 82670.0686672607, 81646.9088044215, 80632.826109327, 79627.7720985403, 78631.6983820215, 77644.5566288927, 76666.2986022276,\n 75696.8762101083, 74736.2413760207, 73784.3460742183, 72841.1424820254, 71906.5828275339, 70980.6193544154, 70063.204453007, 69154.2906087673,\n 68253.8303673938, 67361.7763702768, 66478.0814065118, 65602.6982808849, 64735.579849904, 63876.6791769641, 63025.949377182, 62183.3435815507,\n 61348.8150705175, 60522.317221463, 59703.8034731598, 58893.2273619179, 58090.542574603, 57295.7028140736, 56508.6618359204, 55729.3736066619,\n 54957.7921455483, 54193.8714880201, 53437.5658219265, 52688.8294339685, 51947.6166734621, 51213.8819892159, 50487.5799836161, 49768.6652753625,\n 49057.0925369589, 48352.8166561238, 47655.7925743794, 46965.9752497711, 46283.3197958903, 45607.7814272171, 44939.3154221458, 44277.8771606411,\n 43623.4221794563, 42975.9060319998, 42335.2843266238, 41701.5128914544, 41074.5476095614, 40454.3443808926, 39840.8592642808, 39234.0484216161,\n 38633.868080086, 38040.274570662, 37453.2243845259, 36872.6740298782, 36298.5800710783, 35730.8992971221, 35169.5885531642, 34614.6047016162,\n 34065.9047673383, 33523.4458805626, 32987.1852382959, 32457.0801436939, 31933.0880637786, 31415.1664829788, 30903.2729431824, 30397.3652161132,\n 29897.4011309542, 29403.3385345499, 28915.135440009, 28432.7499682896, 27956.1403087067, 27485.2647592591, 27020.0817857294, 26560.5498717234,\n 26106.6275597007, 25658.2736275352, 25215.4469119545, 24778.106267783, 24346.2107202062, 23919.7194049202, 23498.5915276787, 23082.7864056438,\n 22672.2635279751, 22266.9824021027, 21866.9025958119, 21471.9839183077, 21082.1862391501, 20697.469446465, 20317.7936031548, 19943.1188855014,\n 19573.4055416778, 19208.6139342067, 18848.7046021555, 18493.6381033452, 18143.3750575744, 17797.8763325512, 17457.1028579612, 17121.0155825635,\n 16789.575634671, 16462.7442590809, 16140.4827744667, 15822.752617039, 15509.5154044811, 15200.7327737537, 14896.3664255555, 14596.378315538,\n 14300.7304630906, 14009.3849072278, 13722.3038717104, 13439.4497001572, 13160.7848122237, 12886.2717485712, 12615.8732366994, 12349.552023955,\n 12087.2709233428, 11828.9930105004, 11574.6814267238, 11324.2993335354, 11077.8100828701, 10835.1771502091, 10596.3640894538, 10361.3345792892,\n 10130.0524910676, 9902.48171660711, 9678.58621552163, 9458.33021860609, 9241.67802445076, 9028.59395243382, 8819.04251762682, 8612.98836336888,\n 8410.3962187252, 8211.2309379189, 8015.45756519932, 7823.04116500174, 7633.94687886457, 7448.14015673894, 7265.58652567913, 7086.25151273945,\n 6910.10064497419, 6910.10064497419, 6910.10064497419, 6910.10064497419, 6910.10064497419, 6910.10064497419, 6910.10064497419, 6910.10064497419,\n 6910.10064497419, 6910.10064497419, 6910.10064497419 ]\n\n ρ_range = [\n 1.17051362179725, 1.16251824590627, 1.15313016112624, 1.1426566568231, 1.13140764159872, 1.1196895485888, 1.10780103484676, 1.0960305075832,\n 1.08459745274435, 1.07348344546654, 1.06261389733783, 1.05192695939713, 1.04140273218196, 1.03103180334705, 1.02080494033636, 1.01071303046314,\n 1.00073949487493, 0.990838128123433, 0.980963129632192, 0.971094063252625, 0.961216822759604, 0.951323251200104, 0.941426814080701,\n 0.931545806522527, 0.921697879536919, 0.911900034022987, 0.902164578162167, 0.892487313537494, 0.882860249793038, 0.873275726601859,\n 0.863726411598162, 0.854208105591224, 0.844727922437057, 0.8352954667279, 0.825919969478444, 0.816610286976537, 0.807373343850795,\n 0.798209593686815, 0.789117941453129, 0.780097337389561, 0.771146612530646, 0.762264270399455, 0.753452744759353, 0.74471551136303,\n 0.736055874178972, 0.72747696742561, 0.718980350842218, 0.7105619137831, 0.70221632469523, 0.693938467344086, 0.685723438615901,\n 0.677567861578885, 0.669473709625074, 0.661444143644113, 0.653482167320322, 0.645590627902435, 0.637771283999691, 0.630022091606339,\n 0.622340165988717, 0.614722736979947, 0.607167148420651, 0.599671383580347, 0.592236161810936, 0.584863205035121, 0.577554164218497,\n 0.570310604494631, 0.563133875176496, 0.556024730482117, 0.548983747911074, 0.542011464088805, 0.535108377645103, 0.528275027848979,\n 0.521512222336075, 0.514820790002398, 0.508201503057759, 0.501655076786909, 0.495181994058148, 0.488781994213898, 0.482454636660537,\n 0.476199055919507, 0.470012952419531, 0.463893943289264, 0.457840307961496, 0.451850519497349, 0.445923093407251, 0.440056584487831,\n 0.434249561420135, 0.428500536139493, 0.422808039743611, 0.41717064648228, 0.411586974922854, 0.406055719826046, 0.40057578774956,\n 0.395146674561446, 0.389768253950893, 0.384440397458089, 0.37916287643901, 0.373935102771167, 0.36875641578014, 0.363626174398787,\n 0.358543758947717, 0.353508656205969, 0.348520707739897, 0.343579834253573, 0.338685947906162, 0.333838950271578, 0.329038686145689,\n 0.324284810799633, 0.319576938026155, 0.314914687777366, 0.310297688380068, 0.305725572307512, 0.301197591585183, 0.296712545171569,\n 0.292269251091959, 0.287866570849034, 0.283503406532332, 0.279178694598262, 0.27489141222907, 0.270640579366203, 0.266425260909222,\n 0.262244566294465, 0.258097650982816, 0.25398371429053, 0.249901999425458, 0.245851790781137, 0.241832414286703, 0.237843244370408,\n 0.233883781756254, 0.229953637104228, 0.226052458592054, 0.222179928272777, 0.218335761470863, 0.214519707224664, 0.210731549222293,\n 0.206971102905131, 0.203238226271481, 0.199532888953038, 0.195855122091922, 0.192204980437394, 0.188582544697456, 0.184987920053912,\n 0.181421234993557, 0.177882641628151, 0.174372316360016, 0.170890455921146, 0.167437037955477, 0.164011457709294, 0.160613093688105,\n 0.157241413931045, 0.15389597846219, 0.150576437441233, 0.147282529608045, 0.144014082002788, 0.140771009961674, 0.137553374518383,\n 0.134362580606833, 0.131201172349935, 0.128071588633207, 0.124976110750845, 0.121916861907264, 0.121948919193911, 0.121966894890357,\n 0.121972083953105, 0.121965790399749, 0.121949291096701, 0.121923436104917, 0.121888780280552, 0.121845875238272, 0.121795274570523,\n 0.121737533130168]\n\n dρ_range = [\n -3.59937517425616e-05, -8.74236524613322e-05, -9.98245979099523e-05, -0.00010912795442958, -0.000115340910679925, -0.000118522777615024,\n -0.000118766703148947, -0.000116181583122055, -0.00011260864631165, -0.000109795778726265, -0.000107714951035596, -0.000106039943861968,\n -0.000104460270820561, -0.000102973683474333, -0.000101578762228199, -0.000100274450938608, -9.92861102161202e-05, -9.88281212364185e-05,\n -9.86963271464388e-05, -9.87085122902761e-05, -9.88588230322522e-05, -9.8981065311496e-05, -9.89171829920543e-05, -9.86735488551799e-05,\n -9.82566574041453e-05, -9.76729921159812e-05, -9.70500064028184e-05, -9.65086017246105e-05, -9.60454400625446e-05, -9.56572439955241e-05,\n -9.5340728706457e-05, -9.5008763743794e-05, -9.45788919208772e-05, -9.40548370115828e-05, -9.3440348194004e-05, -9.27391611164088e-05,\n -9.20016033219379e-05, -9.12752248054719e-05, -9.05595697070144e-05, -8.98541831357713e-05, -8.91654815701201e-05, -8.84752762943482e-05,\n -8.77494447742739e-05, -8.69897168393693e-05, -8.61978039593136e-05, -8.53753836002353e-05, -8.45662005091732e-05, -8.38114256450538e-05,\n -8.31088828495735e-05, -8.24564348214017e-05, -8.1851961718636e-05, -8.12540449384666e-05, -8.06237270285041e-05, -7.99625864890717e-05,\n -7.92721952020196e-05, -7.85541037865398e-05, -7.78377768942751e-05, -7.71508783726086e-05, -7.64922509147597e-05, -7.58607550079355e-05,\n -7.5255257575603e-05, -7.46578263440909e-05, -7.40437160173945e-05, -7.34126676822411e-05, -7.27655403821961e-05, -7.21031866632626e-05,\n -7.14303672571544e-05, -7.07515642299482e-05, -7.00671933172211e-05, -6.93776513635009e-05, -6.86833098465315e-05, -6.79822059852417e-05,\n -6.72725239627913e-05, -6.65548356116998e-05, -6.58297114361597e-05, -6.50977145039132e-05, -6.43646771280071e-05, -6.36360540138981e-05,\n -6.29118551436968e-05, -6.22041292322341e-05, -6.15221570675246e-05, -6.08606412866085e-05, -6.02146109367996e-05, -5.95836319123706e-05,\n -5.89672983766289e-05, -5.8365225270403e-05, -5.7777757303502e-05, -5.72051962789751e-05, -5.66471081172719e-05, -5.61030583593601e-05,\n -5.55726043916327e-05, -5.50543285425413e-05, -5.45449966705289e-05, -5.40374636372204e-05, -5.35311627048121e-05, -5.30261985318617e-05,\n -5.25253559659224e-05, -5.20312184290027e-05, -5.15435901963297e-05, -5.10622691485893e-05, -5.05870395077161e-05, -5.01151309694529e-05,\n -4.96439716171912e-05, -4.91736441117515e-05, -4.87042469224382e-05, -4.82358875773372e-05, -4.77700486961607e-05, -4.73081018750182e-05,\n -4.68499874419861e-05, -4.63956358750055e-05, -4.59449597549036e-05, -4.54985344751462e-05, -4.50631284786667e-05, -4.46397691588802e-05,\n -4.42280098023119e-05, -4.38274214139499e-05, -4.34376409961601e-05, -4.30583023124663e-05, -4.26889789379916e-05, -4.23292356052957e-05,\n -4.19786196104028e-05, -4.16366783570754e-05, -4.13029635571471e-05, -4.09770310849877e-05, -4.06584588726149e-05, -4.03468387207095e-05,\n -4.00417403212974e-05, -3.9742477928175e-05, -3.94474194224351e-05, -3.91560587225531e-05, -3.88680420136295e-05, -3.85830382145498e-05,\n -3.83007173685112e-05, -3.80207275262849e-05, -3.77427401195186e-05, -3.74664417323129e-05, -3.71911016163793e-05, -3.69155924693975e-05,\n -3.66396540346307e-05, -3.6363043157771e-05, -3.60854967755032e-05, -3.58067813502568e-05, -3.55266676411338e-05, -3.52449049917705e-05,\n -3.49612716252254e-05, -3.46756639284376e-05, -3.43938806523058e-05, -3.41187600097204e-05, -3.38494075224641e-05, -3.35849193859056e-05,\n -3.33243734832797e-05, -3.3066881627664e-05, -3.28115645135523e-05, -3.25575278845173e-05, -3.23039138342865e-05, -3.20463462400501e-05,\n -3.17652064416084e-05, -3.14588936623928e-05, -3.11289784299643e-05, -3.07770405929711e-05, -1.50046969840055e-05, -2.48015587883801e-07,\n -1.13657079374249e-07, -7.70202792434148e-09, -1.15984516793898e-07, -2.12713494118585e-07, -3.03471417808635e-07, -3.88725590955189e-07,\n -4.68452477358432e-07, -5.42635999340642e-07, -6.1126695981267e-07]\n #! format: on\n\n init_T = Spline1D(z_range, T_range)\n init_qt = Spline1D(z_range, q_range)\n init_p = Spline1D(z_range, p_range)\n init_ρ = Spline1D(z_range, ρ_range)\n init_dρ = Spline1D(z_range, dρ_range)\n\n driver_config, ode_solver_type = config_kinematic_eddy(\n FT,\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n wmax,\n θ_0,\n p_0,\n p_1000,\n qt_0,\n z_0,\n periodicity_x,\n periodicity_y,\n periodicity_z,\n idx_bc_left,\n idx_bc_right,\n idx_bc_front,\n idx_bc_back,\n idx_bc_bottom,\n idx_bc_top,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t_ini,\n t_end,\n driver_config,\n (init_T, init_qt, init_p, init_ρ, init_dρ),\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n init_on_cpu = true,\n #Courant_number = CFL,\n )\n\n model = driver_config.bl\n\n mpicomm = MPI.COMM_WORLD\n\n # get state variables indices for filtering\n ρq_liq_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_liq)\n ρq_ice_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_ice)\n ρq_rai_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_rai)\n ρq_sno_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_sno)\n # get aux variables indices for testing\n q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)\n q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)\n q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)\n q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)\n q_rai_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_rai)\n q_sno_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_sno)\n S_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :S_liq)\n S_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :S_ice)\n rain_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :rain_w)\n snow_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :snow_w)\n\n # filter out negative values\n cb_tmar_filter =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:ρq_tot, :ρq_liq, :ρq_ice, :ρq_rai, :ρq_sno),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n cb_boyd_filter =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:ρq_tot, :ρq_liq, :ρq_ice, :ρq_rai, :ρq_sno, :ρe, :ρ),\n solver_config.dg.grid,\n BoydVandevenFilter(solver_config.dg.grid, 1, 8),\n )\n end\n\n # output for paraview\n\n # initialize base output prefix directory from rank 0\n vtkdir = abspath(joinpath(ClimateMachine.Settings.output_dir, \"vtk\"))\n if MPI.Comm_rank(mpicomm) == 0\n mkpath(vtkdir)\n end\n MPI.Barrier(mpicomm)\n\n vtkstep = [0]\n cb_vtk =\n GenericCallbacks.EveryXSimulationSteps(output_freq) do (init = false)\n out_dirname = @sprintf(\n \"microphysics_test_4_mpirank%04d_step%04d\",\n MPI.Comm_rank(mpicomm),\n vtkstep[1]\n )\n out_path_prefix = joinpath(vtkdir, out_dirname)\n @info \"doing VTK output\" out_path_prefix\n writevtk(\n out_path_prefix,\n solver_config.Q,\n solver_config.dg,\n flattenednames(vars_state(model, Prognostic(), FT)),\n solver_config.dg.state_auxiliary,\n flattenednames(vars_state(model, Auxiliary(), FT)),\n )\n vtkstep[1] += 1\n nothing\n end\n\n # output for netcdf\n boundaries = [\n FT(0) FT(0) FT(0)\n xmax ymax zmax\n ]\n interpol = ClimateMachine.InterpolationConfiguration(\n driver_config,\n boundaries,\n resolution,\n )\n dgngrps = [\n setup_dump_state_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n setup_dump_aux_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n ]\n dgn_config = ClimateMachine.DiagnosticsConfiguration(dgngrps)\n\n # call solve! function for time-integrator\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (cb_boyd_filter, cb_tmar_filter),\n check_euclidean_distance = true,\n )\n\n max_q_tot = maximum(abs.(solver_config.dg.state_auxiliary[:, q_tot_ind, :]))\n @test !isnan(max_q_tot)\n\n max_q_vap = maximum(abs.(solver_config.dg.state_auxiliary[:, q_vap_ind, :]))\n @test !isnan(max_q_vap)\n\n max_q_liq = maximum(abs.(solver_config.dg.state_auxiliary[:, q_liq_ind, :]))\n @test !isnan(max_q_liq)\n\n max_q_ice = maximum(abs.(solver_config.dg.state_auxiliary[:, q_ice_ind, :]))\n @test !isnan(max_q_ice)\n\n max_q_rai = maximum(abs.(solver_config.dg.state_auxiliary[:, q_rai_ind, :]))\n @test !isnan(max_q_rai)\n\n max_q_sno = maximum(abs.(solver_config.dg.state_auxiliary[:, q_sno_ind, :]))\n @test !isnan(max_q_sno)\n\nend\n\nmain()\n" }, { "path": "test/Atmos/Parameterizations/Microphysics/KM_saturation_adjustment.jl", "content": "include(\"KinematicModel.jl\")\n\nfunction vars_state(m::KinematicModel, ::Prognostic, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n ρq_tot::FT\n end\nend\n\nfunction vars_state(m::KinematicModel, ::Auxiliary, FT)\n @vars begin\n # defined in init_state_auxiliary\n p::FT\n x_coord::FT\n z_coord::FT\n # defined in update_aux\n u::FT\n w::FT\n q_tot::FT\n q_vap::FT\n q_liq::FT\n q_ice::FT\n e_tot::FT\n e_kin::FT\n e_pot::FT\n e_int::FT\n T::FT\n S_liq::FT\n RH::FT\n end\nend\n\nfunction init_kinematic_eddy!(eddy_model, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n FT = eltype(state)\n\n _grav::FT = grav(param_set)\n\n dc = eddy_model.data_config\n @inbounds begin\n # density\n q_pt_0 = PhasePartition(dc.qt_0)\n R_m, cp_m, cv_m, γ = gas_constants(param_set, q_pt_0)\n T::FT = dc.θ_0 * (aux.p / dc.p_1000)^(R_m / cp_m)\n ρ::FT = aux.p / R_m / T\n state.ρ = ρ\n\n # moisture\n state.ρq_tot = ρ * dc.qt_0\n\n # velocity (derivative of streamfunction)\n ρu::FT =\n dc.wmax * dc.xmax / dc.zmax *\n cos(FT(π) * z / dc.zmax) *\n cos(2 * FT(π) * x / dc.xmax)\n ρw::FT =\n 2 * dc.wmax * sin(FT(π) * z / dc.zmax) * sin(2 * π * x / dc.xmax)\n state.ρu = SVector(ρu, FT(0), ρw)\n u::FT = ρu / ρ\n w::FT = ρw / ρ\n\n # energy\n e_kin::FT = 1 // 2 * (u^2 + w^2)\n e_pot::FT = _grav * z\n e_int::FT = internal_energy(param_set, T, q_pt_0)\n e_tot::FT = e_kin + e_pot + e_int\n state.ρe = ρ * e_tot\n end\n return nothing\nend\n\nfunction nodal_update_auxiliary_state!(\n m::KinematicModel,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n FT = eltype(state)\n _grav::FT = grav(param_set)\n @inbounds begin\n aux.u = state.ρu[1] / state.ρ\n aux.w = state.ρu[3] / state.ρ\n\n aux.q_tot = state.ρq_tot / state.ρ\n\n aux.e_tot = state.ρe / state.ρ\n aux.e_kin = 1 // 2 * (aux.u^2 + aux.w^2)\n aux.e_pot = _grav * aux.z_coord\n aux.e_int = aux.e_tot - aux.e_kin - aux.e_pot\n\n # saturation adjustment happens here\n ts = PhaseEquil_ρeq(param_set, state.ρ, aux.e_int, aux.q_tot)\n q = PhasePartition(ts)\n\n aux.T = ts.T\n aux.q_vap = vapor_specific_humidity(q)\n aux.q_liq = q.liq\n aux.q_ice = q.ice\n\n # TODO: add super_saturation method in moist thermo\n #aux.S = max(0, aux.q_vap / q_vap_saturation(ts) - FT(1)) * FT(100)\n aux.S_liq = max(0, supersaturation(ts, Liquid()))\n aux.RH = relative_humidity(ts)\n end\nend\n\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n bctype,\n m::KinematicModel,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n t,\n args...,\n) end\n\n@inline function wavespeed(\n m::KinematicModel,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n @inbounds u = state.ρu / state.ρ\n return abs(dot(nM, u))\nend\n\n@inline function flux_first_order!(\n m::KinematicModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n FT = eltype(state)\n @inbounds begin\n # advect moisture ...\n flux.ρq_tot = SVector(\n state.ρu[1] * state.ρq_tot / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_tot / state.ρ,\n )\n # ... energy ...\n flux.ρe = SVector(\n state.ρu[1] / state.ρ * (state.ρe + aux.p),\n FT(0),\n state.ρu[3] / state.ρ * (state.ρe + aux.p),\n )\n # ... and don't advect momentum (kinematic setup)\n end\nend\n\nsource!(::KinematicModel, _...) = nothing\n\nfunction main()\n # Working precision\n FT = Float64\n # DG polynomial order\n N = 4\n # Domain resolution and size\n Δx = FT(20)\n Δy = FT(1)\n Δz = FT(20)\n resolution = (Δx, Δy, Δz)\n # Domain extents\n xmax = 1500\n ymax = 10\n zmax = 1500\n # initial configuration\n wmax = FT(0.6) # max velocity of the eddy [m/s]\n θ_0 = FT(289) # init. theta value (const) [K]\n p_0 = FT(101500) # surface pressure [Pa]\n p_1000 = FT(100000) # reference pressure in theta definition [Pa]\n qt_0 = FT(7.5 * 1e-3) # init. total water specific humidity (const) [kg/kg]\n z_0 = FT(0) # surface height\n\n # time stepping\n t_ini = FT(0)\n t_end = FT(60 * 30)\n dt = 40\n output_freq = 9\n interval = \"9steps\"\n\n # periodicity and boundary numbers\n periodicity_x = true\n periodicity_y = true\n periodicity_z = false\n idx_bc_left = 0\n idx_bc_right = 0\n idx_bc_front = 0\n idx_bc_back = 0\n idx_bc_bottom = 1\n idx_bc_top = 2\n\n driver_config, ode_solver_type = config_kinematic_eddy(\n FT,\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n wmax,\n θ_0,\n p_0,\n p_1000,\n qt_0,\n z_0,\n periodicity_x,\n periodicity_y,\n periodicity_z,\n idx_bc_left,\n idx_bc_right,\n idx_bc_front,\n idx_bc_back,\n idx_bc_bottom,\n idx_bc_top,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t_ini,\n t_end,\n driver_config;\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n init_on_cpu = true,\n #Courant_number = CFL,\n )\n\n mpicomm = MPI.COMM_WORLD\n\n # output for paraview\n # initialize base prefix directory from rank 0\n vtkdir = abspath(joinpath(ClimateMachine.Settings.output_dir, \"vtk\"))\n if MPI.Comm_rank(mpicomm) == 0\n mkpath(vtkdir)\n end\n MPI.Barrier(mpicomm)\n\n model = driver_config.bl\n vtkstep = [0]\n cbvtk = GenericCallbacks.EveryXSimulationSteps(output_freq) do\n out_dirname = @sprintf(\n \"new_ex_1_mpirank%04d_step%04d\",\n MPI.Comm_rank(mpicomm),\n vtkstep[1]\n )\n out_path_prefix = joinpath(vtkdir, out_dirname)\n @info \"doing VTK output\" out_path_prefix\n writevtk(\n out_path_prefix,\n solver_config.Q,\n solver_config.dg,\n flattenednames(vars_state(model, Prognostic(), FT)),\n solver_config.dg.state_auxiliary,\n flattenednames(vars_state(model, Auxiliary(), FT)),\n )\n vtkstep[1] += 1\n nothing\n end\n\n # output for netcdf\n boundaries = [\n FT(0) FT(0) FT(0)\n xmax ymax zmax\n ]\n interpol = ClimateMachine.InterpolationConfiguration(\n driver_config,\n boundaries,\n resolution,\n )\n dgngrps = [\n setup_dump_state_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n setup_dump_aux_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n ]\n dgn_config = ClimateMachine.DiagnosticsConfiguration(dgngrps)\n\n # get aux variables indices for testing\n q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)\n q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)\n q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)\n q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)\n S_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :S_liq)\n\n # call solve! function for time-integrator\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (cbvtk,),\n check_euclidean_distance = true,\n )\n\n # no supersaturation\n max_S_liq = maximum(abs.(solver_config.dg.state_auxiliary[:, S_liq_ind, :]))\n @test isequal(max_S_liq, FT(0))\n\n # qt is conserved\n max_q_tot = maximum(abs.(solver_config.dg.state_auxiliary[:, q_tot_ind, :]))\n min_q_tot = minimum(abs.(solver_config.dg.state_auxiliary[:, q_tot_ind, :]))\n @test isapprox(max_q_tot, qt_0; rtol = 1e-3)\n @test isapprox(min_q_tot, qt_0; rtol = 1e-3)\n\n # q_vap + q_liq = q_tot\n max_water_diff = maximum(abs.(\n solver_config.dg.state_auxiliary[:, q_tot_ind, :] .-\n solver_config.dg.state_auxiliary[:, q_vap_ind, :] .-\n solver_config.dg.state_auxiliary[:, q_liq_ind, :],\n ))\n @test isequal(max_water_diff, FT(0))\n\n # no ice\n max_q_ice = maximum(abs.(solver_config.dg.state_auxiliary[:, q_ice_ind, :]))\n @test isequal(max_q_ice, FT(0))\n\n # q_liq ∈ reference range\n max_q_liq = maximum(solver_config.dg.state_auxiliary[:, q_liq_ind, :])\n min_q_liq = minimum(solver_config.dg.state_auxiliary[:, q_liq_ind, :])\n @test max_q_liq < FT(1e-3)\n @test isequal(min_q_liq, FT(0))\nend\n\nmain()\n" }, { "path": "test/Atmos/Parameterizations/Microphysics/KM_warm_rain.jl", "content": "include(\"KinematicModel.jl\")\n\nfunction vars_state(m::KinematicModel, ::Prognostic, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n ρq_tot::FT\n ρq_liq::FT\n ρq_ice::FT\n ρq_rai::FT\n end\nend\n\nfunction vars_state(m::KinematicModel, ::Auxiliary, FT)\n @vars begin\n # defined in init_state_auxiliary\n p::FT\n x_coord::FT\n z_coord::FT\n # defined in update_aux\n u::FT\n w::FT\n q_tot::FT\n q_vap::FT\n q_liq::FT\n q_ice::FT\n q_rai::FT\n e_tot::FT\n e_kin::FT\n e_pot::FT\n e_int::FT\n T::FT\n S_liq::FT\n RH::FT\n rain_w::FT\n # more diagnostics\n src_cloud_liq::FT\n src_cloud_ice::FT\n src_acnv::FT\n src_accr::FT\n src_rain_evap::FT\n flag_rain::FT\n flag_cloud_liq::FT\n flag_cloud_ice::FT\n end\nend\n\nfunction init_kinematic_eddy!(eddy_model, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n FT = eltype(state)\n\n _grav::FT = grav(param_set)\n\n dc = eddy_model.data_config\n\n @inbounds begin\n # density\n q_pt_0 = PhasePartition(dc.qt_0)\n R_m, cp_m, cv_m, γ = gas_constants(param_set, q_pt_0)\n T::FT = dc.θ_0 * (aux.p / dc.p_1000)^(R_m / cp_m)\n ρ::FT = aux.p / R_m / T\n state.ρ = ρ\n\n # moisture\n state.ρq_tot = ρ * dc.qt_0\n state.ρq_liq = ρ * q_pt_0.liq\n state.ρq_ice = ρ * q_pt_0.ice\n state.ρq_rai = ρ * FT(0)\n\n # velocity (derivative of streamfunction)\n ρu::FT =\n dc.wmax * dc.xmax / dc.zmax *\n cos(π * z / dc.zmax) *\n cos(2 * π * x / dc.xmax)\n ρw::FT = 2 * dc.wmax * sin(π * z / dc.zmax) * sin(2 * π * x / dc.xmax)\n state.ρu = SVector(ρu, FT(0), ρw)\n u::FT = ρu / ρ\n w::FT = ρw / ρ\n\n # energy\n e_kin::FT = 1 // 2 * (u^2 + w^2)\n e_pot::FT = _grav * z\n e_int::FT = internal_energy(param_set, T, q_pt_0)\n e_tot::FT = e_kin + e_pot + e_int\n state.ρe = ρ * e_tot\n end\n return nothing\nend\n\nfunction nodal_update_auxiliary_state!(\n m::KinematicModel,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n FT = eltype(state)\n\n _grav::FT = grav(param_set)\n\n @inbounds begin\n # velocity\n aux.u = state.ρu[1] / state.ρ\n aux.w = state.ρu[3] / state.ρ\n # water\n aux.q_tot = state.ρq_tot / state.ρ\n aux.q_liq = state.ρq_liq / state.ρ\n aux.q_ice = state.ρq_ice / state.ρ\n aux.q_rai = state.ρq_rai / state.ρ\n q = PhasePartition(aux.q_tot, aux.q_liq, aux.q_ice)\n aux.q_vap = vapor_specific_humidity(q)\n # energy\n aux.e_tot = state.ρe / state.ρ\n aux.e_kin = 1 // 2 * (aux.u^2 + aux.w^2)\n aux.e_pot = _grav * aux.z_coord\n aux.e_int = aux.e_tot - aux.e_kin - aux.e_pot\n # supersaturation\n q = PhasePartition(aux.q_tot, aux.q_liq, aux.q_ice)\n aux.T = air_temperature(param_set, aux.e_int, q)\n ts_neq = PhaseNonEquil_ρTq(param_set, state.ρ, aux.T, q)\n aux.S_liq = max(0, supersaturation(ts_neq, Liquid()))\n aux.RH = relative_humidity(ts_neq) * FT(100)\n\n aux.rain_w =\n terminal_velocity(param_set, CM1M.RainType(), state.ρ, aux.q_rai)\n\n # more diagnostics\n ts_eq = PhaseEquil_ρTq(param_set, state.ρ, aux.T, aux.q_tot)\n q_eq = PhasePartition(ts_eq)\n\n aux.src_cloud_liq =\n conv_q_vap_to_q_liq_ice(param_set, CM1M.LiquidType(), q_eq, q)\n aux.src_cloud_ice =\n conv_q_vap_to_q_liq_ice(param_set, CM1M.IceType(), q_eq, q)\n aux.src_acnv = conv_q_liq_to_q_rai(param_set, aux.q_liq)\n aux.src_accr = accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.RainType(),\n aux.q_liq,\n aux.q_rai,\n state.ρ,\n )\n aux.src_rain_evap = evaporation_sublimation(\n param_set,\n CM1M.RainType(),\n q,\n aux.q_rai,\n state.ρ,\n aux.T,\n )\n aux.flag_cloud_liq = FT(0)\n aux.flag_cloud_ice = FT(0)\n aux.flag_rain = FT(0)\n if (aux.q_liq >= FT(0))\n aux.flag_cloud_liq = FT(1)\n end\n if (aux.q_ice >= FT(0))\n aux.flag_cloud_ice = FT(1)\n end\n if (aux.q_rai >= FT(0))\n aux.flag_rain = FT(1)\n end\n end\nend\n\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n bctype,\n m::KinematicModel,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n FT = eltype(state⁻)\n\n #state⁺.ρu -= 2 * dot(state⁻.ρu, n) .* SVector(n)\n\n #state⁺.ρq_rai = -state⁻.ρq_rai\n @inbounds state⁺.ρq_rai = FT(0)\n\nend\n\n@inline function wavespeed(\n m::KinematicModel,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n FT = eltype(state)\n\n @inbounds begin\n u = state.ρu / state.ρ\n q_rai::FT = state.ρq_rai / state.ρ\n\n rain_w = terminal_velocity(param_set, CM1M.RainType(), state.ρ, q_rai)\n nu = nM[1] * u[1] + nM[3] * max(u[3], rain_w, u[3] - rain_w)\n end\n return abs(nu)\nend\n\n@inline function flux_first_order!(\n m::KinematicModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n FT = eltype(state)\n\n @inbounds begin\n q_rai::FT = state.ρq_rai / state.ρ\n rain_w = terminal_velocity(param_set, CM1M.RainType(), state.ρ, q_rai)\n\n # advect moisture ...\n flux.ρq_tot = SVector(\n state.ρu[1] * state.ρq_tot / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_tot / state.ρ,\n )\n flux.ρq_liq = SVector(\n state.ρu[1] * state.ρq_liq / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_liq / state.ρ,\n )\n flux.ρq_ice = SVector(\n state.ρu[1] * state.ρq_ice / state.ρ,\n FT(0),\n state.ρu[3] * state.ρq_ice / state.ρ,\n )\n flux.ρq_rai = SVector(\n state.ρu[1] * state.ρq_rai / state.ρ,\n FT(0),\n (state.ρu[3] / state.ρ - rain_w) * state.ρq_rai,\n )\n # ... energy ...\n flux.ρe = SVector(\n state.ρu[1] / state.ρ * (state.ρe + aux.p),\n FT(0),\n state.ρu[3] / state.ρ * (state.ρe + aux.p),\n )\n # ... and don't advect momentum (kinematic setup)\n end\nend\n\nfunction source!(\n m::KinematicModel,\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n FT = eltype(state)\n _grav::FT = grav(param_set)\n _e_int_v0::FT = e_int_v0(param_set)\n _cv_v::FT = cv_v(param_set)\n _cv_d::FT = cv_d(param_set)\n _T_0::FT = T_0(param_set)\n\n @inbounds begin\n e_tot = state.ρe / state.ρ\n q_tot = state.ρq_tot / state.ρ\n q_liq = state.ρq_liq / state.ρ\n q_ice = state.ρq_ice / state.ρ\n q_rai = state.ρq_rai / state.ρ\n u = state.ρu[1] / state.ρ\n w = state.ρu[3] / state.ρ\n e_int = e_tot - 1 // 2 * (u^2 + w^2) - _grav * aux.z_coord\n\n q = PhasePartition(q_tot, q_liq, q_ice)\n T = air_temperature(param_set, e_int, q)\n # equilibrium state at current T\n ts_eq = PhaseEquil_ρTq(param_set, state.ρ, T, q_tot)\n q_eq = PhasePartition(ts_eq)\n\n # zero out the source terms\n source.ρq_tot = FT(0)\n source.ρq_liq = FT(0)\n source.ρq_ice = FT(0)\n source.ρq_rai = FT(0)\n source.ρe = FT(0)\n\n # cloud water and ice condensation/evaporation\n source.ρq_liq +=\n state.ρ *\n conv_q_vap_to_q_liq_ice(param_set, CM1M.LiquidType(), q_eq, q)\n source.ρq_ice +=\n state.ρ *\n conv_q_vap_to_q_liq_ice(param_set, CM1M.IceType(), q_eq, q)\n\n # tendencies from rain\n src_q_rai_acnv = conv_q_liq_to_q_rai(param_set, q_liq)\n src_q_rai_accr = accretion(\n param_set,\n CM1M.LiquidType(),\n CM1M.RainType(),\n q_liq,\n q_rai,\n state.ρ,\n )\n src_q_rai_evap = evaporation_sublimation(\n param_set,\n CM1M.RainType(),\n q,\n q_rai,\n state.ρ,\n T,\n )\n\n src_q_rai_tot = src_q_rai_acnv + src_q_rai_accr + src_q_rai_evap\n\n source.ρq_liq -= state.ρ * (src_q_rai_acnv + src_q_rai_accr)\n source.ρq_rai += state.ρ * src_q_rai_tot\n source.ρq_tot -= state.ρ * src_q_rai_tot\n source.ρe -=\n state.ρ * src_q_rai_tot * (_e_int_v0 - (_cv_v - _cv_d) * (T - _T_0))\n end\nend\n\nfunction main()\n # Working precision\n FT = Float64\n # DG polynomial order\n N = 4\n # Domain resolution and size\n Δx = FT(20)\n Δy = FT(1)\n Δz = FT(20)\n resolution = (Δx, Δy, Δz)\n # Domain extents\n xmax = 1500\n ymax = 10\n zmax = 1500\n # initial configuration\n wmax = FT(0.6) # max velocity of the eddy [m/s]\n θ_0 = FT(289) # init. theta value (const) [K]\n p_0 = FT(101500) # surface pressure [Pa]\n p_1000 = FT(100000) # reference pressure in theta definition [Pa]\n qt_0 = FT(7.5 * 1e-3) # init. total water specific humidity (const) [kg/kg]\n z_0 = FT(0) # surface height\n\n # time stepping\n t_ini = FT(0)\n t_end = FT(30 * 60)\n dt = FT(5)\n #CFL = FT(1.75)\n filter_freq = 1\n output_freq = 72\n interval = \"9steps\"\n\n # periodicity and boundary numbers\n periodicity_x = true\n periodicity_y = true\n periodicity_z = false\n idx_bc_left = 0\n idx_bc_right = 0\n idx_bc_front = 0\n idx_bc_back = 0\n idx_bc_bottom = 1\n idx_bc_top = 2\n\n driver_config, ode_solver_type = config_kinematic_eddy(\n FT,\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n wmax,\n θ_0,\n p_0,\n p_1000,\n qt_0,\n z_0,\n periodicity_x,\n periodicity_y,\n periodicity_z,\n idx_bc_left,\n idx_bc_right,\n idx_bc_front,\n idx_bc_back,\n idx_bc_bottom,\n idx_bc_top,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t_ini,\n t_end,\n driver_config;\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n init_on_cpu = true,\n #Courant_number = CFL,\n )\n\n model = driver_config.bl\n\n mpicomm = MPI.COMM_WORLD\n\n # get state variables indices for filtering\n ρq_liq_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_liq)\n ρq_ice_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_ice)\n ρq_rai_ind = varsindex(vars_state(model, Prognostic(), FT), :ρq_rai)\n # get aux variables indices for testing\n q_tot_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_tot)\n q_vap_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_vap)\n q_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_liq)\n q_ice_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_ice)\n q_rai_ind = varsindex(vars_state(model, Auxiliary(), FT), :q_rai)\n S_liq_ind = varsindex(vars_state(model, Auxiliary(), FT), :S_liq)\n rain_w_ind = varsindex(vars_state(model, Auxiliary(), FT), :rain_w)\n\n # filter out negative values\n cb_tmar_filter =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:ρq_liq, :ρq_ice, :ρq_rai),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n\n # output for paraview\n\n # initialize base output prefix directory from rank 0\n vtkdir = abspath(joinpath(ClimateMachine.Settings.output_dir, \"vtk\"))\n if MPI.Comm_rank(mpicomm) == 0\n mkpath(vtkdir)\n end\n MPI.Barrier(mpicomm)\n\n # vtk output\n vtkstep = [0]\n cb_vtk =\n GenericCallbacks.EveryXSimulationSteps(output_freq) do (init = false)\n out_dirname = @sprintf(\n \"microphysics_test_3_mpirank%04d_step%04d\",\n MPI.Comm_rank(mpicomm),\n vtkstep[1]\n )\n out_path_prefix = joinpath(vtkdir, out_dirname)\n @info \"doing VTK output\" out_path_prefix\n writevtk(\n out_path_prefix,\n solver_config.Q,\n solver_config.dg,\n flattenednames(vars_state(model, Prognostic(), FT)),\n solver_config.dg.state_auxiliary,\n flattenednames(vars_state(model, Auxiliary(), FT)),\n )\n vtkstep[1] += 1\n nothing\n end\n\n # output for netcdf\n boundaries = [\n FT(0) FT(0) FT(0)\n xmax ymax zmax\n ]\n interpol = ClimateMachine.InterpolationConfiguration(\n driver_config,\n boundaries,\n resolution,\n )\n dgngrps = [\n setup_dump_state_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n setup_dump_aux_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n ]\n dgn_config = ClimateMachine.DiagnosticsConfiguration(dgngrps)\n\n # call solve! function for time-integrator\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (cb_tmar_filter, cb_vtk),\n check_euclidean_distance = true,\n )\n\n # supersaturation in the model\n max_S_liq = maximum(abs.(solver_config.dg.state_auxiliary[:, S_liq_ind, :]))\n @test max_S_liq < FT(0.25)\n @test max_S_liq > FT(0)\n\n # qt < reference number\n max_q_tot = maximum(abs.(solver_config.dg.state_auxiliary[:, q_tot_ind, :]))\n @test max_q_tot < FT(0.0077)\n\n # no ice\n max_q_ice = maximum(abs.(solver_config.dg.state_auxiliary[:, q_ice_ind, :]))\n @test isequal(max_q_ice, FT(0))\n\n # q_liq ∈ reference range\n max_q_liq = maximum(solver_config.dg.state_auxiliary[:, q_liq_ind, :])\n min_q_liq = minimum(solver_config.dg.state_auxiliary[:, q_liq_ind, :])\n @test max_q_liq < FT(1e-3)\n @test abs(min_q_liq) < FT(1e-5)\n\n # q_rai ∈ reference range\n max_q_rai = maximum(solver_config.dg.state_auxiliary[:, q_rai_ind, :])\n min_q_rai = minimum(solver_config.dg.state_auxiliary[:, q_rai_ind, :])\n @test max_q_rai < FT(3e-5)\n @test abs(min_q_rai) < FT(7e-8)\n\n # terminal velocity ∈ reference range\n max_rain_w = maximum(solver_config.dg.state_auxiliary[:, rain_w_ind, :])\n min_rain_w = minimum(solver_config.dg.state_auxiliary[:, rain_w_ind, :])\n @test max_rain_w < FT(4)\n @test isequal(min_rain_w, FT(0))\nend\n\nmain()\n" }, { "path": "test/Atmos/Parameterizations/Microphysics/KinematicModel.jl", "content": "# The set-up was designed for the\n# 8th International Cloud Modelling Workshop\n# ([Muhlbauer2013](@cite))\n#\n# See chapter 2 in [Arabas2015](@cite) for setup details:\n\nusing Dates\nusing DocStringExtensions\nusing LinearAlgebra\nusing Logging\nusing MPI\nusing Printf\nusing StaticArrays\nusing Test\n\nusing ClimateMachine\nClimateMachine.init(diagnostics = \"default\")\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.Grids\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.Mesh.Topologies\nusing Thermodynamics:\n gas_constants,\n PhaseEquil,\n PhaseEquil_ρeq,\n PhasePartition_equil,\n PhasePartition,\n internal_energy,\n q_vap_saturation,\n relative_humidity,\n PhaseEquil_ρTq,\n PhaseNonEquil_ρTq,\n air_temperature,\n latent_heat_fusion,\n Liquid,\n Ice,\n supersaturation,\n vapor_specific_humidity\n\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CloudMicrophysics.Microphysics_0M\nusing CloudMicrophysics.Microphysics_1M\nimport CloudMicrophysics\nconst CM1M = CloudMicrophysics.Microphysics_1M\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet:\n R_d, cp_d, cv_d, cv_v, cv_l, cv_i, T_0, T_freeze, e_int_v0, e_int_i0, grav\n\nusing CLIMAParameters.Atmos.Microphysics\n\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, Hyperdiffusive\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n nodal_init_state_auxiliary!,\n nodal_update_auxiliary_state!,\n flux_first_order!,\n flux_second_order!,\n wavespeed,\n parameter_set,\n boundary_conditions,\n boundary_state!,\n source!\n\nimport ClimateMachine.DGMethods: DGModel\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nstruct KinematicModelConfig{FT}\n xmax::FT\n ymax::FT\n zmax::FT\n wmax::FT\n θ_0::FT\n p_0::FT\n p_1000::FT\n qt_0::FT\n z_0::FT\n periodicity_x::Bool\n periodicity_y::Bool\n periodicity_z::Bool\n idx_bc_left::Int\n idx_bc_right::Int\n idx_bc_front::Int\n idx_bc_back::Int\n idx_bc_bottom::Int\n idx_bc_top::Int\nend\n\nstruct KinematicModel{FT, PS, O, M, P, S, BC, IS, DC} <: BalanceLaw\n param_set::PS\n orientation::O\n moisture::M\n precipitation::P\n source::S\n boundarycondition::BC\n init_state_prognostic::IS\n data_config::DC\nend\n\nparameter_set(m::KinematicModel) = m.param_set\n\nfunction KinematicModel{FT}(\n ::Type{AtmosLESConfigType},\n param_set::AbstractParameterSet;\n orientation::O = FlatOrientation(),\n moisture::M = nothing,\n precipitation::P = nothing,\n source::S = nothing,\n boundarycondition::BC = nothing,\n init_state_prognostic::IS = nothing,\n data_config::DC = nothing,\n) where {FT <: AbstractFloat, O, M, P, S, BC, IS, DC}\n\n @assert param_set ≠ nothing\n @assert init_state_prognostic ≠ nothing\n\n atmos = (\n param_set,\n orientation,\n moisture,\n precipitation,\n source,\n boundarycondition,\n init_state_prognostic,\n data_config,\n )\n\n return KinematicModel{FT, typeof.(atmos)...}(atmos...)\nend\n\nvars_state(m::KinematicModel, ::Gradient, FT) = @vars()\n\nvars_state(m::KinematicModel, ::GradientFlux, FT) = @vars()\n\nfunction nodal_init_state_auxiliary!(\n m::KinematicModel,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n FT = eltype(aux)\n x, y, z = geom.coord\n dc = m.data_config\n param_set = parameter_set(m)\n\n _R_d::FT = R_d(param_set)\n _cp_d::FT = cp_d(param_set)\n _grav::FT = grav(param_set)\n\n # TODO - should R_d and cp_d here be R_m and cp_m?\n R_m, cp_m, cv_m, γ = gas_constants(param_set, PhasePartition(dc.qt_0))\n\n # Pressure profile assuming hydrostatic and constant θ and qt profiles.\n # It is done this way to be consistent with Arabas paper.\n # It's not necessarily the best way to initialize with our model variables.\n p =\n dc.p_1000 *\n (\n (dc.p_0 / dc.p_1000)^(_R_d / _cp_d) -\n _R_d / _cp_d * _grav / dc.θ_0 / R_m * (z - dc.z_0)\n )^(_cp_d / _R_d)\n\n @inbounds begin\n aux.p = p\n aux.x_coord = x\n aux.z_coord = z\n end\nend\n\nfunction init_state_prognostic!(\n m::KinematicModel,\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n args...,\n)\n m.init_state_prognostic(m, state, aux, localgeo, t, args...)\nend\nboundary_conditions(::KinematicModel) = (1, 2, 3, 4, 5, 6)\n\nfunction boundary_state!(\n ::CentralNumericalFluxSecondOrder,\n bctype,\n m::KinematicModel,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n t,\n args...,\n) end\n\n@inline function flux_second_order!(\n m::KinematicModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n) end\n\n@inline function flux_second_order!(\n m::KinematicModel,\n flux::Grad,\n state::Vars,\n τ,\n d_h_tot,\n) end\n\nfunction config_kinematic_eddy(\n FT,\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n wmax,\n θ_0,\n p_0,\n p_1000,\n qt_0,\n z_0,\n periodicity_x,\n periodicity_y,\n periodicity_z,\n idx_bc_left,\n idx_bc_right,\n idx_bc_front,\n idx_bc_back,\n idx_bc_bottom,\n idx_bc_top,\n)\n # Choose explicit solver\n ode_solver = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n kmc = KinematicModelConfig(\n FT(xmax),\n FT(ymax),\n FT(zmax),\n FT(wmax),\n FT(θ_0),\n FT(p_0),\n FT(p_1000),\n FT(qt_0),\n FT(z_0),\n Bool(periodicity_x),\n Bool(periodicity_y),\n Bool(periodicity_z),\n Int(idx_bc_left),\n Int(idx_bc_right),\n Int(idx_bc_front),\n Int(idx_bc_back),\n Int(idx_bc_bottom),\n Int(idx_bc_top),\n )\n\n # Set up the model\n model = KinematicModel{FT}(\n AtmosLESConfigType,\n param_set;\n init_state_prognostic = init_kinematic_eddy!,\n data_config = kmc,\n )\n\n config = ClimateMachine.AtmosLESConfiguration(\n \"KinematicModel\",\n N,\n resolution,\n FT(xmax),\n FT(ymax),\n FT(zmax),\n param_set,\n init_kinematic_eddy!,\n boundary = (\n (Int(idx_bc_left), Int(idx_bc_right)),\n (Int(idx_bc_front), Int(idx_bc_back)),\n (Int(idx_bc_bottom), Int(idx_bc_top)),\n ),\n periodicity = (\n Bool(periodicity_x),\n Bool(periodicity_y),\n Bool(periodicity_z),\n ),\n xmin = FT(0),\n ymin = FT(0),\n zmin = FT(0),\n model = model,\n )\n\n return config, ode_solver\nend\n" }, { "path": "test/Atmos/prog_prim_conversion/runtests.jl", "content": "module TestPrimitivePrognosticConversion\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet\nusing Test\nusing StaticArrays\nusing UnPack\n\nusing ClimateMachine\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)))\n\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.VariableTemplates\nusing Thermodynamics\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics.TestedProfiles\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws:\n prognostic_to_primitive!, primitive_to_prognostic!\nconst BL = BalanceLaws\n\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\natol_temperature = 5e-1\natol_energy = cv_d(param_set) * atol_temperature\n\nimport ClimateMachine.BalanceLaws: vars_state\n\n# Assign aux.ref_state different than state for general testing\nfunction assign!(aux, bl, nt, st::Auxiliary)\n @unpack p, ρ, e_pot, = nt\n aux.ref_state.ρ = ρ / 2\n aux.ref_state.p = p / 2\n aux.orientation.Φ = e_pot\nend\n\nfunction assign!(state, bl, nt, st::Prognostic, aux)\n @unpack u, v, w, ρ, e_kin, e_pot, T, q_pt = nt\n state.ρ = ρ\n assign!(state, bl, compressibility_model(bl), nt, st, aux)\n state.ρu = SVector(state.ρ * u, state.ρ * v, state.ρ * w)\n param_set = parameter_set(bl)\n state.energy.ρe = state.ρ * total_energy(param_set, e_kin, e_pot, T, q_pt)\n assign!(state, bl, moisture_model(bl), nt, st)\nend\nassign!(state, bl, moisture::DryModel, nt, ::Prognostic) = nothing\nassign!(state, bl, moisture::EquilMoist, nt, ::Prognostic) =\n (state.moisture.ρq_tot = state.ρ * nt.q_pt.tot)\nfunction assign!(state, bl, moisture::NonEquilMoist, nt, ::Prognostic)\n state.moisture.ρq_tot = state.ρ * nt.q_pt.tot\n state.moisture.ρq_liq = state.ρ * nt.q_pt.liq\n state.moisture.ρq_ice = state.ρ * nt.q_pt.ice\nend\n# Assign prog.ρ = aux.ref_state.ρ in anelastic1D\nassign!(state, bl, ::Compressible, nt, ::Prognostic, aux) = nothing\nfunction assign!(state, bl, ::Anelastic1D, nt, ::Prognostic, aux)\n state.ρ = aux.ref_state.ρ\nend\n\nfunction assign!(state, bl, nt, st::Primitive, aux)\n @unpack u, v, w, ρ, p = nt\n state.ρ = ρ\n state.u = SVector(u, v, w)\n state.p = p\n assign!(state, bl, compressibility_model(bl), nt, st, aux)\n assign!(state, bl, moisture_model(bl), nt, st)\nend\nassign!(state, bl, moisture::DryModel, nt, ::Primitive) = nothing\nassign!(state, bl, moisture::EquilMoist, nt, ::Primitive) =\n (state.moisture.q_tot = nt.q_pt.tot)\nfunction assign!(state, bl, moisture::NonEquilMoist, nt, ::Primitive)\n state.moisture.q_tot = nt.q_pt.tot\n state.moisture.q_liq = nt.q_pt.liq\n state.moisture.q_ice = nt.q_pt.ice\nend\n# Assign prim.p = aux.ref_state.p in anelastic1D\nassign!(state, bl, ::Compressible, nt, ::Primitive, aux) = nothing\nfunction assign!(state, bl, ::Anelastic1D, nt, ::Primitive, aux)\n state.p = aux.ref_state.p\nend\n\n@testset \"Prognostic-Primitive conversion (dry)\" begin\n FT = Float64\n compressibility = (Anelastic1D(), Compressible())\n for comp in compressibility\n physics = AtmosPhysics{FT}(\n param_set;\n moisture = DryModel(),\n compressibility = comp,\n )\n bl = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = x -> x,\n )\n vs_prog = vars_state(bl, Prognostic(), FT)\n vs_prim = vars_state(bl, Primitive(), FT)\n vs_aux = vars_state(bl, Auxiliary(), FT)\n prog_arr = zeros(varsize(vs_prog))\n prim_arr = zeros(varsize(vs_prim))\n aux_arr = zeros(varsize(vs_aux))\n prog = Vars{vs_prog}(prog_arr)\n prim = Vars{vs_prim}(prim_arr)\n aux = Vars{vs_aux}(aux_arr)\n for nt in TestedProfiles.PhaseDryProfiles(param_set, ArrayType)\n assign!(aux, bl, nt, Auxiliary())\n\n # Test prognostic_to_primitive! identity\n assign!(prog, bl, nt, Prognostic(), aux)\n prog_0 = deepcopy(parent(prog))\n prim_arr .= 0\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test all(parent(prog) .≈ prog_0)\n\n # Test primitive_to_prognostic! identity\n assign!(prim, bl, nt, Primitive(), aux)\n prim_0 = deepcopy(parent(prim))\n prog_arr .= 0\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test all(parent(prim) .≈ prim_0)\n end\n end\nend\n\n@testset \"Prognostic-Primitive conversion (EquilMoist)\" begin\n FT = Float64\n compressibility = (Compressible(),) # Anelastic1D() does not converge\n for comp in compressibility\n physics = AtmosPhysics{FT}(\n param_set;\n moisture = EquilMoist(; maxiter = 5), # maxiter=3 does not converge\n compressibility = comp,\n )\n bl = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = x -> x,\n )\n vs_prog = vars_state(bl, Prognostic(), FT)\n vs_prim = vars_state(bl, Primitive(), FT)\n vs_aux = vars_state(bl, Auxiliary(), FT)\n prog_arr = zeros(varsize(vs_prog))\n prim_arr = zeros(varsize(vs_prim))\n aux_arr = zeros(varsize(vs_aux))\n prog = Vars{vs_prog}(prog_arr)\n prim = Vars{vs_prim}(prim_arr)\n aux = Vars{vs_aux}(aux_arr)\n err_max_fwd = 0\n err_max_bwd = 0\n for nt in TestedProfiles.PhaseEquilProfiles(param_set, ArrayType)\n assign!(aux, bl, nt, Auxiliary())\n\n # Test prognostic_to_primitive! identity\n assign!(prog, bl, nt, Prognostic(), aux)\n prog_0 = deepcopy(parent(prog))\n prim_arr .= 0\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test all(parent(prog)[1:4] .≈ prog_0[1:4])\n @test isapprox(parent(prog)[5], prog_0[5]; atol = atol_energy)\n # @test all(parent(prog)[5] .≈ prog_0[5]) # fails\n @test all(parent(prog)[6] .≈ prog_0[6])\n err_max_fwd = max(abs(parent(prog)[5] .- prog_0[5]), err_max_fwd)\n\n # Test primitive_to_prognostic! identity\n assign!(prim, bl, nt, Primitive(), aux)\n prim_0 = deepcopy(parent(prim))\n prog_arr .= 0\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test all(parent(prim)[1:4] .≈ prim_0[1:4])\n # @test all(parent(prim)[5] .≈ prim_0[5]) # fails\n @test isapprox(parent(prim)[5], prim_0[5]; atol = atol_energy)\n @test all(parent(prim)[6] .≈ prim_0[6])\n err_max_bwd = max(abs(parent(prim)[5] .- prim_0[5]), err_max_bwd)\n end\n end\n # We may want/need to improve this later, so leaving debug info:\n # @show err_max_fwd\n # @show err_max_bwd\nend\n\n@testset \"Prognostic-Primitive conversion (NonEquilMoist)\" begin\n FT = Float64\n compressibility = (Anelastic1D(), Compressible())\n for comp in compressibility\n physics = AtmosPhysics{FT}(\n param_set;\n moisture = NonEquilMoist(),\n compressibility = comp,\n )\n bl = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = x -> x,\n )\n vs_prog = vars_state(bl, Prognostic(), FT)\n vs_prim = vars_state(bl, Primitive(), FT)\n vs_aux = vars_state(bl, Auxiliary(), FT)\n prog_arr = zeros(varsize(vs_prog))\n prim_arr = zeros(varsize(vs_prim))\n aux_arr = zeros(varsize(vs_aux))\n prog = Vars{vs_prog}(prog_arr)\n prim = Vars{vs_prim}(prim_arr)\n aux = Vars{vs_aux}(aux_arr)\n for nt in TestedProfiles.PhaseEquilProfiles(param_set, ArrayType)\n assign!(aux, bl, nt, Auxiliary())\n\n # Test prognostic_to_primitive! identity\n assign!(prog, bl, nt, Prognostic(), aux)\n prog_0 = deepcopy(parent(prog))\n prim_arr .= 0\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test all(parent(prog) .≈ prog_0)\n\n # Test primitive_to_prognostic! identity\n assign!(prim, bl, nt, Primitive(), aux)\n prim_0 = deepcopy(parent(prim))\n prog_arr .= 0\n primitive_to_prognostic!(bl, prog, prim, aux)\n @test !all(parent(prim) .≈ parent(prog)) # ensure not calling fallback\n prognostic_to_primitive!(bl, prim, prog, aux)\n @test all(parent(prim) .≈ prim_0)\n end\n end\nend\n\n@testset \"Prognostic-Primitive conversion (array interface)\" begin\n FT = Float64\n physics = AtmosPhysics{FT}(param_set; moisture = DryModel())\n bl = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = x -> x,\n )\n vs_prog = vars_state(bl, Prognostic(), FT)\n vs_prim = vars_state(bl, Primitive(), FT)\n vs_aux = vars_state(bl, Auxiliary(), FT)\n prog_arr = zeros(varsize(vs_prog))\n prim_arr = zeros(varsize(vs_prim))\n aux_arr = zeros(varsize(vs_aux))\n prog = Vars{vs_prog}(prog_arr)\n prim = Vars{vs_prim}(prim_arr)\n aux = Vars{vs_aux}(aux_arr)\n for nt in TestedProfiles.PhaseDryProfiles(param_set, ArrayType)\n assign!(aux, bl, nt, Auxiliary())\n\n # Test prognostic_to_primitive! identity\n assign!(prog, bl, nt, Prognostic(), aux)\n prog_0 = deepcopy(parent(prog))\n prim_arr .= 0\n BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)\n @test !all(prog_arr .≈ prim_arr) # ensure not calling fallback\n BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)\n @test all(parent(prog) .≈ prog_0)\n\n # Test primitive_to_prognostic! identity\n assign!(prim, bl, nt, Primitive(), aux)\n prim_0 = deepcopy(parent(prim))\n prog_arr .= 0\n BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)\n @test !all(prog_arr .≈ prim_arr) # ensure not calling fallback\n BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)\n @test all(parent(prim) .≈ prim_0)\n end\nend\n\nend\n" }, { "path": "test/Atmos/runtests.jl", "content": "using Test, Pkg\n\n@testset \"Atmos\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"Model\", \"prog_prim_conversion\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"Atmos/$submodule\" in ARGS ||\n \"Atmos\" in ARGS\n include_test(submodule)\n end\n end\nend\n" }, { "path": "test/BalanceLaws/runtests.jl", "content": "using Test\nusing Random\nusing StaticArrays: SVector\nRandom.seed!(1234)\n\nusing ClimateMachine.VariableTemplates: @vars, varsize, Vars\nusing ClimateMachine.BalanceLaws\nconst BL = BalanceLaws\nimport ClimateMachine.BalanceLaws: vars_state, eq_tends, prognostic_vars\n\nstruct TestBL <: BalanceLaw end\nstruct X <: AbstractPrognosticVariable end\nstruct Y <: AbstractPrognosticVariable end\nstruct F1 <: TendencyDef{Flux{FirstOrder}} end\nstruct F2 <: TendencyDef{Flux{SecondOrder}} end\nstruct S <: TendencyDef{Source} end\n\nprognostic_vars(::TestBL) = (X(), Y())\neq_tends(::X, ::TestBL, ::Flux{FirstOrder}) = (F1(),)\neq_tends(::Y, ::TestBL, ::Flux{FirstOrder}) = (F1(),)\neq_tends(::X, ::TestBL, ::Flux{SecondOrder}) = (F2(),)\neq_tends(::Y, ::TestBL, ::Flux{SecondOrder}) = (F2(),)\neq_tends(::X, ::TestBL, ::Source) = (S(),)\neq_tends(::Y, ::TestBL, ::Source) = (S(),)\n\n@testset \"BalanceLaws\" begin\n bl = TestBL()\n @test prognostic_vars(bl) == (X(), Y())\n show_tendencies(bl)\n show_tendencies(bl; include_module = true)\n show_tendencies(bl; table_complete = true)\nend\n\nvars_state(bl::TestBL, st::Prognostic, FT) = @vars begin\n ρ::FT\n ρu::SVector{3, FT}\nend\n\nvars_state(bl::TestBL, st::Auxiliary, FT) = @vars()\n\n@testset \"Prognostic-Primitive conversion (identity)\" begin\n FT = Float64\n bl = TestBL()\n vs_prog = vars_state(bl, Prognostic(), FT)\n vs_prim = vars_state(bl, Primitive(), FT)\n vs_aux = vars_state(bl, Auxiliary(), FT)\n prim_arr = zeros(varsize(vs_prim))\n prog_arr = zeros(varsize(vs_prog))\n aux_arr = zeros(varsize(vs_aux))\n\n # Test prognostic_to_primitive! identity\n prog_arr .= rand(varsize(vs_prog))\n prog_0 = deepcopy(prog_arr)\n prim_arr .= 0\n BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)\n BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)\n @test all(prog_arr .≈ prog_0)\n\n # Test primitive_to_prognostic! identity\n prim_arr .= rand(varsize(vs_prim))\n prim_0 = deepcopy(prim_arr)\n prog_arr .= 0\n BL.primitive_to_prognostic!(bl, prog_arr, prim_arr, aux_arr)\n BL.prognostic_to_primitive!(bl, prim_arr, prog_arr, aux_arr)\n @test all(prim_arr .≈ prim_0)\nend\n" }, { "path": "test/Common/CartesianDomains/runtests.jl", "content": "using ClimateMachine\n\nClimateMachine.init()\n\nusing ClimateMachine.CartesianDomains\n\n@testset \"Cartesian domains\" begin\n\n for FT in (Float64, Float32)\n domain = RectangularDomain(\n FT,\n Ne = (16, 24, 1),\n Np = 4,\n x = (0, π),\n y = (0, 1.1),\n z = (-1, 0),\n periodicity = (false, false, false),\n )\n\n @test eltype(domain) == FT\n @test domain.Ne == (x = 16, y = 24, z = 1)\n @test domain.Np == 4\n @test domain.L.x == FT(π)\n @test domain.L.y == FT(1.1)\n @test domain.L.z == FT(1)\n end\nend\n" }, { "path": "test/Common/CartesianFields/runtests.jl", "content": "using ClimateMachine\n\nClimateMachine.init()\n\nusing ClimateMachine.CartesianFields: RectangularElement, assemble\n\n@testset \"Cartesian fields\" begin\n\n Nx = 3\n Ny = 4\n Nz = 5\n\n data = rand(Nx, Ny, Nz)\n x = repeat(range(0.0, 1.0, length = Nx), 1, Ny, Nz)\n y = repeat(range(0.0, 1.1, length = Ny), Nx, 1, Nz)\n z = repeat(range(0.0, 1.2, length = Nz), Nx, Ny, 1)\n\n element = RectangularElement(data, x, y, z)\n\n @test size(element) == (Nx, Ny, Nz)\n @test maximum(element) == maximum(data)\n @test minimum(element) == minimum(data)\n @test maximum(abs, element) == maximum(abs, data)\n @test minimum(abs, element) == minimum(abs, data)\n @test element[1, 1, 1] == data[1, 1, 1]\n\n east_element = RectangularElement(2 .* data, x .+ x[end, 1, 1], y, z)\n north_element = RectangularElement(3 .* data, x, y .+ y[1, end, 1], z)\n top_element = RectangularElement(4 .* data, x, y, z .+ z[1, 1, end])\n\n west_east = assemble(Val(1), element, east_element)\n south_north = assemble(Val(2), element, north_element)\n bottom_top = assemble(Val(3), element, top_element)\n\n @test west_east.x[1, 1, 1] == x[1, 1, 1]\n @test west_east.x[end, 1, 1] == 2 * x[end, 1, 1]\n\n @test south_north.y[1, 1, 1] == y[1, 1, 1]\n @test south_north.y[1, end, 1] == 2 * y[1, end, 1]\n\n @test bottom_top.z[1, 1, 1] == z[1, 1, 1]\n @test bottom_top.z[1, 1, end] == 2 * z[1, 1, end]\n\n northeast_element =\n RectangularElement(5 .* data, x .+ x[end, 1, 1], y .+ y[1, end, 1], z)\n\n # Remember that matrix literals are transposed, so \"northeast\"\n # is the bottom right corner (for example).\n four_elements = [\n element north_element\n east_element northeast_element\n ]\n\n four_elements = reshape(four_elements, 2, 2, 1)\n\n four_way = assemble(four_elements)\n\n @test four_way.x[end, 1, 1] == west_east.x[end, 1, 1]\n @test four_way.y[1, end, 1] == south_north.y[1, end, 1]\n @test four_way.z[1, 1, end] == z[1, 1, end]\nend\n" }, { "path": "test/Common/Spectra/gcm_standalone_visual_test.jl", "content": "# Standalone test file that tests spectra visually\n#using Plots # uncomment when using the plotting code below\n\nusing ClimateMachine.ConfigTypes\n\nusing ClimateMachine.Spectra:\n compute_gaussian!,\n compute_legendre!,\n SpectralSphericalMesh,\n trans_grid_to_spherical!,\n power_spectrum_1d,\n power_spectrum_2d,\n compute_wave_numbers\nusing FFTW\n\n\ninclude(\"spherical_helper_test.jl\")\n\nFT = Float64\n# -- TEST 1: power_spectrum_1d(AtmosGCMConfigType(), var_grid, z, lat, lon, weight)\nnlats = 32\n\n# Setup grid\nsinθ, wts = compute_gaussian!(nlats)\nyarray = asin.(sinθ) .* 180 / π\nxarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0\nz = 1\n\n# Setup variable\nmass_weight = ones(Float64, length(z));\nvar_grid =\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 5.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',\n length(xarray),\n length(yarray),\n 1,\n ) +\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 10.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',\n length(xarray),\n length(yarray),\n 1,\n )\nnm_spectrum, wave_numbers = power_spectrum_1d(\n AtmosGCMConfigType(),\n var_grid,\n z,\n yarray,\n xarray,\n mass_weight,\n);\n\n\n# Check visually\nplot(wave_numbers[:, 16, 1], nm_spectrum[:, 16, 1], xlims = (0, 20))\ncontourf(var_grid[:, :, 1])\ncontourf(nm_spectrum[2:20, :, 1])\n\n# -- TEST 2: power_spectrum_gcm_2d\n# Setup grid\nsinθ, wts = compute_gaussian!(nlats)\nyarray = asin.(sinθ) .* 180 / π\nxarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0\nz = 1\n\n# Setup variable: use an example analytical P_nm function\nP_32 = sqrt(105 / 8) * (sinθ .- sinθ .^ 3)\nvar_grid =\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 3.0 * π) .* P_32',\n length(xarray),\n length(yarray),\n 1,\n )\n\nmass_weight = ones(Float64, z);\nspectrum, wave_numbers, spherical, mesh =\n power_spectrum_2d(AtmosGCMConfigType(), var_grid, mass_weight)\n\n# Grid to spherical to grid reconstruction\nreconstruction = trans_spherical_to_grid!(mesh, spherical)\n\n# Check visually\ncontourf(var_grid[:, :, 1])\ncontourf(reconstruction[:, :, 1])\ncontourf(var_grid[:, :, 1] .- reconstruction[:, :, 1])\n\n# Spectrum\ncontourf(\n collect(0:1:(mesh.num_fourier - 1))[:],\n collect(0:1:(mesh.num_spherical - 1))[:],\n (spectrum[:, :, 1])',\n xlabel = \"m\",\n ylabel = \"n\",\n)\n\n# Check magnitude\nprintln(0.5 .* sum(spectrum))\n\ndθ = π / length(wts)\ncosθ = sqrt.(1 .- sinθ .^ 2)\narea_factor = reshape(cosθ .* dθ .^ 2 / 4π, (1, length(cosθ)))\n\nprintln(sum(0.5 .* var_grid[:, :, 1] .^ 2 .* area_factor))\n\n# NB: can verify against published packages, e.g., https://github.com/jswhit/pyspharm \n" }, { "path": "test/Common/Spectra/runtests.jl", "content": "module TestSpectra\n\nusing Test\nusing FFTW\n\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Spectra\nusing ClimateMachine.Spectra:\n compute_gaussian!,\n compute_legendre!,\n SpectralSphericalMesh,\n trans_grid_to_spherical!,\n compute_wave_numbers\n\ninclude(\"spherical_helper_test.jl\")\n\n@testset \"power_spectrum_1d (GCM)\" begin\n FT = Float64\n # -- TEST 1: power_spectrum_1d(AtmosGCMConfigType(), var_grid, z, lat, lon, weight)\n nlats = 32\n\n # Setup grid\n sinθ, wts = compute_gaussian!(nlats)\n yarray = asin.(sinθ) .* 180 / π\n xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0\n z = 1\n\n # Setup variable\n mass_weight = ones(Float64, length(z))\n var_grid =\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 5.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',\n length(xarray),\n length(yarray),\n 1,\n ) +\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 10.0 * 2π) .* (yarray .* 0.0 .+ 1.0)',\n length(xarray),\n length(yarray),\n 1,\n )\n nm_spectrum, wave_numbers = power_spectrum_1d(\n AtmosGCMConfigType(),\n var_grid,\n z,\n yarray,\n xarray,\n mass_weight,\n )\n\n nm_spectrum_ = nm_spectrum[:, 10, 1]\n var_grid_ = var_grid[:, 10, 1]\n sum_spec = sum(nm_spectrum_)\n sum_grid = sum(var_grid_ .^ 2) / length(var_grid_)\n\n sum_res = (sum_spec - sum_grid) / sum_grid\n\n @test sum_res < 0.1\nend\n\n@testset \"power_spectrum_2d (GCM)\" begin\n # -- TEST 2: power_spectrum_2d\n # Setup grid\n FT = Float64\n nlats = 32\n sinθ, wts = compute_gaussian!(nlats)\n cosθ = sqrt.(1 .- sinθ .^ 2)\n yarray = asin.(sinθ) .* 180 / π\n xarray = 180.0 ./ nlats * collect(FT, 1:1:(2nlats))[:] .- 180.0\n z = 1\n\n # Setup variable: use an example analytical P_nm function\n P_32 = sqrt(105 / 8) * (sinθ .- sinθ .^ 3)\n var_grid =\n 1.0 * reshape(\n sin.(xarray / xarray[end] * 3.0 * π) .* P_32',\n length(xarray),\n length(yarray),\n 1,\n )\n\n mass_weight = ones(Float64, z)\n spectrum, wave_numbers, spherical, mesh =\n power_spectrum_2d(AtmosGCMConfigType(), var_grid, mass_weight)\n\n # Grid to spherical to grid reconstruction\n reconstruction = trans_spherical_to_grid!(mesh, spherical)\n\n sum_spec = sum((0.5 * spectrum))\n dθ = π / length(wts)\n area_factor = reshape(cosθ .* dθ .^ 2 / 4π, (1, length(cosθ)))\n\n sum_grid = sum(0.5 .* var_grid[:, :, 1] .^ 2 .* area_factor) # scaled to average over Earth's area (units: m2/s2)\n sum_reco = sum(0.5 .* reconstruction[:, :, 1] .^ 2 .* area_factor)\n\n sum_res_1 = (sum_spec - sum_grid) / sum_grid\n sum_res_2 = (sum_reco - sum_grid) / sum_grid\n\n @test abs(sum_res_1) < 0.1\n @test abs(sum_res_2) < 0.1\nend\n\nend\n" }, { "path": "test/Common/Spectra/spherical_helper_test.jl", "content": "# additional helper functions for spherical harmonics spectra\n\n\"\"\"\n TransSphericalToGrid!(mesh, snm )\n\nTransforms a variable expressed in spherical harmonics (var_spherical[num_fourier+1, num_spherical+1]) onto a Gaussian grid (pfield[nλ, nθ])\n\n[THIS IS USED FOR TESTING ONLY]\n\n With F_{m,n} = (-1)^m F_{-m,n}*\n P_{m,n} = (-1)^m P_{-m,n}\n\n F(λ, η) = ∑_{m= -N}^{N} ∑_{n=|m|}^{N} F_{m,n} P_{m,n}(η) e^{imλ}\n = ∑_{m= 0}^{N} ∑_{n=m}^{N} F_{m,n} P_{m,n} e^{imλ} + ∑_{m= 1}^{N} ∑_{n=m}^{N} F_{-m,n} P_{-m,n} e^{-imλ}\n\n Here η = sinθ, N = num_fourier, and denote\n ! extra coeffients in snm n > N are not used.\n\n ∑_{n=m}^{N} F_{m,n} P_{m,n} = g_{m}(η) m = 1, ... N\n ∑_{n=m}^{N} F_{m,n} P_{m,n}/2.0 = g_{m}(η) m = 0\n\n We have\n\n F(λ, η) = ∑_{m= 0}^{N} g_{m}(η) e^{imλ} + ∑_{m= 0}^{N} g_{m}(η)* e^{-imλ}\n = 2real{ ∑_{m= 0}^{N} g_{m}(η) e^{imλ} }\n\n snm = F_{m,n} # Complex{Float64} [num_fourier+1, num_spherical+1]\n qnm = P_{m,n,η} # Float64[num_fourier+1, num_spherical+1, nθ]\n fourier_g = g_{m, η} # Complex{Float64} nλ×nθ with padded 0s fourier_g[num_fourier+2, :] == 0.0\n pfiled = F(λ, η) # Float64[nλ, nθ]\n\n ! use all spherical harmonic modes\n\n\n# Arguments\n- mesh: struct with mesh information\n- snm: spherical variable\n\n# References\n- Ehrendorfer, M., Spectral Numerical Weather Prediction Models, Appendix B, Society for Industrial and Applied Mathematics, 2011\n\"\"\"\nfunction trans_spherical_to_grid!(mesh, snm)\n num_fourier, num_spherical = mesh.num_fourier, mesh.num_spherical\n nλ, nθ, nd = mesh.nλ, mesh.nθ, mesh.nd\n\n qnm = mesh.qnm\n\n fourier_g = mesh.var_fourier .* 0.0\n fourier_s = mesh.var_fourier .* 0.0\n\n @assert(nθ % 2 == 0)\n nθ_half = div(nθ, 2)\n for m in 1:(num_fourier + 1)\n for n in m:num_spherical\n snm_t = transpose(snm[m, n, :, 1:nθ_half]) #snm[m,n, :] is complex number\n if (n - m) % 2 == 0\n fourier_s[m, 1:nθ_half, :] .+=\n qnm[m, n, 1:nθ_half] .* sum(snm_t, dims = 1) #even function part\n else\n fourier_s[m, (nθ_half + 1):nθ, :] .+=\n qnm[m, n, 1:nθ_half] .* sum(snm_t, dims = 1) #odd function part\n end\n end\n end\n fourier_g[:, 1:nθ_half, :] .=\n fourier_s[:, 1:nθ_half, :] .+ fourier_s[:, (nθ_half + 1):nθ, :]\n fourier_g[:, nθ:-1:(nθ_half + 1), :] .=\n fourier_s[:, 1:nθ_half, :] .- fourier_s[:, (nθ_half + 1):nθ, :] # this got ignored...\n\n fourier_g[1, :, :] ./= 2.0\n pfield = zeros(Float64, nλ, nθ, nd)\n for j in 1:nθ\n pfield[:, j, :] .= 2.0 * nλ * real.(ifft(fourier_g[:, j, :], 1)) #fourier for the first dimension\n end\n return pfield\nend\n" }, { "path": "test/Common/runtests.jl", "content": "using Test, Pkg\n\n@testset \"Common\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"CartesianDomains\", \"CartesianFields\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"Common/$submodule\" in ARGS ||\n \"Common\" in ARGS\n include_test(submodule)\n end\n end\n\nend\n" }, { "path": "test/Diagnostics/Debug/test_statecheck.jl", "content": "using Test\n\n# # State debug statistics\n#\n# Set up a basic environment\nusing MPI\nusing StaticArrays\nusing Random\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.StateCheck\nusing ClimateMachine.GenericCallbacks\n\n@testset \"$(@__FILE__)\" begin\n\n ClimateMachine.init()\n FT = Float64\n\n # Define some dummy vector and tensor abstract variables with associated types\n # and dimensions\n F1 = @vars begin\n ν∇u::SMatrix{3, 2, FT, 6}\n κ∇θ::SVector{3, FT}\n end\n F2 = @vars begin\n u::SVector{2, FT}\n θ::SVector{1, FT}\n end\n\n # Create ```MPIStateArray``` variables with arrays to hold elements of the \n # vectors and tensors\n Q1 = MPIStateArray{Float32, F1}(\n MPI.COMM_WORLD,\n ClimateMachine.array_type(),\n 4,\n 9,\n 8,\n )\n Q2 = MPIStateArray{Float64, F2}(\n MPI.COMM_WORLD,\n ClimateMachine.array_type(),\n 4,\n 3,\n 8,\n )\n\n # ### Create a call-back\n cb = ClimateMachine.StateCheck.sccreate(\n [(Q1, \"My gradients\"), (Q2, \"My fields\")],\n 1;\n prec = 15,\n )\n\n # ### Invoke the call-back\n # Compare on local \"realdata\", fill via broadcast to keep GPU happy.\n Q1.data = rand(MersenneTwister(0), Float32, size(Q1.data))\n Q2.data = rand(MersenneTwister(0), Float64, size(Q2.data))\n Q1.realdata .= Q1.data\n Q2.realdata .= Q2.data\n GenericCallbacks.init!(cb, nothing, nothing, nothing, nothing)\n GenericCallbacks.call!(cb, nothing, nothing, nothing, nothing)\n\n # ### Check with reference values\n include(\"test_statecheck_refvals.jl\")\n\n # This should return true (MPI rank != 0 always returns true)\n test_stat_01 = ClimateMachine.StateCheck.scdocheck(cb, (varr, parr))\n\n # This should return false (MPI rank != 0 always returns true)\n varr[1][3] = varr[1][3] * 10.0\n test_stat_02 = ClimateMachine.StateCheck.scdocheck(cb, (varr, parr))\n if MPI.Comm_rank(MPI.COMM_WORLD) != 0\n test_stat_02 = false\n end\n\n test_stat = test_stat_01 && !test_stat_02\n\n @test test_stat\n\nend\n" }, { "path": "test/Diagnostics/Debug/test_statecheck_refvals.jl", "content": "#\n# Example of a set of reference values that StateCheck will use to test current vlaues against.\n# These are generated by the function scprintref() and parsed by the scdocheck()\n#\n\n#! format: off\nvarr = [\n [ \"My gradients\", \"ν∇u[1]\", 1.34348869323730468750e-04, 9.84732866287231445313e-01, 5.23545503616333007813e-01, 3.08209930764271777814e-01 ],\n [ \"My gradients\", \"ν∇u[2]\", 1.16317868232727050781e-01, 9.92088317871093750000e-01, 4.83800649642944335938e-01, 2.83350456014221541157e-01 ],\n [ \"My gradients\", \"ν∇u[3]\", 1.05845928192138671875e-03, 9.51775908470153808594e-01, 4.65474426746368408203e-01, 2.73615551085745090099e-01 ],\n [ \"My gradients\", \"ν∇u[4]\", 5.97668886184692382813e-02, 9.68048095703125000000e-01, 5.42618036270141601563e-01, 2.81570862027933854765e-01 ],\n [ \"My gradients\", \"ν∇u[5]\", 8.31030607223510742188e-02, 9.35931921005249023438e-01, 5.05405902862548828125e-01, 2.46073509972619536290e-01 ],\n [ \"My gradients\", \"ν∇u[6]\", 3.09681892395019531250e-02, 9.98341441154479980469e-01, 4.54375565052032470703e-01, 3.09461067853178561915e-01 ],\n [ \"My gradients\", \"κ∇θ[1]\", 8.47448110580444335938e-02, 9.94180679321289062500e-01, 5.27157366275787353516e-01, 2.92455951648181833313e-01 ],\n [ \"My gradients\", \"κ∇θ[2]\", 1.20514631271362304688e-02, 9.93527650833129882813e-01, 4.71063584089279174805e-01, 2.96449027197666359346e-01 ],\n [ \"My gradients\", \"κ∇θ[3]\", 8.14980268478393554688e-02, 9.55443382263183593750e-01, 5.05038917064666748047e-01, 2.77201022741208891187e-01 ],\n [ \"My fields\", \"u[1]\", 3.53445491472876849315e-02, 9.73118774570230771204e-01, 4.53566376673364857197e-01, 3.28622448223506280485e-01 ],\n [ \"My fields\", \"u[2]\", 4.23016659320296639635e-02, 9.67799553619200114696e-01, 5.40307230728526155517e-01, 2.94603153327737565803e-01 ],\n [ \"My fields\", \"θ[1]\", 6.23675581701588210848e-02, 9.76550123041147521974e-01, 5.16046312418334207628e-01, 2.81983244164903890105e-01 ],\n]\nparr = [\n [ \"My gradients\", \"ν∇u[1]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[2]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[3]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[4]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[5]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[6]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[1]\", 16, 16, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[2]\", 16, 16, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[3]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"u[1]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"u[2]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"θ[1]\", 16, 16, 16, 0 ],\n]\n#! format: on\n" }, { "path": "test/Diagnostics/diagnostic_fields_test.jl", "content": "using Test, MPI\nusing Random\nusing StaticArrays\nusing Test\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport ClimateMachine.Mesh.Grids: _x1, _x2, _x3\nimport ClimateMachine.BalanceLaws: vars_state\nimport ClimateMachine.VariableTemplates.varsindex\n\n# ------------------------ Description ------------------------- #\n# 1) Dry Rising Bubble (circular potential temperature perturbation)\n# 2) Boundaries - `All Walls` : Impenetrable(FreeSlip())\n# Laterally periodic\n# 3) Domain - 2500m[horizontal] x 2500m[horizontal] x 2500m[vertical]\n# 4) Timeend - 1000s\n# 5) Mesh Aspect Ratio (Effective resolution) 1:1\n# 7) Overrides defaults for\n# `init_on_cpu`\n# `solver_type`\n# `sources`\n# `C_smag`\n# 8) Default settings can be found in `src/Driver/Configurations.jl`\n# ------------------------ Description ------------------------- #\nfunction init_risingbubble!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n FT = eltype(state)\n param_set = parameter_set(bl)\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n γ::FT = c_p / c_v\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n\n xc::FT = 1250\n yc::FT = 1250\n zc::FT = 1000\n r = sqrt((x - xc)^2 + (y - yc)^2 + (z - zc)^2)\n rc::FT = 500\n θ_ref::FT = 300\n Δθ::FT = 0\n\n if r <= rc\n Δθ = FT(5) * cospi(r / rc / 2)\n end\n\n #Perturbed state:\n θ = θ_ref + Δθ # potential temperature\n π_exner = FT(1) - _grav / (c_p * θ) * z # exner pressure\n ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas) # density\n q_tot = FT(0)\n ts = PhaseEquil_ρθq(param_set, ρ, θ, q_tot)\n q_pt = PhasePartition(ts)\n\n ρu = SVector(FT(0), FT(0), FT(0))\n\n #State (prognostic) variable assignment\n e_kin = FT(0)\n e_pot = gravitational_potential(bl.orientation, aux)\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts)\n state.ρ = ρ\n state.ρu = ρu\n state.energy.ρe = ρe_tot\n state.moisture.ρq_tot = ρ * q_pt.tot\nend\n\nfunction config_risingbubble(FT, N, resolution, xmax, ymax, zmax)\n\n # Set up the model\n T_profile = DryAdiabaticProfile{FT}(param_set)\n C_smag = FT(0.23)\n ref_state = HydrostaticState(T_profile)\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = SmagorinskyLilly{FT}(C_smag),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = init_risingbubble!,\n source = (Gravity(),),\n )\n\n # Problem configuration\n config = ClimateMachine.AtmosLESConfiguration(\n \"DryRisingBubble\",\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n param_set,\n init_risingbubble!,\n model = model,\n )\n return config\nend\n\nfunction config_diagnostics(driver_config)\n interval = \"10000steps\"\n dgngrp = setup_atmos_default_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n )\n return ClimateMachine.DiagnosticsConfiguration([dgngrp])\nend\n#-------------------------------------------------------------------------\nfunction run_brick_diagostics_fields_test()\n DA = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n root = 0\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n toler = Dict(Float64 => 1e-8, Float32 => 1e-4)\n # Working precision\n for FT in (Float32, Float64)\n # DG polynomial order\n N = (4, 5)\n # Domain resolution and size\n Δh = FT(50)\n Δv = FT(50)\n resolution = (Δh, Δh, Δv)\n # Domain extents\n xmax = FT(2500)\n ymax = FT(2500)\n zmax = FT(2500)\n # Simulation time\n t0 = FT(0)\n timeend = FT(1000)\n\n # Courant number\n CFL = FT(20)\n\n driver_config = config_risingbubble(FT, N, resolution, xmax, ymax, zmax)\n\n # Choose explicit multirate solver\n ode_solver_type = ClimateMachine.MultirateSolverType(\n fast_model = AtmosAcousticGravityLinearModel,\n slow_method = LSRK144NiegemannDiehlBusch,\n fast_method = LSRK144NiegemannDiehlBusch,\n timestep_ratio = 10,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config;\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFL,\n )\n #------------------------------------------------------------------------\n\n model = driver_config.bl\n Q = solver_config.Q\n dg = solver_config.dg\n grid = dg.grid\n vgeo = grid.vgeo\n Nel = size(Q.realdata, 3)\n Npl = size(Q.realdata, 1)\n\n ind = [\n varsindex(vars_state(model, Prognostic(), FT), :ρ)\n varsindex(vars_state(model, Prognostic(), FT), :ρu)\n ]\n _ρ, _ρu, _ρv, _ρw = ind[1], ind[2], ind[3], ind[4]\n\n\n x1 = view(vgeo, :, _x1, 1:Nel)\n x2 = view(vgeo, :, _x2, 1:Nel)\n x3 = view(vgeo, :, _x3, 1:Nel)\n\n fcn0(x, y, z, xmax, ymax, zmax) =\n sin.(pi * x ./ xmax) .* cos.(pi * y ./ ymax) .* cos.(pi * z ./ zmax) # sample function\n\n fcnx(x, y, z, xmax, ymax, zmax) =\n cos.(pi * x ./ xmax) .* cos.(pi * y ./ ymax) .*\n cos.(pi * z ./ zmax) .* pi ./ xmax # ∂/∂x\n fcny(x, y, z, xmax, ymax, zmax) =\n -sin.(pi * x ./ xmax) .* sin.(pi * y ./ ymax) .*\n cos.(pi * z ./ zmax) .* pi ./ ymax # ∂/∂y\n fcnz(x, y, z, xmax, ymax, zmax) =\n -sin.(pi * x ./ xmax) .* cos.(pi * y ./ ymax) .*\n sin.(pi * z ./ zmax) .* pi ./ zmax # ∂/∂z\n\n Q.data[:, _ρ, 1:Nel] .= 1.0 .+ fcn0(x1, x2, x3, xmax, ymax, zmax) * 5.0\n Q.data[:, _ρu, 1:Nel] .=\n Q.data[:, _ρ, 1:Nel] .* fcn0(x1, x2, x3, xmax, ymax, zmax)\n Q.data[:, _ρv, 1:Nel] .=\n Q.data[:, _ρ, 1:Nel] .* fcn0(x1, x2, x3, xmax, ymax, zmax)\n Q.data[:, _ρw, 1:Nel] .=\n Q.data[:, _ρ, 1:Nel] .* fcn0(x1, x2, x3, xmax, ymax, zmax)\n #-----------------------------------------------------------------------\n vgrad = Diagnostics.VectorGradients(dg, Q)\n vort = Diagnostics.Vorticity(dg, vgrad)\n #----------------------------------------------------------------------------\n Ω₁_exact =\n fcny(x1, x2, x3, xmax, ymax, zmax) -\n fcnz(x1, x2, x3, xmax, ymax, zmax)\n Ω₂_exact =\n fcnz(x1, x2, x3, xmax, ymax, zmax) -\n fcnx(x1, x2, x3, xmax, ymax, zmax)\n Ω₃_exact =\n fcnx(x1, x2, x3, xmax, ymax, zmax) -\n fcny(x1, x2, x3, xmax, ymax, zmax)\n\n err = zeros(FT, 12)\n\n err[1] = maximum(abs.(fcnx(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₁u₁))\n err[2] = maximum(abs.(fcny(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₂u₁))\n err[3] = maximum(abs.(fcnz(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₃u₁))\n\n err[4] = maximum(abs.(fcnx(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₁u₂))\n err[5] = maximum(abs.(fcny(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₂u₂))\n err[6] = maximum(abs.(fcnz(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₃u₂))\n\n err[7] = maximum(abs.(fcnx(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₁u₃))\n err[8] = maximum(abs.(fcny(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₂u₃))\n err[9] = maximum(abs.(fcnz(x1, x2, x3, xmax, ymax, zmax) - vgrad.∂₃u₃))\n\n err[10] = maximum(abs.(vort.Ω₁ - Ω₁_exact))\n err[11] = maximum(abs.(vort.Ω₂ - Ω₂_exact))\n err[12] = maximum(abs.(vort.Ω₃ - Ω₃_exact))\n\n errg = MPI.Allreduce(err, max, mpicomm)\n @test maximum(errg) < toler[FT]\n end\nend\n#----------------------------------------------------------------------------\n@testset \"Diagnostics Fields tests\" begin\n run_brick_diagostics_fields_test()\nend\n#------------------------------------------------\n" }, { "path": "test/Diagnostics/dm_tests.jl", "content": "using Dates\nusing FileIO\nusing KernelAbstractions\nusing MPI\nusing NCDatasets\nusing Printf\nusing Random\nusing StaticArrays\nusing Test\n\nusing ClimateMachine\nClimateMachine.init(diagnostics = \"1steps\")\nusing ClimateMachine.Atmos\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.DiagnosticsMachine\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.MPIStateArrays\nusing Thermodynamics\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# Need to import these to define new diagnostic variables and groups.\nimport ..DiagnosticsMachine:\n Settings,\n dv_name,\n dv_attrib,\n dv_args,\n dv_project,\n dv_scale,\n dv_PointwiseDiagnostic,\n dv_HorizontalAverage\n\n# Define some new diagnostic variables.\n@horizontal_average(\n AtmosLESConfigType,\n yvel,\n) do (atmos::AtmosModel, states::States, curr_time, cache)\n states.prognostic.ρu[2] / states.prognostic.ρ\nend\n\n@pointwise_diagnostic(\n AtmosLESConfigType,\n zvel,\n) do (atmos::AtmosModel, states::States, curr_time, cache)\n states.prognostic.ρu[3] / states.prognostic.ρ\nend\n\n# Define a new diagnostics group with some pre-defined diagnostic\n# variables as well as the new ones above.\n@diagnostics_group(\n \"DMTest\",\n AtmosLESConfigType,\n Nothing,\n (_...) -> nothing,\n NoInterpolation,\n u,\n v,\n w,\n rho,\n yvel,\n zvel,\n)\n\n# Make sure DiagnosticsMachine did what it was supposed to do.\n@testset \"DiagnosticsMachine interface\" begin\n @test_throws UndefVarError ALESCT_HA_foo()\n yv = ALESCT_HA_yvel()\n @test yv isa HorizontalAverage\n yvname = dv_name(AtmosLESConfigType(), yv)\n @test yvname == \"yvel\"\n @test yvname ∈\n keys(DiagnosticsMachine.AllDiagnosticVars[AtmosLESConfigType])\n\n @test_throws UndefVarError ALESCT_PD_bar()\n zv = ALESCT_PD_zvel()\n @test zv isa PointwiseDiagnostic\n zvname = dv_name(AtmosLESConfigType(), zv)\n @test zvname == \"zvel\"\n @test zvname ∈\n keys(DiagnosticsMachine.AllDiagnosticVars[AtmosLESConfigType])\nend\n\n# Set up a simple experiment to run the diagnostics group.\n\ninclude(\"sin_init.jl\")\n\nfunction main()\n FT = Float64\n\n # DG polynomial order\n N = 4\n\n # Domain resolution and size\n Δh = FT(25)\n Δv = FT(25)\n resolution = (Δh, Δh, Δv)\n\n xmax = FT(500)\n ymax = FT(500)\n zmax = FT(500)\n\n t0 = FT(0)\n dt = FT(0.01)\n timeend = dt\n\n driver_config = ClimateMachine.AtmosLESConfiguration(\n \"DMTest\",\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n param_set,\n init_sin_test!,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n init_on_cpu = true,\n )\n dm_dgngrp = DMTest(\"1steps\", driver_config.name)\n dgn_config = ClimateMachine.DiagnosticsConfiguration([dm_dgngrp])\n\n ClimateMachine.invoke!(solver_config, diagnostics_config = dgn_config)\n\n # Check the output from the diagnostics group.\n @testset \"DiagnosticsMachine correctness\" begin\n mpicomm = solver_config.mpicomm\n mpirank = MPI.Comm_rank(mpicomm)\n if mpirank == 0\n nm = driver_config.name * \"_DMTest.nc\"\n ds = Dataset(joinpath(ClimateMachine.Settings.output_dir, nm), \"r\")\n ds_u = ds[\"u\"][:]\n ds_yvel = ds[\"yvel\"][:]\n ds_zvel = ds[\"zvel\"][:]\n close(ds)\n end\n Q =\n array_device(solver_config.Q) isa CPU ? solver_config.Q :\n Array(solver_config.Q)\n havg_rho = compute_havg(solver_config, view(Q, :, 1:1, :))\n havg_u = compute_havg(solver_config, view(Q, :, 2:2, :))\n v = view(Q, :, 3:3, :) ./ view(Q, :, 1:1, :)\n havg_v = compute_havg(solver_config, v)\n if mpirank == 0\n havg_u ./= havg_rho\n @test all(ds_u[:, 3] .≈ havg_u)\n @test all(ds_yvel[:, 3] .≈ havg_v)\n\n realelems = solver_config.dg.grid.topology.realelems\n w = view(Q, :, 4, realelems) ./ view(Q, :, 1, realelems)\n @test all(ds_zvel[:, :, 3] .≈ w)\n end\n end\n\n nothing\nend\n\nfunction compute_havg(solver_config, field)\n mpicomm = solver_config.mpicomm\n mpirank = MPI.Comm_rank(mpicomm)\n grid = solver_config.dg.grid\n grid_info = basic_grid_info(grid)\n topl_info = basic_topology_info(grid.topology)\n Nqh = grid_info.Nqh\n Nqk = grid_info.Nqk\n nvertelem = topl_info.nvertelem\n nhorzrealelem = topl_info.nhorzrealelem\n\n function arrange_array(A, dim = :)\n A = array_device(A) isa CPU ? A : Array(A)\n reshape(\n view(A, :, dim, grid.topology.realelems),\n Nqh,\n Nqk,\n nvertelem,\n nhorzrealelem,\n )\n end\n\n field = arrange_array(field)\n MH = arrange_array(grid.vgeo, grid.MHid)\n full_field = MPI.Reduce!(sum(field .* MH, dims = (1, 4))[:], +, 0, mpicomm)\n full_MH = MPI.Reduce!(sum(MH, dims = (1, 4))[:], +, 0, mpicomm)\n if mpirank == 0\n return full_field ./ full_MH\n else\n return nothing\n end\nend\n\nmain()\n" }, { "path": "test/Diagnostics/runtests.jl", "content": "using MPI, Test\n\ninclude(joinpath(\"..\", \"testhelpers.jl\"))\n\n@testset \"Diagnostics\" begin\n runmpi(joinpath(@__DIR__, \"sin_test.jl\"), ntasks = 2)\n runmpi(joinpath(@__DIR__, \"dm_tests.jl\"), ntasks = 2)\n runmpi(joinpath(@__DIR__, \"Debug/test_statecheck.jl\"), ntasks = 2)\nend\n" }, { "path": "test/Diagnostics/sin_init.jl", "content": "function init_sin_test!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n z = FT(z)\n _grav::FT = grav(param_set)\n _MSLP::FT = MSLP(param_set)\n\n # These constants are those used by Stevens et al. (2005)\n qref = FT(9.0e-3)\n q_pt_sfc = PhasePartition(qref)\n Rm_sfc = FT(gas_constant_air(param_set, q_pt_sfc))\n T_sfc = FT(292.5)\n P_sfc = _MSLP\n\n # Specify moisture profiles\n q_liq = FT(0)\n q_ice = FT(0)\n zb = FT(600) # initial cloud bottom\n zi = FT(840) # initial cloud top\n dz_cloud = zi - zb\n q_liq_peak = FT(0.00045) # cloud mixing ratio at z_i\n\n if z > zb && z <= zi\n q_liq = (z - zb) * q_liq_peak / dz_cloud\n end\n\n if z <= zi\n θ_liq = FT(289.0)\n q_tot = qref\n else\n θ_liq = FT(297.5) + (z - zi)^(FT(1 / 3))\n q_tot = FT(1.5e-3)\n end\n\n w = FT(10 + 0.5 * sin(2 * π * ((x / 1500) + (y / 1500))))\n u = (5 + 2 * sin(2 * π * ((x / 1500) + (y / 1500))))\n v = FT(5 + 2 * sin(2 * π * ((x / 1500) + (y / 1500))))\n\n # Pressure\n H = Rm_sfc * T_sfc / _grav\n p = P_sfc * exp(-z / H)\n\n # Density, Temperature\n ts = PhaseEquil_pθq(param_set, p, θ_liq, q_tot)\n #ρ = air_density(ts)\n ρ = one(FT)\n\n e_kin = FT(1 / 2) * FT((u^2 + v^2 + w^2))\n e_pot = _grav * z\n E = ρ * total_energy(e_kin, e_pot, ts)\n\n state.ρ = ρ\n state.ρu = SVector(ρ * u, ρ * v, ρ * w)\n state.energy.ρe = E\n state.moisture.ρq_tot = ρ * q_tot\nend\n" }, { "path": "test/Diagnostics/sin_test.jl", "content": "using Dates\nusing FileIO\nusing MPI\nusing NCDatasets\nusing Printf\nusing Random\nusing StaticArrays\nusing Test\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing Thermodynamics\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Writers\nusing ClimateMachine.GenericCallbacks\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\ninclude(\"sin_init.jl\")\n\nfunction config_sin_test(FT, N, resolution, xmax, ymax, zmax)\n config = ClimateMachine.AtmosLESConfiguration(\n \"Diagnostics SIN test\",\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n param_set,\n init_sin_test!,\n )\n\n return config\nend\n\nfunction config_diagnostics(driver_config)\n interval = \"100steps\"\n dgngrp = setup_atmos_default_diagnostics(\n AtmosLESConfigType(),\n interval,\n replace(driver_config.name, \" \" => \"_\"),\n writer = NetCDFWriter(),\n )\n return ClimateMachine.DiagnosticsConfiguration([dgngrp])\nend\n\nfunction main()\n # Disable driver diagnostics as we're testing it here\n ClimateMachine.Settings.diagnostics = \"never\"\n\n FT = Float64\n\n # DG polynomial order\n N = 4\n\n # Domain resolution and size\n Δh = FT(50)\n Δv = FT(20)\n resolution = (Δh, Δh, Δv)\n\n xmax = FT(1500)\n ymax = FT(1500)\n zmax = FT(1500)\n\n t0 = FT(0)\n dt = FT(0.01)\n timeend = dt\n\n driver_config = config_sin_test(FT, N, resolution, xmax, ymax, zmax)\n\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK54CarpenterKennedy,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n init_on_cpu = true,\n )\n dgn_config = config_diagnostics(driver_config)\n\n mpicomm = solver_config.mpicomm\n dg = solver_config.dg\n Q = solver_config.Q\n solver = solver_config.solver\n\n outdir = mktempdir()\n currtime = ODESolvers.gettime(solver)\n starttime = replace(string(now()), \":\" => \".\")\n Diagnostics.init(mpicomm, param_set, dg, Q, starttime, outdir, false)\n GenericCallbacks.init!(\n dgn_config.groups[1],\n nothing,\n nothing,\n nothing,\n currtime,\n )\n\n ClimateMachine.invoke!(solver_config)\n\n # Check results\n mpirank = MPI.Comm_rank(mpicomm)\n if mpirank == 0\n dgngrp = dgn_config.groups[1]\n nm = @sprintf(\"%s_%s.nc\", dgngrp.out_prefix, dgngrp.name)\n ds = NCDataset(joinpath(outdir, nm), \"r\")\n ds_u = ds[\"u\"][:]\n ds_cov_w_u = ds[\"cov_w_u\"][:]\n N = size(ds_u, 1)\n err = 0\n err1 = 0\n for i in 1:N\n u = ds_u[i]\n cov_w_u = ds_cov_w_u[i]\n err += (cov_w_u - 0.5)^2\n err1 += (u - 5)^2\n end\n close(ds)\n err = sqrt(err / N)\n @test err1 <= 1e-16\n @test err <= 2e-15\n end\nend\n\nmain()\n" }, { "path": "test/Driver/cr_unit_tests.jl", "content": "using MPI\nusing Printf\nusing StaticArrays\nusing Test\nimport KernelAbstractions: CPU\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.ConfigTypes\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Grids\nusing ClimateMachine.ODESolvers\nimport ClimateMachine.MPIStateArrays: array_device\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n # callable to set initial conditions\n FT = eltype(state)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n param_set = parameter_set(bl)\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n return nothing\nend\n\nfunction main()\n FT = Float64\n\n # DG polynomial order\n N = 4\n\n # Domain resolution\n nelem_horz = 4\n nelem_vert = 6\n resolution = (nelem_horz, nelem_vert)\n\n t0 = FT(0)\n timeend = FT(1800)\n # Timestep size (s)\n dt = FT(600)\n\n ode_solver_type = ClimateMachine.MISSolverType(;\n splitting_type = ClimateMachine.SlowFastSplitting(),\n nsubsteps = (20,),\n )\n\n setup = AcousticWaveSetup{FT}()\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n ref_state = HydrostaticState(T_profile)\n turbulence = ConstantDynamicViscosity(FT(0))\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = turbulence,\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n\n driver_config = ClimateMachine.AtmosGCMConfiguration(\n \"Checkpoint unit tests\",\n N,\n resolution,\n setup.domain_height,\n param_set,\n setup;\n model = model,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n )\n\n isdir(ClimateMachine.Settings.checkpoint_dir) ||\n mkpath(ClimateMachine.Settings.checkpoint_dir)\n\n @testset \"Checkpoint/restart unit tests\" begin\n rm_checkpoint(\n ClimateMachine.Settings.checkpoint_dir,\n solver_config.name,\n solver_config.mpicomm,\n 0,\n )\n write_checkpoint(\n solver_config,\n ClimateMachine.Settings.checkpoint_dir,\n solver_config.name,\n solver_config.mpicomm,\n 0,\n )\n\n nm = replace(solver_config.name, \" \" => \"_\")\n cname = @sprintf(\n \"%s_checkpoint_mpirank%04d_num%04d.jld2\",\n nm,\n MPI.Comm_rank(solver_config.mpicomm),\n 0,\n )\n cfull = joinpath(ClimateMachine.Settings.checkpoint_dir, cname)\n @test isfile(cfull)\n\n s_Q, s_aux, s_t = try\n read_checkpoint(\n ClimateMachine.Settings.checkpoint_dir,\n nm,\n driver_config.array_type,\n solver_config.mpicomm,\n 0,\n )\n catch\n (nothing, nothing, nothing)\n end\n @test s_Q !== nothing\n @test s_aux !== nothing\n @test s_t !== nothing\n if Array ∉ typeof(s_Q).parameters\n s_Q = Array(s_Q)\n s_aux = Array(s_aux)\n end\n\n dg = solver_config.dg\n Q = solver_config.Q\n if array_device(Q) isa CPU\n h_Q = Q.realdata\n h_aux = dg.state_auxiliary.realdata\n else\n h_Q = Array(Q.realdata)\n h_aux = Array(dg.state_auxiliary.realdata)\n end\n t = ODESolvers.gettime(solver_config.solver)\n\n @test h_Q == s_Q\n @test h_aux == s_aux\n @test t == s_t\n\n rm_checkpoint(\n ClimateMachine.Settings.checkpoint_dir,\n solver_config.name,\n solver_config.mpicomm,\n 0,\n )\n @test !isfile(cfull)\n end\nend\n\nmain()\n" }, { "path": "test/Driver/gcm_driver_test.jl", "content": "using StaticArrays\nusing Test\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.ConfigTypes\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Grids\nusing ClimateMachine.ODESolvers\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n # callable to set initial conditions\n FT = eltype(state)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n param_set = parameter_set(bl)\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n return nothing\nend\n\nfunction main()\n FT = Float64\n\n # DG polynomial orders\n N = (4, 4)\n\n # Domain resolution\n nelem_horz = 4\n nelem_vert = 6\n resolution = (nelem_horz, nelem_vert)\n\n t0 = FT(0)\n timeend = FT(3600)\n # Timestep size (s)\n dt = FT(1800)\n\n setup = AcousticWaveSetup{FT}()\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n ref_state = HydrostaticState(T_profile)\n turbulence = ConstantDynamicViscosity(FT(0))\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = turbulence,\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n\n ode_solver_type = ClimateMachine.MultirateSolverType(\n splitting_type = ClimateMachine.HEVISplitting(),\n fast_model = AtmosAcousticGravityLinearModel,\n implicit_solver_adjustable = true,\n slow_method = LSRK54CarpenterKennedy,\n fast_method = ARK2ImplicitExplicitMidpoint,\n timestep_ratio = 300,\n )\n driver_config = ClimateMachine.AtmosGCMConfiguration(\n \"GCM Driver test\",\n N,\n resolution,\n setup.domain_height,\n param_set,\n setup;\n model = model,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n ode_dt = dt,\n )\n\n cb_test = 0\n result = ClimateMachine.invoke!(\n solver_config;\n user_info_callback = () -> cb_test += 1,\n )\n @test cb_test > 0\nend\n\nmain()\n" }, { "path": "test/Driver/les_driver_test.jl", "content": "using StaticArrays\nusing Test\n\nusing ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.Mesh.Grids\nusing Thermodynamics\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nfunction init_test!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n z = FT(z)\n _grav::FT = grav(param_set)\n _MSLP::FT = MSLP(param_set)\n\n # These constants are those used by Stevens et al. (2005)\n qref = FT(9.0e-3)\n q_pt_sfc = PhasePartition(qref)\n Rm_sfc = FT(gas_constant_air(param_set, q_pt_sfc))\n T_sfc = FT(290.4)\n P_sfc = _MSLP\n\n # Specify moisture profiles\n q_liq = FT(0)\n q_ice = FT(0)\n\n θ_liq = FT(289.0)\n q_tot = qref\n\n ugeo = FT(7)\n vgeo = FT(-5.5)\n u, v, w = ugeo, vgeo, FT(0)\n\n # Pressure\n H = Rm_sfc * T_sfc / _grav\n p = P_sfc * exp(-z / H)\n\n # Density, Temperature\n ts = PhaseEquil_pθq(param_set, p, θ_liq, q_tot)\n ρ = air_density(ts)\n\n e_kin = FT(1 / 2) * FT((u^2 + v^2 + w^2))\n e_pot = _grav * z\n E = ρ * total_energy(e_kin, e_pot, ts)\n\n state.ρ = ρ\n state.ρu = SVector(ρ * u, ρ * v, ρ * w)\n state.energy.ρe = E\n state.moisture.ρq_tot = ρ * q_tot\n\n return nothing\nend\n\nfunction main()\n @test_throws ArgumentError ClimateMachine.init(dsisable_gpu = true)\n ClimateMachine.init()\n\n FT = Float64\n\n # DG polynomial orders\n N = (4, 4)\n\n # Domain resolution and size\n Δh = FT(40)\n Δv = FT(40)\n resolution = (Δh, Δh, Δv)\n\n xmax = FT(320)\n ymax = FT(320)\n zmax = FT(400)\n\n t0 = FT(0)\n timeend = FT(10)\n CFL = FT(0.4)\n\n driver_config = ClimateMachine.AtmosLESConfiguration(\n \"Driver test\",\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n param_set,\n init_test!,\n )\n ode_solver_type = ClimateMachine.ExplicitSolverType()\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n Courant_number = CFL,\n )\n\n # Test the courant wrapper\n # by default the CFL should be less than what asked for\n CFL_nondiff = ClimateMachine.DGMethods.courant(\n ClimateMachine.Courant.nondiffusive_courant,\n solver_config,\n )\n @test CFL_nondiff < CFL\n CFL_adv = ClimateMachine.DGMethods.courant(\n ClimateMachine.Courant.advective_courant,\n solver_config,\n )\n CFL_adv_v = ClimateMachine.DGMethods.courant(\n ClimateMachine.Courant.advective_courant,\n solver_config;\n direction = VerticalDirection(),\n )\n CFL_adv_h = ClimateMachine.DGMethods.courant(\n ClimateMachine.Courant.advective_courant,\n solver_config;\n direction = HorizontalDirection(),\n )\n\n # compute known advective Courant number (based on initial conditions)\n ugeo_abs = FT(7)\n vgeo_abs = FT(5.5)\n Δt = solver_config.dt\n ca_h = ugeo_abs * (Δt / Δh) + vgeo_abs * (Δt / Δh)\n # vertical velocity is 0\n caᵥ = FT(0.0)\n @test isapprox(CFL_adv_v, caᵥ, atol = 10 * eps(FT))\n @test isapprox(CFL_adv_h, ca_h, atol = 0.0005)\n @test isapprox(CFL_adv, ca_h, atol = 0.0005)\n\n cb_test = 0\n result = ClimateMachine.invoke!(solver_config)\n # cb_test should be zero since user_info_callback not specified\n @test cb_test == 0\n\n result = ClimateMachine.invoke!(\n solver_config,\n user_info_callback = () -> cb_test += 1,\n )\n # cb_test should be greater than one if the user_info_callback got called\n @test cb_test > 0\n\n # Test that if dt is not adjusted based on final time the CFL is correct\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n Courant_number = CFL,\n timeend_dt_adjust = false,\n )\n\n CFL_nondiff = ClimateMachine.DGMethods.courant(\n ClimateMachine.Courant.nondiffusive_courant,\n solver_config,\n )\n @test CFL_nondiff ≈ CFL\n\n # Test that if dt is not adjusted based on final time the CFL is correct\n for fixed_number_of_steps in (-1, 0, 10)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n Courant_number = CFL,\n fixed_number_of_steps = fixed_number_of_steps,\n )\n\n if fixed_number_of_steps < 0\n @test solver_config.timeend == timeend\n else\n @test solver_config.timeend ==\n solver_config.dt * fixed_number_of_steps\n @test solver_config.numberofsteps == fixed_number_of_steps\n end\n end\nend\n\nmain()\n" }, { "path": "test/Driver/mms3.jl", "content": "using ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.Orientations: NoOrientation\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\n\nimport Thermodynamics: total_specific_enthalpy\nimport ClimateMachine.BalanceLaws: source, prognostic_vars\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport CLIMAParameters\n# Assume zero reference temperature\nCLIMAParameters.Planet.T_0(::EarthParameterSet) = 0\n\nusing LinearAlgebra\nusing MPI\nusing StaticArrays\nusing Test\nusing UnPack\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"test\",\n \"Numerics\",\n \"DGMethods\",\n \"compressible_Navier_Stokes\",\n \"mms_solution_generated.jl\",\n))\n\ntotal_specific_enthalpy(ts::PhaseDry{FT}, e_tot::FT) where {FT <: Real} =\n zero(FT)\n\nfunction mms3_init_state!(problem, bl, state::Vars, aux::Vars, localgeo, t)\n (x1, x2, x3) = localgeo.coord\n state.ρ = ρ_g(t, x1, x2, x3, Val(3))\n state.ρu = SVector(\n U_g(t, x1, x2, x3, Val(3)),\n V_g(t, x1, x2, x3, Val(3)),\n W_g(t, x1, x2, x3, Val(3)),\n )\n state.energy.ρe = E_g(t, x1, x2, x3, Val(3))\nend\n\nstruct MMSSource{N} <: TendencyDef{Source} end\n\nprognostic_vars(::MMSSource{N}) where {N} = (Mass(), Momentum(), Energy())\n\nfunction source(::Mass, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return Sρ_g(t, x1, x2, x3, Val(N))\nend\nfunction source(::Momentum, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return SVector(\n SU_g(t, x1, x2, x3, Val(N)),\n SV_g(t, x1, x2, x3, Val(N)),\n SW_g(t, x1, x2, x3, Val(N)),\n )\nend\nfunction source(::Energy, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return SE_g(t, x1, x2, x3, Val(N))\nend\n\nfunction main()\n FT = Float64\n\n # DG polynomial order\n N = 4\n\n t0 = FT(0)\n timeend = FT(1)\n\n ode_dt = 0.00125\n nsteps = ceil(Int64, timeend / ode_dt)\n ode_dt = timeend / nsteps\n\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK54CarpenterKennedy,\n )\n\n expected_result = FT(3.403104838700577e-02)\n\n problem = AtmosProblem(\n boundaryconditions = (InitStateBC(),),\n init_state_prognostic = mms3_init_state!,\n )\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(μ_exact), WithDivergence()),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (MMSSource{3}(),),\n )\n\n brickrange = (\n range(FT(0); length = 5, stop = 1),\n range(FT(0); length = 5, stop = 1),\n range(FT(0); length = 5, stop = 1),\n )\n topl = BrickTopology(\n MPI.COMM_WORLD,\n brickrange,\n periodicity = (false, false, false),\n connectivity = :face,\n )\n warpfun =\n (x1, x2, x3) -> begin\n (\n x1 + (x1 - 1 / 2) * cos(2 * π * x2 * x3) / 4,\n x2 + exp(sin(2π * (x1 * x2 + x3))) / 20,\n x3 + x1 / 4 + x2^2 / 2 + sin(x1 * x2 * x3),\n )\n end\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = N,\n meshwarp = warpfun,\n )\n driver_config = ClimateMachine.DriverConfiguration(\n AtmosLESConfigType(),\n \"MMS3\",\n (N, N),\n FT,\n ClimateMachine.array_type(),\n param_set,\n model,\n MPI.COMM_WORLD,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n nothing,\n nothing, # filter\n ClimateMachine.AtmosLESSpecificInfo(),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n ode_dt = ode_dt,\n )\n Q₀ = solver_config.Q\n\n # turn on checkpointing\n ClimateMachine.Settings.checkpoint = \"300steps\"\n ClimateMachine.Settings.checkpoint_keep_one = false\n\n # run the simulation\n ClimateMachine.invoke!(solver_config)\n\n # turn off checkpointing and set up a restart\n ClimateMachine.Settings.checkpoint = \"never\"\n ClimateMachine.Settings.restart_from_num = 2\n\n # the solver configuration is where the restart is set up\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n ode_dt = ode_dt,\n )\n\n # run the restarted simulation\n ClimateMachine.invoke!(solver_config)\n\n # test correctness\n dg = DGModel(driver_config)\n Qe = init_ode_state(dg, timeend)\n result = euclidean_distance(Q₀, Qe)\n @test result ≈ expected_result\nend\n\nmain()\n" }, { "path": "test/Driver/runtests.jl", "content": "using MPI, Test\n\ninclude(\"../testhelpers.jl\")\n\n@testset \"Driver\" begin\n runmpi(joinpath(@__DIR__, \"cr_unit_tests.jl\"), ntasks = 1)\n runmpi(joinpath(@__DIR__, \"les_driver_test.jl\"), ntasks = 1)\n runmpi(joinpath(@__DIR__, \"mms3.jl\"), ntasks = 2)\nend\n" }, { "path": "test/InputOutput/VTK/runtests.jl", "content": "module TestVTK\n\nusing Test\nusing GaussQuadrature: legendre, both\nusing ClimateMachine.VTK: writemesh_highorder, writemesh_raw\n\n@testset \"VTK\" begin\n for writemesh in (writemesh_highorder, writemesh_raw)\n Ns = writemesh == writemesh_raw ? ((4, 4, 4), (3, 5, 7)) : ((4, 4, 4),)\n for N in Ns\n for dim in 1:3\n nelem = 3\n T = Float64\n\n Nq = N .+ 1\n r = legendre(T, Nq[1], both)[1]\n s = dim < 2 ? [0] : legendre(T, Nq[2], both)[1]\n t = dim < 3 ? [0] : legendre(T, Nq[3], both)[1]\n Nq1 = length(r)\n Nq2 = length(s)\n Nq3 = length(t)\n Np = Nq1 * Nq2 * Nq3\n\n x1 = Array{T, 4}(undef, Nq1, Nq2, Nq3, nelem)\n x2 = Array{T, 4}(undef, Nq1, Nq2, Nq3, nelem)\n x3 = Array{T, 4}(undef, Nq1, Nq2, Nq3, nelem)\n\n for e in 1:nelem, k in 1:Nq3, j in 1:Nq2, i in 1:Nq1\n xoffset = nelem + 1 - 2e\n x1[i, j, k, e], x2[i, j, k, e], x3[i, j, k, e] =\n r[i] - xoffset, s[j], t[k]\n end\n\n if dim == 1\n x1 = x1 .^ 3\n elseif dim == 2\n x1, x2 = x1 + sin.(π * x2) / 5, x2 + exp.(-x1 .^ 2)\n else\n x1, x2, x3 = x1 + sin.(π * x2) / 5,\n x2 + exp.(-hypot.(x1, x3) .^ 2),\n x3 + sin.(π * x1) / 5\n end\n d = exp.(sin.(hypot.(x1, x2, x3)))\n s = copy(d)\n\n if dim == 1\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1;\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1;\n x2 = x2,\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1;\n x2 = x2,\n x3 = x3,\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n elseif dim == 2\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1,\n x2;\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1,\n x2;\n x3 = x3,\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n elseif dim == 3\n @test \"test$(dim)d.vtu\" == writemesh(\n \"test$(dim)d\",\n x1,\n x2,\n x3;\n fields = ((\"d\", d), (\"s\", s)),\n )[1]\n end\n end\n end\n end\nend\n\n\nusing MPI\nMPI.Initialized() || MPI.Init()\nusing ClimateMachine.Mesh.Topologies: BrickTopology\nusing ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid\nusing ClimateMachine.VTK: writevtk_helper\n\nlet\n mpicomm = MPI.COMM_SELF\n for FT in (Float64,) #Float32)\n for dim in 2:3\n for _N in ((2, 3, 4), (0, 2, 5), (3, 0, 0), (0, 0, 0))\n N = _N[1:dim]\n if dim == 2\n Ne = (4, 5)\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false),\n connectivity = :face,\n )\n warpfun =\n (x1, x2, _) -> begin\n (x1 + sin(x1 * x2), x2 + sin(2 * x1 * x2), 0)\n end\n elseif dim == 3\n Ne = (3, 4, 5)\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n range(FT(0); length = Ne[3] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false, false),\n connectivity = :face,\n )\n warpfun =\n (x1, x2, x3) -> begin\n (\n x1 + (x1 - 1 / 2) * cos(2 * π * x2 * x3) / 4,\n x2 + exp(sin(2π * (x1 * x2 + x3))) / 20,\n x3 + x1 / 4 + x2^2 / 2 + sin(x1 * x2 * x3),\n )\n end\n end\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = Array,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n Q = rand(FT, prod(N .+ 1), 3, prod(Ne))\n prefix = \"test$(dim)d_raw$(prod(ntuple(i->\"_$(N[i])\", dim)))\"\n @test \"$(prefix).vtu\" == writevtk_helper(\n prefix,\n grid.vgeo,\n Q,\n grid,\n (\"a\", \"b\", \"c\");\n number_sample_points = 0,\n )[1]\n prefix = \"test$(dim)d_high_order$(prod(ntuple(i->\"_$(N[i])\", dim)))\"\n @test \"$(prefix).vtu\" == writevtk_helper(\n prefix,\n grid.vgeo,\n Q,\n grid,\n (\"a\", \"b\", \"c\");\n number_sample_points = 10,\n )[1]\n end\n end\n end\nend\n\nend #module TestVTK\n" }, { "path": "test/InputOutput/Writers/runtests.jl", "content": "module TestWriters\n\nusing Dates\nusing NCDatasets\nusing OrderedCollections\nusing Test\nusing ClimateMachine.Writers\n\n@testset \"Writers\" begin\n odims = OrderedDict(\n \"x\" => (collect(1:5), Dict()),\n \"y\" => (collect(1:5), Dict()),\n \"z\" => (collect(1010:10:1050), Dict()),\n )\n ovartypes = OrderedDict(\n \"v1\" => ((\"x\", \"y\", \"z\"), Float64, Dict()),\n \"v2\" => ((\"x\", \"y\", \"z\"), Float64, Dict()),\n )\n vals1 = rand(5, 5, 5)\n vals2 = rand(5, 5, 5)\n\n nc = NetCDFWriter()\n nfn, _ = mktemp()\n nfull = full_name(nc, nfn)\n\n touch(nfull)\n @test_throws ErrorException init_data(nc, nfn, true, odims, ovartypes)\n rm(nfull)\n\n init_data(nc, nfn, false, odims, ovartypes)\n append_data(nc, OrderedDict(\"v1\" => vals1, \"v2\" => vals2), 2.0)\n\n NCDataset(nfull, \"r\") do nds\n xdim = nds[\"x\"][:]\n ydim = nds[\"y\"][:]\n zdim = nds[\"z\"][:]\n @test xdim == odims[\"x\"][1]\n @test ydim == odims[\"y\"][1]\n @test zdim == odims[\"z\"][1]\n @test try\n adim = nds[\"a\"][:]\n adim == ones(5)\n false\n catch e\n true\n end\n t = nds[\"time\"][:]\n v1 = nds[\"v1\"][:]\n v2 = nds[\"v2\"][:]\n @test length(t) == 1\n @test t[1] == DateTime(1900, 1, 1, 0, 0, 2)\n @test v1[:, :, :, 1] == vals1\n @test v2[:, :, :, 1] == vals2\n end\nend\n\nend #module TestWriters\n" }, { "path": "test/InputOutput/runtests.jl", "content": "using Test, Pkg\n\n@testset \"InputOutput\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"VTK\", \"Writers\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"InputOutput/$submodule\" in ARGS ||\n \"InputOutput\" in ARGS\n include_test(submodule)\n end\n end\n\nend\n" }, { "path": "test/Land/Model/Artifacts.toml", "content": "[richards]\ngit-tree-sha1 = \"ff73fa6a0b6a807e71a6921f7ef7d0befe776edd\"\n\n[richards_sand]\ngit-tree-sha1 = \"b0dc82dd02159c646e909bfb61170d3b9dc347f3\"\n\n[tiltedv]\ngit-tree-sha1 = \"db27235cb7ce2b7674607876da15d1635906b512\"" }, { "path": "test/Land/Model/freeze_thaw_alone.jl", "content": "# Test that freeze thaw alone conserves water mass\n\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\nusing CLIMAParameters.Planet: ρ_cloud_liq\nusing CLIMAParameters.Planet: ρ_cloud_ice\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n@testset \"Freeze thaw alone\" begin\n struct tmp_model <: BalanceLaw end\n struct tmp_param_set <: AbstractParameterSet end\n\n function get_grid_spacing(\n N_poly::Int64,\n nelem_vert::Int64,\n zmax::FT,\n zmin::FT,\n ) where {FT}\n test_config = ClimateMachine.SingleStackConfiguration(\n \"TmpModel\",\n N_poly,\n nelem_vert,\n zmax,\n tmp_param_set(),\n tmp_model();\n zmin = zmin,\n )\n\n Δ = min_node_distance(test_config.grid)\n return Δ\n end\n\n function init_soil_water!(land, state, aux, coordinates, time)\n ϑ_l = eltype(state)(land.soil.water.initialϑ_l(aux))\n θ_i = eltype(state)(land.soil.water.initialθ_i(aux))\n state.soil.water.ϑ_l = ϑ_l\n state.soil.water.θ_i = θ_i\n param_set = land.param_set\n\n θ_l =\n volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_ds = land.soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n\n state.soil.heat.ρe_int = volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n )\n end\n\n FT = Float64\n ClimateMachine.init()\n\n N_poly = 1\n nelem_vert = 30\n zmax = FT(0)\n zmin = FT(-1)\n t0 = FT(0)\n timeend = FT(60 * 60 * 24)\n dt = FT(1800)\n\n Δ = get_grid_spacing(N_poly, nelem_vert, zmax, zmin)\n freeze_thaw_source = PhaseChange{FT}(Δz = Δ)\n ρp = FT(2700) # kg/m^3\n ρc_ds = FT(2e06) # J/m^3/K\n\n Ksat = FT(0.0)\n S_s = FT(1e-4)\n wpf = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n soil_param_functions = SoilParamFunctions(\n FT;\n porosity = 0.75,\n ν_ss_gravel = 0.0,\n ν_ss_om = 0.0,\n ν_ss_quartz = 0.5,\n ρc_ds = ρc_ds,\n ρp = ρp,\n κ_solid = 1.0,\n κ_sat_unfrozen = 1.0,\n κ_sat_frozen = 1.0,\n water = wpf,\n )\n\n\n bottom_flux = (aux, t) -> eltype(aux)(0.0)\n surface_flux = (aux, t) -> eltype(aux)(0.0)\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_water = Neumann(bottom_flux),\n soil_heat = Dirichlet((aux, t) -> eltype(aux)(280)),\n ),\n surface_bc = LandComponentBC(\n soil_water = Neumann(surface_flux),\n soil_heat = Dirichlet((aux, t) -> eltype(aux)(290)),\n ),\n )\n ϑ_l0 = (aux) -> eltype(aux)(1e-10)\n θ_i0 = (aux) -> eltype(aux)(0.33)\n\n soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0, initialθ_i = θ_i0)\n\n T_init = (aux) -> eltype(aux)(aux.z * 10 + 290)\n soil_heat_model = SoilHeatModel(FT; initialT = T_init)\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = (freeze_thaw_source,)\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n state_types = (Prognostic(), Auxiliary())\n initial =\n Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\n ClimateMachine.invoke!(solver_config)\n final =\n Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\n m_init =\n ρ_cloud_liq(param_set) * sum(initial[1][\"soil.water.ϑ_l\"]) .+\n ρ_cloud_ice(param_set) * sum(initial[1][\"soil.water.θ_i\"])\n m_final =\n ρ_cloud_liq(param_set) * sum(final[1][\"soil.water.ϑ_l\"]) .+\n ρ_cloud_ice(param_set) * sum(final[1][\"soil.water.θ_i\"])\n\n @test abs(m_final - m_init) < 1e-10\nend\n" }, { "path": "test/Land/Model/haverkamp_test.jl", "content": "# Test that Richard's equation agrees with solution from Bonan's book,\n# simulation 8.2\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Dierckx\nusing Test\nusing Pkg.Artifacts\nusing DelimitedFiles\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\nusing ArtifactWrappers\n\nhaverkamp_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"richards\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/dfijf07io7h5dk1k87saaewgsg9apq8d.csv\",\n filename = \"bonan_haverkamp_data.csv\",\n ),],\n)\nhaverkamp_dataset_path = get_data_folder(haverkamp_dataset)\n\nbonan_sand_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"richards_sand\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/2vk7bvyjah8xd5b7wxcqy72yfd2myjss.csv\",\n filename = \"sand_bonan_sp801.csv\",\n ),],\n)\nbonan_sand_dataset_path = get_data_folder(bonan_sand_dataset)\n\n@testset \"Richard's equation - Haverkamp test\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(aux)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n soil_heat_model = PrescribedTemperatureModel()\n Ksat = FT(0.0443 / (3600 * 100))\n S_s = FT(1e-3)\n water_param_functions = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n\n soil_param_functions =\n SoilParamFunctions(FT; porosity = 0.495, water = water_param_functions)\n surface_state = (aux, t) -> eltype(aux)(0.494)\n bottom_flux = (aux, t) -> aux.soil.water.K * eltype(aux)(-1)\n ϑ_l0 = (aux) -> eltype(aux)(0.24)\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),\n )\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = Haverkamp(FT;),\n initialϑ_l = ϑ_l0,\n )\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60 * 24)\n\n dt = FT(6)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n\n ClimateMachine.invoke!(solver_config)\n ϑ_l_ind = varsindex(vars_state(m, Prognostic(), FT), :soil, :water, :ϑ_l)\n ϑ_l = Array(Q[:, ϑ_l_ind, :][:])\n z_ind = varsindex(vars_state(m, Auxiliary(), FT), :z)\n z = Array(aux[:, z_ind, :][:])\n\n # Compare with Bonan simulation data at 1 day.\n data = joinpath(haverkamp_dataset_path, \"bonan_haverkamp_data.csv\")\n ds_bonan = readdlm(data, ',')\n bonan_moisture = reverse(ds_bonan[:, 1])\n bonan_z = reverse(ds_bonan[:, 2]) ./ 100.0\n\n\n # Create an interpolation from the Bonan data\n bonan_moisture_continuous = Spline1D(bonan_z, bonan_moisture)\n bonan_at_clima_z = bonan_moisture_continuous.(z)\n MSE = mean((bonan_at_clima_z .- ϑ_l) .^ 2.0)\n @test MSE < 1e-5\nend\n\n\n@testset \"Richard's equation - Haverkamp implicit test\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(aux)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n soil_heat_model = PrescribedTemperatureModel()\n Ksat = FT(0.0443 / (3600 * 100))\n S_s = FT(1e-3)\n water_param_functions = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n soil_param_functions =\n SoilParamFunctions(FT; porosity = 0.495, water = water_param_functions)\n surface_value = FT(0.494)\n bottom_flux_multiplier = FT(1.0)\n initial_moisture = FT(0.24)\n\n surface_state = (aux, t) -> surface_value\n bottom_flux = (aux, t) -> aux.soil.water.K * bottom_flux_multiplier\n ϑ_l0 = (aux) -> initial_moisture\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),\n )\n\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = Haverkamp(FT;),\n initialϑ_l = ϑ_l0,\n )\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60 * 24)\n\n dt = FT(200)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n\n #################### Change the ode_solver\n\n dg = solver_config.dg\n Q = solver_config.Q\n\n vdg = DGModel(\n driver_config;\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n\n\n linearsolver = BatchedGeneralizedMinimalResidual(\n dg,\n Q;\n max_subspace_size = 30,\n atol = -1.0,\n rtol = 1e-9,\n )\n\n \"\"\"\n N(q)(Q) = Qhat => F(Q) = N(q)(Q) - Qhat\n\n F(Q) == 0\n ||F(Q^i) || / ||F(Q^0) || < tol\n\n \"\"\"\n nonlinearsolver =\n JacobianFreeNewtonKrylovSolver(Q, linearsolver; tol = 1e-9)\n\n ode_solver = ARK548L2SA2KennedyCarpenter(\n dg,\n vdg,\n NonLinearBackwardEulerSolver(\n nonlinearsolver;\n isadjustable = true,\n preconditioner_update_freq = 100,\n ),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n variant = NaiveVariant(),\n )\n\n solver_config.solver = ode_solver\n\n #######################################\n\n mygrid = solver_config.dg.grid\n aux = solver_config.dg.state_auxiliary\n ClimateMachine.invoke!(solver_config)\n ϑ_l_ind = varsindex(vars_state(m, Prognostic(), FT), :soil, :water, :ϑ_l)\n ϑ_l = Array(Q[:, ϑ_l_ind, :][:])\n z_ind = varsindex(vars_state(m, Auxiliary(), FT), :z)\n z = Array(aux[:, z_ind, :][:])\n\n # Compare with Bonan simulation data at 1 day.\n data = joinpath(haverkamp_dataset_path, \"bonan_haverkamp_data.csv\")\n ds_bonan = readdlm(data, ',')\n bonan_moisture = reverse(ds_bonan[:, 1])\n bonan_z = reverse(ds_bonan[:, 2]) ./ 100.0\n\n\n # Create an interpolation from the Bonan data\n bonan_moisture_continuous = Spline1D(bonan_z, bonan_moisture)\n bonan_at_clima_z = bonan_moisture_continuous.(z)\n MSE = mean((bonan_at_clima_z .- ϑ_l) .^ 2.0)\n @test MSE < 1e-5\nend\n\n\n\n@testset \"Richard's equation - Sand van Genuchten test\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(aux)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n soil_heat_model = PrescribedTemperatureModel()\n wpf = WaterParamFunctions(\n FT;\n Ksat = 34 / (3600 * 100),\n θ_r = 0.075,\n S_s = 1e-3,\n )\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.287, water = wpf)\n\n surface_state = (aux, t) -> eltype(aux)(0.267)\n bottom_flux = (aux, t) -> aux.soil.water.K * eltype(aux)(-1)\n ϑ_l0 = (aux) -> eltype(aux)(0.1)\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),\n )\n\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = vanGenuchten(FT; n = 3.96, α = 2.7),\n initialϑ_l = ϑ_l0,\n )\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 1\n nelem_vert = 150\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1.5)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60 * 0.8)\n\n dt = FT(0.5)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n\n ClimateMachine.invoke!(solver_config)\n ϑ_l_ind = varsindex(vars_state(m, Prognostic(), FT), :soil, :water, :ϑ_l)\n ϑ_l = Array(Q[:, ϑ_l_ind, :][:])\n z_ind = varsindex(vars_state(m, Auxiliary(), FT), :z)\n z = Array(aux[:, z_ind, :][:])\n\n # Compare with Bonan simulation data at 0.8 h\n data = joinpath(bonan_sand_dataset_path, \"sand_bonan_sp801.csv\")\n ds_bonan = readdlm(data, ',')\n bonan_moisture = reverse(ds_bonan[:, 1])\n bonan_z = reverse(ds_bonan[:, 2]) ./ 100.0\n\n\n # Create an interpolation from the Bonan data\n bonan_moisture_continuous = Spline1D(bonan_z, bonan_moisture)\n bonan_at_clima_z = bonan_moisture_continuous.(z)\n MSE = mean((bonan_at_clima_z .- ϑ_l) .^ 2.0)\n @test MSE < 1e-5\nend\n" }, { "path": "test/Land/Model/heat_analytic_unit_test.jl", "content": "# Test heat equation agrees with analytic solution to problem 55 on page 28 in https://ocw.mit.edu/courses/mathematics/18-303-linear-partial-differential-equations-fall-2006/lecture-notes/heateqni.pdf\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Dierckx\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws\n\nfunction init_soil!(land, state, aux, localgeo, time)\n FT = eltype(state)\n ϑ_l, θ_i = get_water_content(land.soil.water, aux, state, time)\n θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n param_set = parameter_set(land)\n ρc_s = volumetric_heat_capacity(\n θ_l,\n θ_i,\n land.soil.param_functions.ρc_ds,\n param_set,\n )\n\n state.soil.heat.ρe_int = FT(volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n ))\nend\n\nconst FT = Float64\n\n@testset \"Heat analytic unit test\" begin\n ClimateMachine.init()\n\n soil_param_functions = SoilParamFunctions(\n FT;\n porosity = 0.495,\n ν_ss_gravel = 0.1,\n ν_ss_om = 0.1,\n ν_ss_quartz = 0.1,\n ρc_ds = 0.43314518988433487,\n κ_solid = 8.0,\n ρp = 2700.0,\n κ_sat_unfrozen = 0.57,\n κ_sat_frozen = 2.29,\n )\n\n heat_surface_state = (aux, t) -> eltype(aux)(0.0)\n\n tau = FT(1) # period (sec)\n A = FT(5) # amplitude (K)\n ω = FT(2 * pi / tau)\n heat_bottom_state = (aux, t) -> A * cos(ω * t)\n T_init = (aux) -> eltype(aux)(0.0)\n\n soil_water_model = PrescribedWaterModel(\n (aux, t) -> eltype(aux)(0.0),\n (aux, t) -> eltype(aux)(0.0),\n )\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_heat = Dirichlet(heat_bottom_state)),\n surface_bc = LandComponentBC(soil_heat = Dirichlet(heat_surface_state)),\n )\n soil_heat_model = SoilHeatModel(FT; initialT = T_init)\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n )\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(1)\n zmin = FT(0)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(2)\n dt = FT(1e-4)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n aux = solver_config.dg.state_auxiliary\n ClimateMachine.invoke!(solver_config)\n t = ODESolvers.gettime(solver_config.solver)\n\n z_ind = varsindex(vars_state(m, Auxiliary(), FT), :z)\n z = Array(aux[:, z_ind, :][:])\n\n T_ind = varsindex(vars_state(m, Auxiliary(), FT), :soil, :heat, :T)\n T = Array(aux[:, T_ind, :][:])\n\n num =\n exp.(sqrt(ω / 2) * (1 + im) * (1 .- z)) .-\n exp.(-sqrt(ω / 2) * (1 + im) * (1 .- z))\n denom = exp(sqrt(ω / 2) * (1 + im)) - exp.(-sqrt(ω / 2) * (1 + im))\n analytic_soln = real(num .* A * exp(im * ω * timeend) / denom)\n MSE = mean((analytic_soln .- T) .^ 2.0)\n @test eltype(aux) == FT\n @test MSE < 1e-5\nend\n" }, { "path": "test/Land/Model/prescribed_twice.jl", "content": "# Test that the land model still runs, even with the lowest/simplest\n# version of soil (prescribed heat and prescribed water - no state\n# variables)\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n@testset \"Prescribed Models\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time) end\n\n soil_water_model = PrescribedWaterModel()\n soil_heat_model = PrescribedTemperatureModel()\n soil_param_functions = nothing\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n\n t0 = FT(0)\n timeend = FT(60)\n dt = FT(1)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n\n ClimateMachine.invoke!(solver_config;)\n\n t = ODESolvers.gettime(solver_config.solver)\n state_vars = SingleStackUtils.get_vars_from_nodal_stack(\n mygrid,\n Q,\n vars_state(m, Prognostic(), FT),\n )\n aux_vars = SingleStackUtils.get_vars_from_nodal_stack(\n mygrid,\n aux,\n vars_state(m, Auxiliary(), FT),\n )\n #Make sure it runs, and that there are no state variables, and only \"x,y,z\" as aux.\n @test t == timeend\n @test size(Base.collect(keys(aux_vars)))[1] == 3\n @test size(Base.collect(keys(state_vars)))[1] == 0\nend\n" }, { "path": "test/Land/Model/runtests.jl", "content": "using Test, Pkg\n@testset \"Land\" begin\n include(\"test_heat_parameterizations.jl\")\n include(\"test_water_parameterizations.jl\")\n include(\"prescribed_twice.jl\")\n include(\"freeze_thaw_alone.jl\")\n include(\"test_overland_flow_analytic.jl\")\n include(\"test_physical_bc.jl\")\n include(\"test_radiative_energy_flux_functions.jl\")\nend\n" }, { "path": "test/Land/Model/soil_heterogeneity.jl", "content": "using MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Dierckx\nusing Test\nusing Pkg.Artifacts\nusing DelimitedFiles\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n\n@testset \"hydrostatic test 1\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(aux)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n soil_heat_model = PrescribedTemperatureModel()\n\n Ksat = (aux) -> eltype(aux)(0.0443 / (3600 * 100))\n S_s = (aux) -> eltype(aux)((1e-3) * exp(-0.2 * aux.z))\n vgn = FT(2)\n wpf = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.495, water = wpf)\n bottom_flux = (aux, t) -> eltype(aux)(0.0)\n surface_flux = bottom_flux\n ϑ_l0 = (aux) -> eltype(aux)(0.494)\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Neumann(surface_flux)),\n )\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = vanGenuchten(FT; n = vgn),\n initialϑ_l = ϑ_l0,\n )\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 2\n nelem_vert = 20\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-10)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60 * 24 * 200)\n\n dt = FT(500)\n n_outputs = 3\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n aux = solver_config.dg.state_auxiliary\n state_types = (Prognostic(), Auxiliary())\n dons_arr =\n Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\n time_data = FT[0]\n callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n z = dons_arr[1][\"z\"]\n interface_z = -1.0395\n function hydrostatic_profile(z, zm, porosity, n, α)\n myf = eltype(z)\n m = FT(1 - 1 / n)\n S = FT((FT(1) + (α * (z - zm))^n)^(-m))\n return FT(S * porosity)\n end\n function soln(z, interface, porosity, n, α, δ, S_s)\n if z < interface\n return porosity + S_s * (interface - z) * exp(-δ * z)\n else\n return hydrostatic_profile(z, interface, porosity, n, α)\n end\n end\n\n MSE = mean(\n (\n soln.(z, interface_z, 0.495, vgn, 2.6, 0.2, 1e-3) .-\n dons_arr[4][\"soil.water.ϑ_l\"]\n ) .^ 2.0,\n )\n @test MSE < 1e-4\nend\n\n\n@testset \"hydrostatic test 2\" begin\n ClimateMachine.init()\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(aux)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n soil_heat_model = PrescribedTemperatureModel()\n\n Ksat = (4.42 / 3600 / 100)\n S_s = 1e-3\n wpf = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n\n soil_param_functions = SoilParamFunctions(FT, porosity = 0.6, water = wpf)\n bottom_flux = (aux, t) -> eltype(aux)(0.0)\n surface_flux = bottom_flux\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Neumann(surface_flux)),\n )\n sigmoid(x, offset, width) =\n typeof(x)(exp((x - offset) / width) / (1 + exp((x - offset) / width)))\n function soln(z::f, interface::f, porosity::f) where {f}\n function hydrostatic_profile(\n z::f,\n interface::f,\n porosity::f,\n n::f,\n α::f,\n m::f,\n )\n ψ_interface = f(-1)\n ψ = -(z - interface) + ψ_interface\n S = (f(1) + (-α * ψ)^n)^(-m)\n return S * porosity\n end\n if z < interface\n return hydrostatic_profile(\n z,\n interface,\n porosity,\n f(1.31),\n f(1.9),\n f(1) - f(1) / f(1.31),\n )\n else\n return hydrostatic_profile(\n z,\n interface,\n porosity,\n f(1.89),\n f(7.5),\n f(1) - f(1) / f(1.89),\n )\n end\n end\n ϑ_l0 = (aux) -> soln(aux.z, -1.0, 0.6)\n\n vgα(aux) = aux.z < -1.0 ? 1.9 : 7.5\n vgn(aux) = aux.z < -1.0 ? 1.31 : 1.89\n\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = vanGenuchten(FT; n = vgn, α = vgα),\n initialϑ_l = ϑ_l0,\n )\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 1\n nelem_vert = 80\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-2)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60 * 12)\n\n dt = FT(5)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n\n ClimateMachine.invoke!(solver_config)\n\n state_types = (Prognostic(), Auxiliary())\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n z = dons[\"z\"]\n\n RMSE =\n mean(\n (soln.(z, Ref(-1.0), Ref(0.6)) .- dons[\"soil.water.ϑ_l\"]) .^ 2.0,\n )^0.5\n @test RMSE < 2.0 * eps(FT)\nend\n" }, { "path": "test/Land/Model/test_bc.jl", "content": "# Test that the way we specify boundary conditions works as expected\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n@testset \"Boundary condition functions\" begin\n ClimateMachine.init()\n\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n wpf = WaterParamFunctions(FT; Ksat = 1e-7, S_s = 1e-3)\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.75, water = wpf)\n bottom_flux_amplitude = FT(-3.0)\n f = FT(pi * 2.0 / 300.0)\n bottom_flux =\n (aux, t) -> bottom_flux_amplitude * sin(f * t) * aux.soil.water.K\n surface_state = (aux, t) -> eltype(aux)(0.2)\n ϑ_l0 = (aux) -> eltype(aux)(0.2)\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),\n )\n soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0)\n soil_heat_model = PrescribedTemperatureModel()\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n nelem_vert = 50\n\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n\n t0 = FT(0)\n timeend = FT(300)\n dt = FT(0.05)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n grads = solver_config.dg.state_gradient_flux\n K∇h_vert_ind =\n varsindex(vars_state(m, GradientFlux(), FT), :soil, :water)[3]\n K_ind = varsindex(vars_state(m, Auxiliary(), FT), :soil, :water, :K)\n n_outputs = 30\n\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n dons_arr = Dict([k => Dict() for k in 1:n_outputs]...)\n\n iostep = [1]\n callback = GenericCallbacks.EveryXSimulationTime(\n every_x_simulation_time,\n ) do (init = false)\n t = ODESolvers.gettime(solver_config.solver)\n K = aux[:, K_ind, :]\n K∇h_vert = grads[:, K∇h_vert_ind, :]\n all_vars = Dict{String, Array}(\n \"t\" => [t],\n \"K\" => K,\n \"K∇h_vert\" => K∇h_vert,\n )\n dons_arr[iostep[1]] = all_vars\n iostep[1] += 1\n nothing\n end\n\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n t = ODESolvers.gettime(solver_config.solver)\n K = aux[:, K_ind, :]\n K∇h_vert = grads[:, K∇h_vert_ind, :]\n all_vars = Dict{String, Array}(\"t\" => [t], \"K\" => K, \"K∇h_vert\" => K∇h_vert)\n dons_arr[n_outputs] = all_vars\n\n\n computed_bottom_∇h =\n [dons_arr[k][\"K∇h_vert\"][1] for k in 1:n_outputs] ./ [dons_arr[k][\"K\"][1] for k in 1:n_outputs]\n\n\n t = [dons_arr[k][\"t\"][1] for k in 1:n_outputs]\n # we need a -1 out in front here because the flux BC is on -K∇h\n prescribed_bottom_∇h = t -> FT(-1) * FT(-3.0 * sin(pi * 2.0 * t / 300.0))\n\n MSE = mean((prescribed_bottom_∇h.(t) .- computed_bottom_∇h) .^ 2.0)\n @test MSE < 1e-7\nend\n" }, { "path": "test/Land/Model/test_bc_3d.jl", "content": "# Test that the way we specify boundary conditions works as expected\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n@testset \"Boundary condition functions\" begin\n ClimateMachine.init()\n\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n wpf = WaterParamFunctions(FT; Ksat = 1e-7, S_s = 1e-3)\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.75, water = wpf)\n bottom_flux_amplitude = FT(-3.0)\n f = FT(pi * 2.0 / 300.0)\n # nota bene: the flux is -K∇h\n bottom_flux =\n (aux, t) -> bottom_flux_amplitude * sin(f * t) * aux.soil.water.K\n surface_state = (aux, t) -> eltype(aux)(0.2)\n lateral_state = (aux, t) -> eltype(aux)(0.0)\n ϑ_l0 = (aux) -> eltype(aux)(0.2)\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Dirichlet(surface_state)),\n minx_bc = LandComponentBC(soil_water = Neumann(lateral_state)),\n maxx_bc = LandComponentBC(soil_water = Neumann(lateral_state)),\n miny_bc = LandComponentBC(soil_water = Neumann(lateral_state)),\n maxy_bc = LandComponentBC(soil_water = Neumann(lateral_state)),\n )\n soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0)\n soil_heat_model = PrescribedTemperatureModel()\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n xres = FT(0.2)\n yres = FT(0.2)\n zres = FT(0.01)\n # Specify the domain boundaries.\n zmax = FT(0)\n zmin = FT(-1)\n xmax = FT(1)\n ymax = FT(1)\n\n driver_config = ClimateMachine.MultiColumnLandModel(\n \"LandModel\",\n (N_poly, N_poly),\n (xres, yres, zres),\n xmax,\n ymax,\n zmax,\n param_set,\n m;\n zmin = zmin,\n )\n\n\n t0 = FT(0)\n timeend = FT(300)\n dt = FT(0.5)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n n_outputs = 30\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n state_types = (Auxiliary(), GradientFlux())\n dons_arr =\n Dict[dict_of_nodal_states(solver_config, state_types; interp = false)]\n time_data = FT[0] # store time data\n\n # We specify a function which evaluates `every_x_simulation_time` and returns\n # the state vector, appending the variables we are interested in into\n # `all_data`.\n\n callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = false)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n # # Run the integration\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n computed_bottom_∇h =\n [dons_arr[k][\"soil.water.K∇h[3]\"][1] for k in 2:n_outputs] ./ [dons_arr[k][\"soil.water.K\"][1] for k in 2:n_outputs]\n\n\n t = time_data[2:n_outputs]\n # we need a -1 out in front here because the flux BC is on -K∇h\n prescribed_bottom_∇h = t -> FT(-1) * FT(-3.0 * sin(pi * 2.0 * t / 300.0))\n\n MSE = mean((prescribed_bottom_∇h.(t) .- computed_bottom_∇h) .^ 2.0)\n computed_y1_∇h = maximum(abs.(\n [dons_arr[k][\"soil.water.K∇h[2]\"][1] for k in 2:n_outputs] ./ [dons_arr[k][\"soil.water.K\"][1] for k in 2:n_outputs],\n ))\n computed_x1_∇h = maximum(abs.(\n [dons_arr[k][\"soil.water.K∇h[1]\"][1] for k in 2:n_outputs] ./ [dons_arr[k][\"soil.water.K\"][1] for k in 2:n_outputs],\n ))\n @test MSE < 1e-4\n @test computed_x1_∇h < 1e-10\n @test computed_y1_∇h < 1e-10\nend\n" }, { "path": "test/Land/Model/test_heat_parameterizations.jl", "content": "using MPI\nusing OrderedCollections\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: ρ_cloud_liq, ρ_cloud_ice, cp_l, cp_i, T_0, LH_f0\nusing CLIMAParameters.Atmos.Microphysics: K_therm\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Land\n\n@testset \"Land heat parameterizations\" begin\n FT = Float64\n # Density of liquid water (kg/m``^3``)\n _ρ_l = FT(ρ_cloud_liq(param_set))\n # Density of ice water (kg/m``^3``)\n _ρ_i = FT(ρ_cloud_ice(param_set))\n # Volum. isobaric heat capacity liquid water (J/m3/K)\n _ρcp_l = FT(cp_l(param_set) * _ρ_l)\n # Volumetric isobaric heat capacity ice (J/m3/K)\n _ρcp_i = FT(cp_i(param_set) * _ρ_i)\n # Reference temperature (K)\n _T_ref = FT(T_0(param_set))\n # Latent heat of fusion at ``T_0`` (J/kg)\n _LH_f0 = FT(LH_f0(param_set))\n # Thermal conductivity of dry air\n κ_air = FT(K_therm(param_set))\n\n @test temperature_from_ρe_int(5.4e7, 0.05, 2.1415e6, param_set) ==\n FT(_T_ref + (5.4e7 + 0.05 * _ρ_i * _LH_f0) / 2.1415e6)\n\n @test volumetric_heat_capacity(0.25, 0.05, 1e6, param_set) ==\n FT(1e6 + 0.25 * _ρcp_l + 0.05 * _ρcp_i)\n\n @test volumetric_internal_energy(0.05, 2.1415e6, 300.0, param_set) ==\n FT(2.1415e6 * (300.0 - _T_ref) - 0.05 * _ρ_i * _LH_f0)\n\n @test saturated_thermal_conductivity(0.25, 0.05, 0.57, 2.29) ==\n FT(0.57^(0.25 / (0.05 + 0.25)) * 2.29^(0.05 / (0.05 + 0.25)))\n\n @test saturated_thermal_conductivity(0.0, 0.0, 0.57, 2.29) == FT(0.0)\n\n @test relative_saturation(0.25, 0.05, 0.4) == FT((0.25 + 0.05) / 0.4)\n\n # Test branching in kersten_number\n soil_param_functions = SoilParamFunctions(\n FT;\n porosity = 0.2,\n ν_ss_gravel = 0.1,\n ν_ss_om = 0.1,\n ν_ss_quartz = 0.1,\n κ_solid = 0.1,\n ρp = 1.0,\n )\n # ice fraction = 0\n @test kersten_number(0.0, 0.75, soil_param_functions) == FT(\n 0.75^((FT(1) + 0.1 - 0.24 * 0.1 - 0.1) / FT(2)) *\n (\n (FT(1) + exp(-18.1 * 0.75))^(-FT(3)) -\n ((FT(1) - 0.75) / FT(2))^FT(3)\n )^(FT(1) - 0.1),\n )\n\n # ice fraction ~= 0\n @test kersten_number(0.05, 0.75, soil_param_functions) ==\n FT(0.75^(FT(1) + 0.1))\n\n @test thermal_conductivity(1.5, 0.7287, 0.7187) ==\n FT(0.7287 * 0.7187 + (FT(1) - 0.7287) * 1.5)\n\n @test volumetric_internal_energy_liq(300.0, param_set) ==\n FT(_ρcp_l * (300.0 - _T_ref))\n\n @test k_solid(FT(0.5), FT(0.25), FT(2.0), FT(3.0), FT(2.0)) ==\n FT(2)^FT(0.5) * FT(2)^FT(0.25) * FT(3.0)^FT(0.25)\n\n @test ksat_frozen(FT(0.5), FT(0.1), FT(0.4)) ==\n FT(0.5)^FT(0.9) * FT(0.4)^FT(0.1)\n\n @test ksat_unfrozen(FT(0.5), FT(0.1), FT(0.4)) ==\n FT(0.5)^FT(0.9) * FT(0.4)^FT(0.1)\n\n @test k_dry(param_set, soil_param_functions) ==\n ((FT(0.053) * FT(0.1) - κ_air) * FT(0.8) + κ_air * FT(1.0)) /\n (FT(1.0) - (FT(1.0) - FT(0.053)) * FT(0.8))\nend\n" }, { "path": "test/Land/Model/test_overland_flow_analytic.jl", "content": "using MPI\nusing OrderedCollections\nusing StaticArrays\nusing Test\nusing Statistics\nusing DelimitedFiles\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SurfaceFlow\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\nusing ArtifactWrappers\n\n\n# Test that the land model with no surface flow works correctly\n@testset \"NoSurfaceFlow Model\" begin\n function init_land_model!(land, state, aux, localgeo, time) end\n ClimateMachine.init()\n FT = Float64\n\n soil_water_model = PrescribedWaterModel()\n soil_heat_model = PrescribedTemperatureModel()\n soil_param_functions = nothing\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n m_surface = NoSurfaceFlowModel()\n\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n surface = m_surface,\n source = sources,\n init_state_prognostic = init_land_model!,\n )\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(10)\n dt = FT(1)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n\n ClimateMachine.invoke!(solver_config;)\n\n t = ODESolvers.gettime(solver_config.solver)\n state_vars = SingleStackUtils.get_vars_from_nodal_stack(\n mygrid,\n Q,\n vars_state(m, Prognostic(), FT),\n )\n aux_vars = SingleStackUtils.get_vars_from_nodal_stack(\n mygrid,\n aux,\n vars_state(m, Auxiliary(), FT),\n )\n #Make sure it runs, and that there are no state variables, and only \"x,y,z\" as aux.\n @test t == timeend\n @test size(Base.collect(keys(aux_vars)))[1] == 3\n @test size(Base.collect(keys(state_vars)))[1] == 0\nend\n\n\n# Constant slope analytical test case defined as Model 1 / Eqn 6\n# DOI: 10.1061/(ASCE)0733-9429(2007)133:2(217)\n@testset \"Analytical Overland Model\" begin\n function warp_constant_slope(\n xin,\n yin,\n zin;\n topo_max = 0.2,\n zmin = -0.1,\n xmax = 400,\n )\n FT = eltype(xin)\n zmax = FT((FT(1.0) - xin / xmax) * topo_max)\n alpha = FT(1.0) - zmax / zmin\n zout = zmin + (zin - zmin) * alpha\n x, y, z = xin, yin, zout\n return x, y, z\n end\n ClimateMachine.init()\n FT = Float64\n\n soil_water_model = PrescribedWaterModel()\n soil_heat_model = PrescribedTemperatureModel()\n soil_param_functions = nothing\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n\n m_surface = OverlandFlowModel(\n (x, y) -> eltype(x)(-0.0016),\n (x, y) -> eltype(x)(0.0);\n mannings = (x, y) -> eltype(x)(0.025),\n )\n\n\n bc = LandDomainBC(\n minx_bc = LandComponentBC(\n surface = Dirichlet((aux, t) -> eltype(aux)(0)),\n ),\n )\n\n function init_land_model!(land, state, aux, localgeo, time)\n state.surface.height = eltype(state)(0)\n end\n\n # units in m / s \n precip(x, y, t) = t < (30 * 60) ? 1.4e-5 : 0.0\n\n sources = (Precip{FT}(precip),)\n\n m = LandModel(\n param_set,\n m_soil;\n surface = m_surface,\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_land_model!,\n )\n\n N_poly = 1\n xres = FT(2.286)\n yres = FT(0.25)\n zres = FT(0.1)\n # Specify the domain boundaries.\n zmax = FT(0)\n zmin = FT(-0.1)\n xmax = FT(182.88)\n ymax = FT(1.0)\n topo_max = FT(0.0016 * xmax)\n\n driver_config = ClimateMachine.MultiColumnLandModel(\n \"LandModel\",\n (N_poly, N_poly),\n (xres, yres, zres),\n xmax,\n ymax,\n zmax,\n param_set,\n m;\n zmin = zmin,\n meshwarp = (x...) -> warp_constant_slope(\n x...;\n topo_max = topo_max,\n zmin = zmin,\n xmax = xmax,\n ),\n )\n\n t0 = FT(0)\n timeend = FT(60 * 60)\n dt = FT(10)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n\n h_index = varsindex(vars_state(m, Prognostic(), FT), :surface, :height)\n n_outputs = 60\n\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n dons = Dict([k => Dict() for k in 1:n_outputs]...)\n\n iostep = [1]\n callback = GenericCallbacks.EveryXSimulationTime(\n every_x_simulation_time,\n ) do (init = false)\n t = ODESolvers.gettime(solver_config.solver)\n h = Q[:, h_index, :]\n all_vars = Dict{String, Array}(\"t\" => [t], \"h\" => h)\n dons[iostep[1]] = all_vars\n iostep[1] += 1\n return\n end\n\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n # Compare flowrate analytical derivation\n aux = solver_config.dg.state_auxiliary\n\n # get all nodal points at the max X bound of the domain\n mask = Array(aux[:, 1, :] .== 182.88)\n n_outputs = length(dons)\n # get prognostic variable height from nodal state (m^2)\n height = [mean(Array(dons[k][\"h\"])[mask[:]]) for k in 1:n_outputs]\n # get similation timesteps (s)\n time_data = [dons[l][\"t\"][1] for l in 1:n_outputs]\n\n\n alpha = sqrt(0.0016) / 0.025\n i = 1.4e-5\n L = xmax\n m = 5 / 3\n t_c = (L * i^(1 - m) / alpha)^(1 / m)\n t_r = 30 * 60\n\n q = height .^ (m) .* alpha\n\n function g(m, y, i, t_r, L, alpha, t)\n output = L / alpha - y^(m) / i - y^(m - 1) * m * (t - t_r)\n return output\n end\n\n function dg(m, y, i, t_r, L, alpha, t)\n output = -y^(m - 1) * m / i - y^(m - 2) * m * (m - 1) * (t - t_r)\n return output\n end\n\n function analytic(t, alpha, t_c, t_r, i, L, m)\n if t < t_c\n return alpha * (i * t)^(m)\n end\n\n if t <= t_r && t > t_c\n return alpha * (i * t_c)^(m)\n end\n\n if t > t_r\n yL = (i * (t - t_r))\n delta = 1\n error = g(m, yL, i, t_r, L, alpha, t)\n while abs(error) > 1e-4\n delta =\n -g(m, yL, i, t_r, L, alpha, t) /\n dg(m, yL, i, t_r, L, alpha, t)\n yL = yL + delta\n error = g(m, yL, i, t_r, L, alpha, t)\n end\n return alpha * yL^m\n\n end\n\n end\n\n q = Array(q) # copy to host if GPU array\n @test sqrt_rmse_over_max_q =\n sqrt(mean(\n (analytic.(time_data, alpha, t_c, t_r, i, L, m) .- q) .^ 2.0,\n )) / maximum(q) < 3e-3\nend\n" }, { "path": "test/Land/Model/test_overland_flow_vcatchment.jl", "content": "using Test\nusing Statistics\nusing DelimitedFiles\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SurfaceFlow\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\nusing ArtifactWrappers\n\n\n@testset \"V Catchment Maxwell River Model\" begin\n tv_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"tiltedv\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/qi2gftjw2vu2j66b0tyfef427xxj3ug7.csv\",\n filename = \"TiltedVOutput.csv\",\n ),],\n )\n tv_dataset_path = get_data_folder(tv_dataset)\n\n ClimateMachine.init()\n FT = Float64\n\n soil_water_model = PrescribedWaterModel()\n soil_heat_model = PrescribedTemperatureModel()\n soil_param_functions = nothing\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n\n function x_slope(x, y)\n MFT = eltype(x)\n if x < MFT(800)\n MFT(-0.05)\n elseif x <= MFT(820)\n MFT(0)\n else\n MFT(0.05)\n end\n end\n\n function y_slope(x, y)\n MFT = eltype(x)\n MFT(-0.02)\n end\n\n function channel_mannings(x, y)\n MFT = eltype(x)\n return x >= MFT(800) && x <= MFT(820) ? MFT(2.5 * 60 * 10^-3) :\n MFT(2.5 * 60 * 10^-4)\n end\n\n m_surface = OverlandFlowModel(x_slope, y_slope; mannings = channel_mannings)\n\n bc = LandDomainBC(\n miny_bc = LandComponentBC(\n surface = Dirichlet((aux, t) -> eltype(aux)(0)),\n ),\n minx_bc = LandComponentBC(\n surface = Dirichlet((aux, t) -> eltype(aux)(0)),\n ),\n maxx_bc = LandComponentBC(\n surface = Dirichlet((aux, t) -> eltype(aux)(0)),\n ),\n )\n\n function init_land_model!(land, state, aux, localgeo, time)\n state.surface.height = eltype(state)(0)\n end\n\n # units in m / s \n precip(x, y, t) = t < eltype(t)(90 * 60) ? eltype(t)(3e-6) : eltype(t)(0.0)\n\n sources = (Precip{FT}(precip),)\n\n m = LandModel(\n param_set,\n m_soil;\n surface = m_surface,\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_land_model!,\n )\n\n N_poly_hori = 1\n N_poly_vert = 1\n xres = FT(20)\n yres = FT(20)\n zres = FT(1)\n # Specify the domain boundaries.\n zmax = FT(1)\n zmin = FT(0)\n xmax = FT(1620)\n ymax = FT(1000)\n\n\n driver_config = ClimateMachine.MultiColumnLandModel(\n \"LandModel\",\n (N_poly_hori, N_poly_vert),\n (xres, yres, zres),\n xmax,\n ymax,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = RusanovNumericalFlux(),\n )\n\n t0 = FT(0)\n timeend = FT(180 * 60)\n dt = FT(0.5)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n Q = solver_config.Q\n\n h_index = varsindex(vars_state(m, Prognostic(), FT), :surface, :height)\n n_outputs = 60\n\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n\n dons = Dict([k => Dict() for k in 1:n_outputs]...)\n\n iostep = [1]\n callback = GenericCallbacks.EveryXSimulationTime(\n every_x_simulation_time,\n ) do (init = false)\n t = ODESolvers.gettime(solver_config.solver)\n h = Q[:, h_index, :]\n all_vars = Dict{String, Array}(\"t\" => [t], \"h\" => h)\n dons[iostep[1]] = all_vars\n iostep[1] += 1\n return\n end\n\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n aux = solver_config.dg.state_auxiliary\n x = Array(aux[:, 1, :])\n y = Array(aux[:, 2, :])\n z = Array(aux[:, 3, :])\n # Get points at outlet (y = ymax)\n mask2 = (Float64.(y .== 1000.0)) .== 1\n n_outputs = length(dons)\n function compute_Q(h, xv)\n height = max.(h, 0.0)\n v = calculate_velocity(m_surface, xv, 1000.0, height)# put in y = 1000.0\n speed = sqrt(v[1]^2.0 + v[2]^2.0 + v[3]^2.0)\n Q_outlet = speed .* height .* 60.0 # multiply by 60 so it is per minute, not per second\n return Q_outlet\n end\n # We divide by 4 because we have 4 nodal points with the same value at each x (z = 0, 1)\n # Multiply by xres because the solution at each point roughly represents that the solution for that range in x\n Q =\n [\n sum(compute_Q.(Array(dons[k][\"h\"])[:][mask2[:]], x[mask2[:]]))\n for k in 1:n_outputs\n ] ./ 4.0 .* xres\n data = joinpath(tv_dataset_path, \"TiltedVOutput.csv\")\n ds_tv = readdlm(data, ',')\n error = sqrt(mean(Q .- ds_tv))\n @test error < 5e-7\nend\n" }, { "path": "test/Land/Model/test_physical_bc.jl", "content": "# Test that the way we specify boundary conditions works as expected\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.Runoff\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\nimport ClimateMachine.DGMethods.FVReconstructions: FVLinear\n\n@testset \"NoRunoff\" begin\n ClimateMachine.init()\n\n FT = Float64\n\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n wpf = WaterParamFunctions(FT; Ksat = 1e-7, S_s = 1e-3)\n\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.75, water = wpf)\n surface_precip_amplitude = FT(3e-8)\n f = FT(pi * 2.0 / 300.0)\n precip = (t) -> surface_precip_amplitude * sin(f * t)\n ϑ_l0 = (aux) -> eltype(aux)(0.2)\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),\n ),\n surface_bc = LandComponentBC(\n soil_water = SurfaceDrivenWaterBoundaryConditions(\n FT;\n precip_model = DrivenConstantPrecip{FT}(precip),\n runoff_model = NoRunoff(),\n ),\n ),\n )\n\n soil_water_model = SoilWaterModel(FT; initialϑ_l = ϑ_l0)\n soil_heat_model = PrescribedTemperatureModel()\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n N_poly = 5\n nelem_vert = 50\n\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n\n t0 = FT(0)\n timeend = FT(150)\n dt = FT(0.05)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n n_outputs = 60\n mygrid = solver_config.dg.grid\n every_x_simulation_time = ceil(Int, timeend / n_outputs)\n state_types = (Prognostic(), Auxiliary(), GradientFlux())\n dons_arr =\n Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\n time_data = FT[0] # store time data\n\n # We specify a function which evaluates `every_x_simulation_time` and returns\n # the state vector, appending the variables we are interested in into\n # `dons_arr`.\n\n callback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n # # Run the integration\n ClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n\n # Get z-coordinate\n z = get_z(solver_config.dg.grid; rm_dupes = true)\n N = length(dons_arr)\n # Note - we take the indices 2:N here to avoid the t = 0 spot. Gradients\n # are not calculated before the invoke! command, so we cannot compare at t=0.\n computed_surface_flux = [dons_arr[k][\"soil.water.K∇h[3]\"][end] for k in 2:N]\n t = time_data[2:N]\n prescribed_surface_flux = t -> FT(-1) * FT(3e-8 * sin(pi * 2.0 * t / 300.0))\n MSE = mean((prescribed_surface_flux.(t) .- computed_surface_flux) .^ 2.0)\n @test MSE < 5e-7\nend\n\n\n\n@testset \"Explicit Dirichlet BC Comparison\" begin\n # This tests the \"if\" branch of compute_surface_grad_bc\n ClimateMachine.init()\n FT = Float64\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n Ksat = FT(0.0443 / (3600 * 100))\n S_s = FT(1e-4)\n wpf = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.495, water = wpf)\n surface_precip_amplitude = FT(-5e-4)\n\n precip = (t) -> surface_precip_amplitude\n hydraulics = vanGenuchten(FT; n = 2.0)\n function hydrostatic_profile(z, zm, porosity, n, α)\n myf = eltype(z)\n m = FT(1 - 1 / n)\n S = FT((FT(1) + (α * (z - zm))^n)^(-m))\n return FT(S * porosity)\n end\n N_poly = 2\n nelem_vert = 50\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-0.35)\n Δz = abs(FT((zmin - zmax) / nelem_vert / 2))\n\n ϑ_l0 =\n (aux) -> eltype(aux)(hydrostatic_profile(\n aux.z,\n -0.35,\n 0.495,\n hydraulics.n,\n hydraulics.α,\n ))\n\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = hydraulics,\n initialϑ_l = ϑ_l0,\n )\n soil_heat_model = PrescribedTemperatureModel()\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n\n\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),\n ),\n surface_bc = LandComponentBC(\n soil_water = SurfaceDrivenWaterBoundaryConditions(\n FT;\n precip_model = DrivenConstantPrecip{FT}(precip),\n runoff_model = CoarseGridRunoff{FT}(Δz),\n ),\n ),\n )\n\n\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n #fv_reconstruction = FVLinear(),\n )\n\n\n t0 = FT(0)\n timeend = FT(60 * 60)\n dt = FT(0.1)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n state_types = (Prognostic(), Auxiliary(), GradientFlux())\n\n ClimateMachine.invoke!(solver_config;)# user_callbacks = (callback,))\n srf_dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n\n\n ###### Repeat with explicit dirichlet BC\n soil_water_model2 = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = hydraulics,\n initialϑ_l = ϑ_l0,\n )\n soil_heat_model2 = PrescribedTemperatureModel()\n\n m_soil2 =\n SoilModel(soil_param_functions, soil_water_model2, soil_heat_model2)\n bc2 = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),\n ),\n surface_bc = LandComponentBC(\n soil_water = Dirichlet((aux, t) -> eltype(aux)(0.495)),\n ),\n )\n\n\n m2 = LandModel(\n param_set,\n m_soil2;\n boundary_conditions = bc2,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n driver_config2 = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m2;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n # fv_reconstruction = FVLinear(),\n )\n\n solver_config2 = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config2,\n ode_dt = dt,\n )\n\n ClimateMachine.invoke!(solver_config2;)\n dir_dons = dict_of_nodal_states(solver_config2, state_types; interp = true)\n # Here we take the solution near the surface, where changes between the two\n # would occur.\n @test mean(\n (\n srf_dons[\"soil.water.ϑ_l\"][(end - 15):end] .-\n dir_dons[\"soil.water.ϑ_l\"][(end - 15):end]\n ) .^ 2.0,\n ) < 5e-5\nend\n\n@testset \"Explicit Flux BC Comparison\" begin\n # This tests the else branch of compute_surface_grad_bc\n ClimateMachine.init()\n FT = Float64\n function init_soil_water!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n state.soil.water.ϑ_l = myfloat(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = myfloat(land.soil.water.initialθ_i(aux))\n end\n\n Ksat = FT(0.0443 / (3600 * 100))\n S_s = FT(1e-4)\n wpf = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s)\n\n soil_param_functions = SoilParamFunctions(FT; porosity = 0.495, water = wpf)\n surface_precip_amplitude = FT(0)\n\n precip = (t) -> surface_precip_amplitude\n hydraulics = vanGenuchten(FT; n = 2.0)\n function hydrostatic_profile(z, zm, porosity, n, α)\n myf = eltype(z)\n m = FT(1 - 1 / n)\n S = FT((FT(1) + (α * (z - zm))^n)^(-m))\n return FT(S * porosity)\n end\n\n ϑ_l0 =\n (aux) -> eltype(aux)(hydrostatic_profile(\n aux.z,\n -1,\n 0.495,\n hydraulics.n,\n hydraulics.α,\n ))\n\n soil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = hydraulics,\n initialϑ_l = ϑ_l0,\n )\n soil_heat_model = PrescribedTemperatureModel()\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n N_poly = (1, 0)\n nelem_vert = 200\n\n # Specify the domain boundaries\n zmax = FT(0)\n zmin = FT(-1)\n Δz = FT(1 / 200 / 2)\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_water = Neumann((aux, t) -> eltype(aux)(0.0)),\n ),\n surface_bc = LandComponentBC(\n soil_water = SurfaceDrivenWaterBoundaryConditions(\n FT;\n precip_model = DrivenConstantPrecip{FT}(precip),\n runoff_model = CoarseGridRunoff{FT}(Δz),\n ),\n ),\n )\n\n\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n )\n\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n fv_reconstruction = FVLinear(),\n )\n\n\n t0 = FT(0)\n timeend = FT(100000)\n dt = FT(4)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n state_types = (Prognostic(), Auxiliary(), GradientFlux())\n\n ClimateMachine.invoke!(solver_config;)\n\n srf_dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n # Note - we only look at differences between the solutions at the surface, because\n # near the bottom they will agree.\n z = srf_dons[\"z\"][150:end]\n error1 = sqrt(mean(\n (\n srf_dons[\"soil.water.ϑ_l\"][150:end] .-\n hydrostatic_profile.(z, -1, 0.495, hydraulics.n, hydraulics.α)\n ) .^ 2.0,\n ))\n error2 = sqrt(mean(srf_dons[\"soil.water.K∇h[3]\"][150:end] .^ 2.0))\n @test error1 < 1e-5\n @test error2 < eps(FT)\nend\n" }, { "path": "test/Land/Model/test_radiative_energy_flux_functions.jl", "content": "# Test functions used in runoff modeling.\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.RadiativeEnergyFlux\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Land.Runoff\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n@testset \"Radiative energy flux testing\" begin\n F = Float32\n user_nswf = t -> F(2 * t)\n user_swf = t -> F(2 * t)\n user_α = t -> F(0.2 * t)\n prescribed_swf_and_a =\n PrescribedSwFluxAndAlbedo(F; α = user_α, swf = user_swf)\n prescribed_nswf = PrescribedNetSwFlux(F; nswf = user_nswf)\n flux_from_prescribed_swf_and_albedo =\n compute_net_radiative_energy_flux.(\n Ref(prescribed_swf_and_a),\n [1, 2, 3, 4],\n )\n\n flux_from_prescribed_nswf =\n compute_net_radiative_energy_flux.(Ref(prescribed_nswf), [1, 2, 3, 4])\n\n @test flux_from_prescribed_swf_and_albedo ≈\n F.(([0.8, 0.6, 0.4, 0.2]) .* ([2, 4, 6, 8]))\n @test eltype(flux_from_prescribed_swf_and_albedo) == F\n\n @test flux_from_prescribed_nswf ≈ F.([2, 4, 6, 8])\n @test eltype(flux_from_prescribed_nswf) == F\nend\n\n@testset \"Heat analytic unit test\" begin\n ClimateMachine.init()\n\n FT = Float64\n\n function init_soil!(land, state, aux, localgeo, time)\n myfloat = eltype(state)\n ϑ_l, θ_i = get_water_content(land.soil.water, aux, state, time)\n θ_l =\n volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_s = volumetric_heat_capacity(\n θ_l,\n θ_i,\n land.soil.param_functions.ρc_ds,\n land.param_set,\n )\n\n state.soil.heat.ρe_int = myfloat(volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n land.param_set,\n ))\n end\n\n\n soil_param_functions = SoilParamFunctions(\n FT;\n porosity = 0.4,\n ν_ss_gravel = 0.2,\n ν_ss_om = 0.2,\n ν_ss_quartz = 0.2,\n ρc_ds = 1.0, # diffusivity = k/ρc. When no water, ρc = ρc_ds.\n κ_solid = 1.0,\n ρp = 1.0,\n κ_sat_unfrozen = 0.57,\n κ_sat_frozen = 2.29,\n )\n\n # Prescribed net short wave flux, initial temperature\n A = FT(5)\n prescribed_nswf = t -> FT(-A * t)\n T_init = FT(300.0)\n T_init_func = aux -> T_init\n\n # Flux entering = flux leaving soil column\n bc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_heat = Neumann((aux, t) -> FT(-A * t)),\n ),\n surface_bc = LandComponentBC(\n soil_heat = SurfaceDrivenHeatBoundaryConditions(\n FT;\n nswf_model = PrescribedNetSwFlux(FT; nswf = prescribed_nswf),\n ),\n ),\n )\n\n soil_water_model = PrescribedWaterModel()\n soil_heat_model = SoilHeatModel(FT; initialT = T_init_func)\n\n m_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model)\n sources = ()\n m = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n )\n\n N_poly = 5\n nelem_vert = 10\n\n # Specify the domain boundaries\n zmax = FT(1)\n zmin = FT(0)\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n )\n\n t0 = FT(0)\n timeend = FT(1)\n dt = FT(1e-4)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n mygrid = solver_config.dg.grid\n aux = solver_config.dg.state_auxiliary\n ClimateMachine.invoke!(solver_config)\n t = ODESolvers.gettime(solver_config.solver)\n\n z_ind = varsindex(vars_state(m, Auxiliary(), FT), :z)\n z = Array(aux[:, z_ind, :][:])\n\n T_ind = varsindex(vars_state(m, Auxiliary(), FT), :soil, :heat, :T)\n T = Array(aux[:, T_ind, :][:])\n\n k = k_dry(param_set, soil_param_functions)\n diffusivity = k / soil_param_functions.ρc_ds\n\n A = A / k # Because the flux is applied on κ∇T. κ∇T = A*t in clima, whereas ∇T = A*t in the analytic soln.\n approx_sum_term = sum(\n (\n -2 * A / (diffusivity * pi^4) *\n cos.(n * pi * z) *\n (pi * n * sin(pi * n) + cos(pi * n) - 1) *\n (1 - exp(-timeend * diffusivity * (pi * n)^2)) / n^4\n ) for n in 1:100\n )\n\n approx_analytic_soln =\n A * timeend .* z .+ T_init .- A * timeend / 2 .+ approx_sum_term\n\n MSE = mean((approx_analytic_soln .- T) .^ 2.0)\n\n @test eltype(aux) == FT\n @test MSE < 1e-5\nend\n" }, { "path": "test/Land/Model/test_water_parameterizations.jl", "content": "# Test the branching of the SoilWaterParameterizations functions,\n# test the instantiation of them works as expected, and that they\n# return what we expect.\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\n\n@testset \"Land water parameterizations\" begin\n FT = Float64\n test_array = [0.5, 1.0]\n vg_model = vanGenuchten(FT;)\n mm = MoistureDependent{FT}()\n bc_model = BrooksCorey(FT;)\n hk_model = Haverkamp(FT;)\n\n #Use an array to confirm that extra arguments are unused.\n @test viscosity_factor.(Ref(ConstantViscosity{FT}()), test_array) ≈\n [1.0, 1.0]\n @test impedance_factor.(Ref(NoImpedance{FT}()), test_array) ≈ [1.0, 1.0]\n @test moisture_factor.(\n Ref(MoistureIndependent{FT}()),\n Ref(vg_model),\n test_array,\n ) ≈ [1.0, 1.0]\n\n viscosity_model = TemperatureDependentViscosity{FT}(; T_ref = FT(1.0))\n @test viscosity_factor(viscosity_model, FT(1.0)) == 1\n impedance_model = IceImpedance{FT}(; Ω = 2.0)\n @test impedance_factor(impedance_model, 0.5) == FT(0.1)\n\n\n imp_f = impedance_factor(impedance_model, 0.5)\n vis_f = viscosity_factor(viscosity_model, 1.0)\n m_f = moisture_factor(MoistureDependent{FT}(), vg_model, 1.0)\n Ksat = 1.0\n @test hydraulic_conductivity(Ksat, imp_f, vis_f, m_f) == FT(0.1)\n\n @test_throws Exception effective_saturation(0.5, -1.0, 0.0)\n @test effective_saturation(0.5, 0.25, 0.05) - 4.0 / 9.0 < eps(FT)\n\n n = vg_model.n\n m = 1.0 - 1.0 / n\n α = vg_model.α\n S_s = 0.001\n θ_r = 0.0\n ν = 1.0\n @test moisture_factor.(Ref(mm), Ref(vg_model), test_array) ==\n sqrt.(test_array) .*\n (FT(1) .- (FT(1) .- test_array .^ (FT(1) / m)) .^ m) .^ FT(2)\n @test moisture_factor.(Ref(mm), Ref(bc_model), test_array) ==\n test_array .^ (FT(2) * bc_model.m + FT(3))\n ψ =\n -(\n (test_array .^ (-FT(1) / hk_model.m) .- FT(1)) .*\n hk_model.α^(-hk_model.n)\n ) .^ (FT(1) / hk_model.n)\n @test moisture_factor.(Ref(mm), Ref(hk_model), test_array) ==\n hk_model.A ./ (hk_model.A .+ abs.(ψ) .^ hk_model.k)\n\n @test pressure_head.(\n Ref(vg_model),\n Ref(ν),\n Ref(S_s),\n Ref(θ_r),\n test_array,\n Ref(0.0),\n ) ≈ .-((-1 .+ test_array .^ (-1 / m)) .* α^(-n)) .^ (1 / n)\n #test branching in pressure head\n @test pressure_head(vg_model, ν, S_s, θ_r, 1.5, 0.0) == 500\n @test pressure_head.(\n Ref(hk_model),\n Ref(ν),\n Ref(S_s),\n Ref(θ_r),\n test_array,\n Ref(0.0),\n ) ≈ .-((-1 .+ test_array .^ (-1 / m)) .* α^(-n)) .^ (1 / n)\n m = FT(0.5)\n ψb = FT(0.1656)\n @test pressure_head(bc_model, ν, S_s, θ_r, 0.5, 0.0) ≈ -ψb * 0.5^(-m)\n @test volumetric_liquid_fraction.([0.5, 1.5], Ref(0.75)) ≈ [0.5, 0.75]\n\n @test inverse_matric_potential(\n bc_model,\n matric_potential(bc_model, FT(0.5)),\n ) - FT(0.5) < eps(FT)\n @test inverse_matric_potential(\n vg_model,\n matric_potential(vg_model, FT(0.5)),\n ) == FT(0.5)\n @test inverse_matric_potential(\n hk_model,\n matric_potential(hk_model, FT(0.5)),\n ) == FT(0.5)\n\n @test_throws Exception inverse_matric_potential(bc_model, 1)\n @test_throws Exception inverse_matric_potential(hk_model, 1)\n @test_throws Exception inverse_matric_potential(vg_model, 1)\n\n\n # Test that spatial dependence works properly\n vg_spatial = vanGenuchten(FT; n = (x) -> 2 * x, α = (x) -> x^2)\n @test supertype(typeof(vg_spatial.n)) == Function\n vg_point = vg_spatial(FT(2.0))\n @test typeof(vg_point.n) == FT\n @test vg_point.n == 4.0\n @test vg_point.α == 4.0\n @test vg_point.m == 0.75\nend\n" }, { "path": "test/Land/runtests.jl", "content": "using Test, Pkg\n\n@testset \"Land\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"Model\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"Land/$submodule\" in ARGS ||\n \"Land\" in ARGS\n include_test(submodule)\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/Euler/acousticwave_1d_imex.jl", "content": "using ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies:\n StackedCubedSphereTopology, equiangular_cubed_sphere_warp, grid1d\nusing ClimateMachine.Mesh.Grids:\n DiscontinuousSpectralElementGrid, VerticalDirection\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods: DGModel, init_ode_state, remainder_DGModel\nusing ClimateMachine.DGMethods.NumericalFluxes:\n RusanovNumericalFlux,\n CentralNumericalFluxGradient,\n CentralNumericalFluxSecondOrder\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing Thermodynamics:\n air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition\nusing Thermodynamics.TemperatureProfiles: IsothermalProfile\nusing ClimateMachine.Atmos:\n AtmosPhysics,\n AtmosModel,\n DryModel,\n NoPrecipitation,\n NoRadiation,\n NTracers,\n vars_state,\n Gravity,\n HydrostaticState,\n AtmosAcousticGravityLinearModel,\n parameter_set\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.Orientations:\n SphericalOrientation, gravitational_potential, altitude, latitude, longitude\nusing ClimateMachine.VariableTemplates: flattenednames\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: planet_radius\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\nconst output_vtk = false\n\nconst ntracers = 1\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 5\n numelem_horz = 10\n numelem_vert = 5\n\n timeend = 60 * 60\n # timeend = 33 * 60 * 60 # Full simulation\n outputtime = 60 * 60\n\n expected_result = Dict()\n expected_result[Float32] = 9.5066030866432000e+13\n expected_result[Float64] = 9.5073452847149594e+13\n\n for FT in (Float64, Float32)\n for split_explicit_implicit in (false, true)\n result = test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n split_explicit_implicit,\n )\n @test result ≈ expected_result[FT]\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n split_explicit_implicit,\n)\n\n setup = AcousticWaveSetup{FT}()\n\n _planet_radius::FT = planet_radius(param_set)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + setup.domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n δ_χ = @SVector [FT(ii) for ii in 1:ntracers]\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(T_profile),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n tracers = NTracers{length(δ_χ), FT}(δ_χ),\n )\n model = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n linearmodel = AtmosAcousticGravityLinearModel(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n lineardg = DGModel(\n linearmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n\n # determine the time step\n element_size = (setup.domain_height / numelem_vert)\n acoustic_speed = soundspeed_air(param_set, FT(setup.T_ref))\n dt_factor = 445\n dt = dt_factor * element_size / acoustic_speed / polynomialorder^2\n # Adjust the time step so we exactly hit 1 hour for VTK output\n dt = 60 * 60 / ceil(60 * 60 / dt)\n nsteps = ceil(Int, timeend / dt)\n\n Q = init_ode_state(dg, FT(0))\n\n linearsolver = ManyColumnLU()\n # linearsolver = BatchedGeneralizedMinimalResidual(\n # lineardg,\n # Q;\n # atol = 1.0e-6, #sqrt(eps(FT)) * 0.01,\n # rtol = 1.0e-8, #sqrt(eps(FT)) * 0.01,\n # # Maximum number of Krylov iterations in a column\n #)\n\n\n if split_explicit_implicit\n rem_dg = remainder_DGModel(\n dg,\n (lineardg,);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (lineardg.numerical_flux_first_order,),\n ),\n )\n end\n odesolver = ARK2GiraldoKellyConstantinescu(\n split_explicit_implicit ? rem_dg : dg,\n lineardg,\n LinearBackwardEulerSolver(\n linearsolver;\n isadjustable = true,\n preconditioner_update_freq = -1,\n ),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = split_explicit_implicit,\n )\n @test getsteps(odesolver) == 0\n\n filterorder = 18\n filter = ExponentialFilter(grid, 0, filterorder)\n cbfilter = EveryXSimulationSteps(1) do\n Filters.apply!(Q, :, grid, filter, direction = VerticalDirection())\n nothing\n end\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n ArrayType = %s\n FT = %s\n polynomialorder = %d\n numelem_horz = %d\n numelem_vert = %d\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" \"$ArrayType\" \"$FT\" polynomialorder numelem_horz numelem_vert dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(odesolver) runtime energy\n end\n end\n callbacks = (cbinfo, cbfilter)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_acousticwave\" *\n \"_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)\" *\n \"_dt$(dt_factor)x_$(ArrayType)_$(FT)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(odesolver))\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(\n Q,\n odesolver;\n numberofsteps = nsteps,\n adjustfinalstep = false,\n callbacks = callbacks,\n )\n\n @test getsteps(odesolver) == nsteps\n\n # final statistics\n engf = norm(Q)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n \"\"\" engf engf / eng0 engf - eng0\n engf\nend\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n # callable to set initial conditions\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n\n state.tracers.ρχ = @SVector [FT(ii) for ii in 1:ntracers]\n nothing\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n testname = \"acousticwave\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(filename, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...), eltype(Q))\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/acousticwave_mrigark.jl", "content": "using ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies:\n StackedCubedSphereTopology, equiangular_cubed_sphere_warp, grid1d\nusing ClimateMachine.Mesh.Grids:\n DiscontinuousSpectralElementGrid,\n VerticalDirection,\n HorizontalDirection,\n EveryDirection,\n min_node_distance\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods: DGModel, init_ode_state, remainder_DGModel\nusing ClimateMachine.DGMethods.NumericalFluxes:\n RusanovNumericalFlux,\n CentralNumericalFluxGradient,\n CentralNumericalFluxSecondOrder\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing Thermodynamics:\n air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition\nusing Thermodynamics.TemperatureProfiles: IsothermalProfile\nusing ClimateMachine.Atmos:\n AtmosPhysics,\n AtmosModel,\n DryModel,\n NoPrecipitation,\n NoRadiation,\n NTracers,\n ConstantDynamicViscosity,\n vars_state,\n Gravity,\n HydrostaticState,\n AtmosAcousticGravityLinearModel,\n AtmosAcousticLinearModel,\n parameter_set\nusing ClimateMachine.Orientations:\n SphericalOrientation, gravitational_potential, altitude, latitude, longitude\nusing ClimateMachine.VariableTemplates: flattenednames\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: planet_radius\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\nconst output_vtk = false\n\nconst ntracers = 1\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 5\n numelem_horz = 10\n numelem_vert = 5\n\n timeend = 60 * 60\n # timeend = 33 * 60 * 60 # Full simulation\n outputtime = 60 * 60\n\n expected_result = Dict()\n expected_result[Float64, true] = 9.5073337869322578e+13\n expected_result[Float64, false] = 9.5073455070673781e+13\n\n @testset \"acoustic wave\" begin\n for FT in (Float64,)# Float32)\n for explicit in (true, false)\n result = test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n explicit,\n )\n @test result ≈ expected_result[FT, explicit]\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n explicit_solve,\n)\n\n setup = AcousticWaveSetup{FT}()\n\n _planet_radius::FT = planet_radius(param_set)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + setup.domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n hmnd = min_node_distance(grid, HorizontalDirection())\n vmnd = min_node_distance(grid, VerticalDirection())\n\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n δ_χ = @SVector [FT(ii) for ii in 1:ntracers]\n\n fullphysics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(T_profile),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n tracers = NTracers{length(δ_χ), FT}(δ_χ),\n )\n fullmodel = AtmosModel{FT}(\n AtmosLESConfigType,\n fullphysics;\n orientation = SphericalOrientation(),\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n dg = DGModel(\n fullmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n Q = init_ode_state(dg, FT(0))\n\n # The linear model which contains the fast modes\n # acousticmodel = AtmosAcousticLinearModel(fullmodel)\n acousticmodel = AtmosAcousticGravityLinearModel(fullmodel)\n\n vacoustic_dg = DGModel(\n acousticmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n\n # Advection model is the difference between the fullmodel and acousticmodel.\n # This will be handled with explicit substepping (time step in between the\n # vertical and horizontally acoustic models)\n rem_dg = remainder_DGModel(dg, (vacoustic_dg,))\n\n # determine the time step for the model components\n acoustic_speed = soundspeed_air(param_set, FT(setup.T_ref))\n advection_speed = 1 # What's a reasonable number here?\n\n vacoustic_dt = vmnd / acoustic_speed\n remainder_dt = min(min(hmnd, vmnd) / advection_speed, hmnd / acoustic_speed)\n odesolver = if explicit_solve\n remainder_dt = 5vacoustic_dt\n\n nsteps_output = ceil(outputtime / remainder_dt)\n remainder_dt = outputtime / nsteps_output\n nsteps = ceil(Int, timeend / remainder_dt)\n @assert nsteps * remainder_dt ≈ timeend\n\n vacoustic_solver =\n LSRK54CarpenterKennedy(vacoustic_dg, Q; dt = vacoustic_dt)\n rem_solver = MRIGARKERK45aSandu(\n rem_dg,\n vacoustic_solver,\n Q;\n dt = remainder_dt,\n )\n\n rem_solver\n else\n vacoustic_dt = 200vacoustic_dt\n element_size = (setup.domain_height / numelem_vert)\n\n nsteps_output = ceil(outputtime / vacoustic_dt)\n vacoustic_dt = outputtime / nsteps_output\n nsteps = ceil(Int, timeend / vacoustic_dt)\n @assert nsteps * vacoustic_dt ≈ timeend\n\n rem_solver = LSRK54CarpenterKennedy(rem_dg, Q; dt = remainder_dt)\n vacoustic_solver = MRIGARKESDIRK24LSA(\n vacoustic_dg,\n LinearBackwardEulerSolver(ManyColumnLU(); isadjustable = false),\n rem_solver,\n Q;\n dt = vacoustic_dt,\n )\n\n vacoustic_solver\n end\n\n\n filterorder = 18\n filter = ExponentialFilter(grid, 0, filterorder)\n cbfilter = EveryXSimulationSteps(1) do\n Filters.apply!(Q, :, grid, filter, direction = VerticalDirection())\n nothing\n end\n\n eng0 = norm(Q)\n @info @sprintf(\n \"\"\"Starting\n ArrayType = %s\n FT = %s\n polynomialorder = %d\n numelem_horz = %d\n numelem_vert = %d\n acoustic dt = %.16e\n remainder_dt = %.16e\n norm(Q₀) = %.16e\n \"\"\",\n \"$ArrayType\",\n \"$FT\",\n polynomialorder,\n numelem_horz,\n numelem_vert,\n vacoustic_dt,\n remainder_dt,\n eng0\n )\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(odesolver) runtime energy\n end\n end\n callbacks = (cbinfo, cbfilter)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_acousticwave\" *\n \"_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)\" *\n \"_$(ArrayType)_$(FT)\" *\n \"_$(explicit_solve ? \"Explicit_Explicit\" : \"Implicit_Explicit\")\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, fullmodel)\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(odesolver))\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, fullmodel)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(\n Q,\n odesolver;\n numberofsteps = nsteps,\n adjustfinalstep = false,\n callbacks = callbacks,\n )\n\n # final statistics\n engf = norm(Q)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n \"\"\" engf engf / eng0 engf - eng0\n engf\nend\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n # callable to set initial conditions\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n\n state.tracers.ρχ = @SVector [FT(ii) for ii in 1:ntracers]\n nothing\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n testname = \"acousticwave\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(filename, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...), eltype(Q))\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/acousticwave_variable_degree.jl", "content": "using ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies:\n StackedCubedSphereTopology, equiangular_cubed_sphere_warp, grid1d\nusing ClimateMachine.Mesh.Grids:\n DiscontinuousSpectralElementGrid, VerticalDirection\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods: DGModel, init_ode_state, remainder_DGModel\nusing ClimateMachine.DGMethods.NumericalFluxes:\n RusanovNumericalFlux,\n CentralNumericalFluxGradient,\n CentralNumericalFluxSecondOrder\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing Thermodynamics:\n air_density, soundspeed_air, internal_energy, PhaseDry_pT, PhasePartition\nusing Thermodynamics.TemperatureProfiles: IsothermalProfile\nusing ClimateMachine.Atmos:\n AtmosPhysics,\n AtmosModel,\n DryModel,\n NoPrecipitation,\n NoRadiation,\n NTracers,\n vars_state,\n Gravity,\n HydrostaticState,\n AtmosAcousticGravityLinearModel,\n parameter_set\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.Orientations:\n SphericalOrientation, gravitational_potential, altitude, latitude, longitude\nusing ClimateMachine.VariableTemplates: flattenednames\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: planet_radius\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\nconst ntracers = 1\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n # 5th order in the horizontal, cubic in the vertical\n polynomialorder = (5, 3)\n numelem_horz = 10\n numelem_vert = 6\n\n timeend = 60 * 60\n # timeend = 33 * 60 * 60 # Full simulation\n outputtime = 60 * 60\n\n expected_result = Dict()\n expected_result[Float32] = 9.5075065f13\n expected_result[Float64] = 9.507349773781483e13\n\n for FT in (Float64, Float32)\n for split_explicit_implicit in (false, true)\n result = test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n split_explicit_implicit,\n )\n @test result ≈ expected_result[FT]\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n split_explicit_implicit,\n)\n\n setup = AcousticWaveSetup{FT}()\n\n _planet_radius::FT = planet_radius(param_set)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + setup.domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n δ_χ = @SVector [FT(ii) for ii in 1:ntracers]\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(T_profile),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n tracers = NTracers{length(δ_χ), FT}(δ_χ),\n )\n model = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n\n linearmodel = AtmosAcousticGravityLinearModel(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n lineardg = DGModel(\n linearmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n\n # determine the time step\n element_size = (setup.domain_height / numelem_vert)\n acoustic_speed = soundspeed_air(param_set, FT(setup.T_ref))\n dt_factor = 445\n Nmax = maximum(polynomialorder)\n dt = dt_factor * element_size / acoustic_speed / Nmax^2\n # Adjust the time step so we exactly hit 1 hour for VTK output\n dt = 60 * 60 / ceil(60 * 60 / dt)\n nsteps = ceil(Int, timeend / dt)\n\n Q = init_ode_state(dg, FT(0))\n\n linearsolver = ManyColumnLU()\n\n if split_explicit_implicit\n rem_dg = remainder_DGModel(\n dg,\n (lineardg,);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (lineardg.numerical_flux_first_order,),\n ),\n )\n end\n odesolver = ARK2GiraldoKellyConstantinescu(\n split_explicit_implicit ? rem_dg : dg,\n lineardg,\n LinearBackwardEulerSolver(\n linearsolver;\n isadjustable = true,\n preconditioner_update_freq = -1,\n ),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = split_explicit_implicit,\n )\n @test getsteps(odesolver) == 0\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n ArrayType = %s\n FT = %s\n poly order horz = %d\n poly order vert = %d\n numelem_horz = %d\n numelem_vert = %d\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" \"$ArrayType\" \"$FT\" polynomialorder... numelem_horz numelem_vert dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(odesolver) runtime energy\n end\n end\n\n # Look ma, no filters!\n callbacks = (cbinfo,)\n\n solve!(\n Q,\n odesolver;\n numberofsteps = nsteps,\n adjustfinalstep = false,\n callbacks = callbacks,\n )\n\n @test getsteps(odesolver) == nsteps\n\n # final statistics\n engf = norm(Q)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n \"\"\" engf engf / eng0 engf - eng0\n return engf\nend\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n # callable to set initial conditions\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n\n state.tracers.ρχ = @SVector [FT(ii) for ii in 1:ntracers]\n nothing\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/fvm_balance.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.FVReconstructions: FVLinear\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: planet_radius\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing Test, MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = (4, 0)\n FT = Float64\n NumericalFlux = RoeNumericalFlux\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n FT = %s\n NumericalFlux = %s\n \"\"\" ArrayType FT NumericalFlux\n\n numelem_horz = 10\n numelem_vert = 32\n\n @testset for domain_type in (:box, :sphere)\n err = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n NumericalFlux,\n FT,\n domain_type,\n )\n @test err < FT(6e-12)\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n NumericalFlux,\n FT,\n domain_type,\n)\n domain_height = 10e3\n if domain_type === :box\n domain_width = 20e3\n horz_range =\n range(FT(0), length = numelem_horz + 1, stop = FT(domain_width))\n vert_range = range(0, length = numelem_vert + 1, stop = domain_height)\n brickrange = (horz_range, horz_range, vert_range)\n periodicity = (true, true, false)\n topology =\n StackedBrickTopology(mpicomm, brickrange; periodicity = periodicity)\n meshwarp = (x...) -> identity(x)\n elseif domain_type === :sphere\n _planet_radius::FT = planet_radius(param_set)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n meshwarp = equiangular_cubed_sphere_warp\n end\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = meshwarp,\n )\n\n T0 = FT(300)\n temp_profile = IsothermalProfile(param_set, T0)\n ref_state = HydrostaticState(temp_profile; subtract_off = false)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n\n problem = AtmosProblem(;\n physics = physics,\n init_state_prognostic = initialcondition!,\n )\n\n if domain_type === :box\n configtype = AtmosLESConfigType\n source = (Gravity(),)\n elseif domain_type === :sphere\n configtype = AtmosGCMConfigType\n source = (Gravity(), Coriolis())\n end\n\n model =\n AtmosModel{FT}(configtype, physics; problem = problem, source = source)\n\n dg = DGFVModel(\n model,\n grid,\n HBFVReconstruction(model, FVLinear()),\n NumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n timeend = FT(100)\n\n # determine the time step\n cfl = 0.2\n dz = step(vert_range)\n dt = cfl * dz / soundspeed_air(param_set, T0)\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dg, FT(0))\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n domain_type = %s\n numelem_horz = %d\n numelem_vert = %d\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" domain_type numelem_horz numelem_vert dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @views begin\n ρu = extrema(Array(Q.data[:, 2, :]))\n ρv = extrema(Array(Q.data[:, 3, :]))\n ρw = extrema(Array(Q.data[:, 4, :]))\n end\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n ρu = %.16e, %.16e\n ρv = %.16e, %.16e\n ρw = %.16e, %.16e\n norm(Q) = %.16e\n \"\"\" gettime(lsrk) runtime ρu... ρv... ρw... energy\n end\n end\n callbacks = (cbinfo,)\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n errr = errf / engfe\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errr\n errr\nend\n\nfunction initialcondition!(problem, bl, state, aux, coords, t, args...)\n state.ρ = aux.ref_state.ρ\n state.ρu = SVector(0, 0, 0)\n state.energy.ρe = aux.ref_state.ρe\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/fvm_isentropicvortex.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n\n # Just to make it shorter and aligning\n Roe = RoeNumericalFlux\n\n # Float64, Dim 2, degree 4 in the horizontal, FV order 1, refinement level\n expected_error[Float64, FVConstant(), Roe, 1] = 3.5317756615538940e+01\n expected_error[Float64, FVConstant(), Roe, 2] = 2.5104217086071472e+01\n expected_error[Float64, FVConstant(), Roe, 3] = 1.6169569358521223e+01\n expected_error[Float64, FVConstant(), Roe, 4] = 9.4731749125284708e+00\n\n expected_error[Float64, FVLinear(), Roe, 1] = 2.6132907774912638e+01\n expected_error[Float64, FVLinear(), Roe, 2] = 8.8198392537283006e+00\n expected_error[Float64, FVLinear(), Roe, 3] = 2.4517109427575416e+00\n expected_error[Float64, FVLinear(), Roe, 4] = 7.4154384427579900e-01\n\n\n dims = 2\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,)\n for NumericalFlux in (Roe,)\n setup = IsentropicVortexSetup{FT}()\n for fvmethod in (FVConstant(), FVLinear())\n @info @sprintf \"\"\"Configuration\n FT = %s\n ArrayType = %s\n FV Reconstruction = %s\n NumericalFlux = %s\n dims = %d\n \"\"\" FT ArrayType fvmethod NumericalFlux dims\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n\n # Match element numbers\n numelems = (\n 2^(level - 1) * 5,\n 2^(level - 1) * 5 * polynomialorder,\n )\n\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n )\n\n @test errors[level] ≈\n expected_error[FT, fvmethod, NumericalFlux, level]\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(errors[l]) - log2(errors[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n connectivity = dims == 3 ? :full : :face\n\n topology = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n connectivity = connectivity,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = (polynomialorder, 0),\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n NumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n timeend = FT(2 * setup.domain_halflength / 10 / setup.translation_speed)\n\n # Determine the time step\n elementsize = minimum(step.(brickrange))\n dt = elementsize / soundspeed_air(param_set, setup.T∞) / polynomialorder^2\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dgfvm, FT(0))\n lsrk = LSRK54CarpenterKennedy(dgfvm, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(lsrk) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # Create vtk dir\n vtkdir =\n \"vtk_isentropicvortex\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Q, model)\n\n # Setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(lsrk), setup)\n do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Final statistics\n Qe = init_ode_state(dgfvm, timeend, setup)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex\",\n)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init(parse_clargs = true)\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n\n # just to make it shorter and aligning\n Rusanov = RusanovNumericalFlux()\n Central = CentralNumericalFluxFirstOrder()\n Roe = RoeNumericalFlux()\n HLLC = HLLCNumericalFlux()\n RoeMoist = RoeNumericalFluxMoist()\n RoeMoistLM = RoeNumericalFluxMoist(; LM = true)\n RoeMoistHH = RoeNumericalFluxMoist(; HH = true)\n RoeMoistLV = RoeNumericalFluxMoist(; LV = true)\n RoeMoistLVPP = RoeNumericalFluxMoist(; LVPP = true)\n\n expected_error[Float64, 2, Rusanov, 1] = 1.1990999506538110e+01\n expected_error[Float64, 2, Rusanov, 2] = 2.0813000228865612e+00\n expected_error[Float64, 2, Rusanov, 3] = 6.3752572004789149e-02\n expected_error[Float64, 2, Rusanov, 4] = 2.0984975076420455e-03\n\n expected_error[Float64, 2, Central, 1] = 2.0840574601661153e+01\n expected_error[Float64, 2, Central, 2] = 2.9255455365299827e+00\n expected_error[Float64, 2, Central, 3] = 3.6935849488949657e-01\n expected_error[Float64, 2, Central, 4] = 8.3528804679907434e-03\n\n expected_error[Float64, 2, Roe, 1] = 1.2891386634733328e+01\n expected_error[Float64, 2, Roe, 2] = 1.3895805145495934e+00\n expected_error[Float64, 2, Roe, 3] = 6.6174934435569849e-02\n expected_error[Float64, 2, Roe, 4] = 2.1917769287815940e-03\n\n expected_error[Float64, 2, RoeMoist, 1] = 1.2415957884123003e+01\n expected_error[Float64, 2, RoeMoist, 2] = 1.4188653323424882e+00\n expected_error[Float64, 2, RoeMoist, 3] = 6.7913325894248130e-02\n expected_error[Float64, 2, RoeMoist, 4] = 2.0377963111049128e-03\n\n expected_error[Float64, 2, RoeMoistLM, 1] = 1.2316906651180444e+01\n expected_error[Float64, 2, RoeMoistLM, 2] = 1.4359406523244560e+00\n expected_error[Float64, 2, RoeMoistLM, 3] = 6.8650238542101505e-02\n expected_error[Float64, 2, RoeMoistLM, 4] = 2.0000156591586842e-03\n\n expected_error[Float64, 2, RoeMoistHH, 1] = 1.2425625606467793e+01\n expected_error[Float64, 2, RoeMoistHH, 2] = 1.4029458093339020e+00\n expected_error[Float64, 2, RoeMoistHH, 3] = 6.8648208937091268e-02\n expected_error[Float64, 2, RoeMoistHH, 4] = 2.0711985861781648e-03\n\n expected_error[Float64, 2, RoeMoistLV, 1] = 1.2415957884123003e+01\n expected_error[Float64, 2, RoeMoistLV, 2] = 1.4188653323424882e+00\n expected_error[Float64, 2, RoeMoistLV, 3] = 6.7913325894248130e-02\n expected_error[Float64, 2, RoeMoistLV, 4] = 2.0377963111049128e-03\n\n expected_error[Float64, 2, RoeMoistLVPP, 1] = 1.2441813136310969e+01\n expected_error[Float64, 2, RoeMoistLVPP, 2] = 2.0219325767566727e+00\n expected_error[Float64, 2, RoeMoistLVPP, 3] = 6.7716921628626484e-02\n expected_error[Float64, 2, RoeMoistLVPP, 4] = 2.1051129944994005e-03\n\n expected_error[Float64, 2, HLLC, 1] = 1.2889756097329746e+01\n expected_error[Float64, 2, HLLC, 2] = 1.3895808565455936e+00\n expected_error[Float64, 2, HLLC, 3] = 6.6175116756217900e-02\n expected_error[Float64, 2, HLLC, 4] = 2.1917772135679118e-03\n\n expected_error[Float64, 3, Rusanov, 1] = 3.7918869862613858e+00\n expected_error[Float64, 3, Rusanov, 2] = 6.5816485664822677e-01\n expected_error[Float64, 3, Rusanov, 3] = 2.0160333422867591e-02\n expected_error[Float64, 3, Rusanov, 4] = 6.6360317881818034e-04\n\n expected_error[Float64, 3, Central, 1] = 6.5903683487905749e+00\n expected_error[Float64, 3, Central, 2] = 9.2513872939749997e-01\n expected_error[Float64, 3, Central, 3] = 1.1680141169828175e-01\n expected_error[Float64, 3, Central, 4] = 2.6414127301659534e-03\n\n expected_error[Float64, 3, Roe, 1] = 4.0766143963611068e+00\n expected_error[Float64, 3, Roe, 2] = 4.3942394181655547e-01\n expected_error[Float64, 3, Roe, 3] = 2.0926351682882375e-02\n expected_error[Float64, 3, Roe, 4] = 6.9310072176312712e-04\n\n expected_error[Float64, 3, RoeMoist, 1] = 3.9262706246552574e+00\n expected_error[Float64, 3, RoeMoist, 2] = 4.4868461432545598e-01\n expected_error[Float64, 3, RoeMoist, 3] = 2.1476079330305119e-02\n expected_error[Float64, 3, RoeMoist, 4] = 6.4440777504566171e-04\n\n expected_error[Float64, 3, RoeMoistLM, 1] = 3.8949478745407458e+00\n expected_error[Float64, 3, RoeMoistLM, 2] = 4.5408430461734528e-01\n expected_error[Float64, 3, RoeMoistLM, 3] = 2.1709111570716800e-02\n expected_error[Float64, 3, RoeMoistLM, 4] = 6.3246048386171073e-04\n\n expected_error[Float64, 3, RoeMoistHH, 1] = 3.9293278268948622e+00\n expected_error[Float64, 3, RoeMoistHH, 2] = 4.4365041912830866e-01\n expected_error[Float64, 3, RoeMoistHH, 3] = 2.1708469753267460e-02\n expected_error[Float64, 3, RoeMoistHH, 4] = 6.5497050185153694e-04\n\n expected_error[Float64, 3, RoeMoistLV, 1] = 3.9262706246552574e+00\n expected_error[Float64, 3, RoeMoistLV, 2] = 4.4868461432545598e-01\n expected_error[Float64, 3, RoeMoistLV, 3] = 2.1476079330305119e-02\n expected_error[Float64, 3, RoeMoistLV, 4] = 6.4440777504566171e-04\n\n expected_error[Float64, 3, RoeMoistLVPP, 1] = 3.9344467732944071e+00\n expected_error[Float64, 3, RoeMoistLVPP, 2] = 6.3939122178436703e-01\n expected_error[Float64, 3, RoeMoistLVPP, 3] = 2.1413970848172485e-02\n expected_error[Float64, 3, RoeMoistLVPP, 4] = 6.6569517942933988e-04\n\n expected_error[Float64, 3, HLLC, 1] = 4.0760987751605402e+00\n expected_error[Float64, 3, HLLC, 2] = 4.3942404996518236e-01\n expected_error[Float64, 3, HLLC, 3] = 2.0926409337758904e-02\n expected_error[Float64, 3, HLLC, 4] = 6.9310081182571569e-04\n\n expected_error[Float32, 2, Rusanov, 1] = 1.1990781784057617e+01\n expected_error[Float32, 2, Rusanov, 2] = 2.0813269615173340e+00\n expected_error[Float32, 2, Rusanov, 3] = 6.7035309970378876e-02\n expected_error[Float32, 2, Rusanov, 4] = 5.3008597344160080e-02\n\n expected_error[Float32, 2, Central, 1] = 2.0840391159057617e+01\n expected_error[Float32, 2, Central, 2] = 2.9256355762481689e+00\n expected_error[Float32, 2, Central, 3] = 3.7092915177345276e-01\n expected_error[Float32, 2, Central, 4] = 1.1543693393468857e-01\n\n expected_error[Float32, 2, Roe, 1] = 1.2891359329223633e+01\n expected_error[Float32, 2, Roe, 2] = 1.3895936012268066e+00\n expected_error[Float32, 2, Roe, 3] = 6.8037144839763641e-02\n expected_error[Float32, 2, Roe, 4] = 3.8893952965736389e-02\n\n expected_error[Float32, 2, RoeMoist, 1] = 1.2415886878967285e+01\n expected_error[Float32, 2, RoeMoist, 2] = 1.4188879728317261e+00\n expected_error[Float32, 2, RoeMoist, 3] = 6.9743692874908447e-02\n expected_error[Float32, 2, RoeMoist, 4] = 3.7607192993164063e-02\n\n expected_error[Float32, 2, RoeMoistLM, 1] = 1.2316809654235840e+01\n expected_error[Float32, 2, RoeMoistLM, 2] = 4.5408430461734528e+00\n expected_error[Float32, 2, RoeMoistLM, 3] = 7.0370830595493317e-02\n expected_error[Float32, 2, RoeMoistLM, 4] = 3.7792034447193146e-02\n\n expected_error[Float32, 2, RoeMoistHH, 1] = 1.2425449371337891e+01\n expected_error[Float32, 2, RoeMoistHH, 2] = 1.4030106067657471e+00\n expected_error[Float32, 2, RoeMoistHH, 3] = 7.0363849401473999e-02\n expected_error[Float32, 2, RoeMoistHH, 4] = 3.7904966622591019e-02\n\n expected_error[Float32, 2, RoeMoistLV, 1] = 1.2415886878967285e+01\n expected_error[Float32, 2, RoeMoistLV, 2] = 1.4188879728317261e+00\n expected_error[Float32, 2, RoeMoistLV, 3] = 6.9743692874908447e-02\n expected_error[Float32, 2, RoeMoistLV, 4] = 3.7607192993164063e-02\n\n expected_error[Float32, 2, RoeMoistLVPP, 1] = 1.2441481590270996e+01\n expected_error[Float32, 2, RoeMoistLVPP, 2] = 2.0217459201812744e+00\n expected_error[Float32, 2, RoeMoistLVPP, 3] = 7.0483185350894928e-02\n expected_error[Float32, 2, RoeMoistLVPP, 4] = 5.1601748913526535e-02\n\n expected_error[Float32, 2, HLLC, 1] = 1.2889801025390625e+01\n expected_error[Float32, 2, HLLC, 2] = 1.3895059823989868e+00\n expected_error[Float32, 2, HLLC, 3] = 6.8006515502929688e-02\n expected_error[Float32, 2, HLLC, 4] = 3.8637656718492508e-02\n\n expected_error[Float32, 3, Rusanov, 1] = 3.7918186187744141e+00\n expected_error[Float32, 3, Rusanov, 2] = 6.5816193819046021e-01\n expected_error[Float32, 3, Rusanov, 3] = 2.0893247798085213e-02\n expected_error[Float32, 3, Rusanov, 4] = 1.1554701253771782e-02\n\n expected_error[Float32, 3, Central, 1] = 6.5903329849243164e+00\n expected_error[Float32, 3, Central, 2] = 9.2512512207031250e-01\n expected_error[Float32, 3, Central, 3] = 1.1707859486341476e-01\n expected_error[Float32, 3, Central, 4] = 2.1001411601901054e-02\n\n expected_error[Float32, 3, Roe, 1] = 4.0765657424926758e+00\n expected_error[Float32, 3, Roe, 2] = 4.3941807746887207e-01\n expected_error[Float32, 3, Roe, 3] = 2.1365188062191010e-02\n expected_error[Float32, 3, Roe, 4] = 9.3323951587080956e-03\n\n expected_error[Float32, 3, RoeMoist, 1] = 3.9262301921844482e+00\n expected_error[Float32, 3, RoeMoist, 2] = 4.4864514470100403e-01\n expected_error[Float32, 3, RoeMoist, 3] = 2.1889146417379379e-02\n expected_error[Float32, 3, RoeMoist, 4] = 8.8266804814338684e-03\n\n expected_error[Float32, 3, RoeMoistLM, 1] = 3.8948786258697510e+00\n expected_error[Float32, 3, RoeMoistLM, 2] = 4.5405751466751099e-01\n expected_error[Float32, 3, RoeMoistLM, 3] = 2.2112159058451653e-02\n expected_error[Float32, 3, RoeMoistLM, 4] = 8.7371272966265678e-03\n\n expected_error[Float32, 3, RoeMoistHH, 1] = 3.9292929172515869e+00\n expected_error[Float32, 3, RoeMoistHH, 2] = 4.4363334774971008e-01\n expected_error[Float32, 3, RoeMoistHH, 3] = 2.2118536755442619e-02\n expected_error[Float32, 3, RoeMoistHH, 4] = 8.9262928813695908e-03\n\n expected_error[Float32, 3, RoeMoistLV, 1] = 3.9262151718139648e+00\n expected_error[Float32, 3, RoeMoistLV, 2] = 4.4865489006042480e-01\n expected_error[Float32, 3, RoeMoistLV, 3] = 2.1889505907893181e-02\n expected_error[Float32, 3, RoeMoistLV, 4] = 8.8385939598083496e-03\n\n expected_error[Float32, 3, RoeMoistLVPP, 1] = 3.9343423843383789e+00\n expected_error[Float32, 3, RoeMoistLVPP, 2] = 6.3935810327529907e-01\n expected_error[Float32, 3, RoeMoistLVPP, 3] = 2.1930629387497902e-02\n expected_error[Float32, 3, RoeMoistLVPP, 4] = 1.0632344521582127e-02\n\n expected_error[Float32, 3, HLLC, 1] = 4.0760631561279297e+00\n expected_error[Float32, 3, HLLC, 2] = 4.3940672278404236e-01\n expected_error[Float32, 3, HLLC, 3] = 2.1352596580982208e-02\n expected_error[Float32, 3, HLLC, 4] = 9.2315869405865669e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32), dims in (2, 3)\n for NumericalFlux in (\n Rusanov,\n Central,\n Roe,\n HLLC,\n RoeMoist,\n RoeMoistLM,\n RoeMoistHH,\n RoeMoistLV,\n RoeMoistLVPP,\n )\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n FT = %s\n NumericalFlux = %s\n dims = %d\n \"\"\" ArrayType \"$FT\" \"$NumericalFlux\" dims\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n )\n\n rtol = sqrt(eps(FT))\n # increase rtol for comparing with GPU results using Float32\n if FT === Float32 && ArrayType !== Array\n rtol *= 10 # why does this factor have to be so big :(\n end\n @test isapprox(\n errors[level],\n expected_error[FT, dims, NumericalFlux, level];\n rtol = rtol,\n )\n end\n\n rates = @. log2(\n first(errors[1:(numlevels - 1)]) /\n first(errors[2:numlevels]),\n )\n numlevels > 1 && @info \"Convergence rates\\n\" * join(\n [\n \"rate for levels $l → $(l + 1) = $(rates[l])\"\n for l in 1:(numlevels - 1)\n ],\n \"\\n\",\n )\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n if NumericalFlux isa RoeNumericalFluxMoist\n moisture = EquilMoist()\n else\n moisture = DryModel()\n end\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = moisture,\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n\n dg = DGModel(\n model,\n grid,\n NumericalFlux,\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n timeend = FT(2 * setup.domain_halflength / 10 / setup.translation_speed)\n\n # determine the time step\n elementsize = minimum(step.(brickrange))\n dt = elementsize / soundspeed_air(param_set, setup.T∞) / polynomialorder^2\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dg, FT(0))\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(lsrk) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk), setup)\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend, setup)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_imex.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n\n expected_error[Float64, false, 1] = 2.3225467541870387e+01\n expected_error[Float64, false, 2] = 5.2663709730295070e+00\n expected_error[Float64, false, 3] = 1.2183770894070467e-01\n expected_error[Float64, false, 4] = 2.8660813871243937e-03\n\n expected_error[Float64, true, 1] = 2.3225467618783981e+01\n expected_error[Float64, true, 2] = 5.2663709730207771e+00\n expected_error[Float64, true, 3] = 1.2183770891083319e-01\n expected_error[Float64, true, 4] = 2.8660813810759854e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,), dims in 2\n for split_explicit_implicit in (false, true)\n let\n split =\n split_explicit_implicit ? \"(Nonlinear, Linear)\" :\n \"(Full, Linear)\"\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n FT = %s\n dims = %d\n splitting = %s\n \"\"\" ArrayType \"$FT\" dims split\n end\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n split_explicit_implicit,\n FT,\n dims,\n level,\n )\n\n @test errors[level] ≈\n expected_error[FT, split_explicit_implicit, level]\n end\n\n rates = @. log2(\n first(errors[1:(numlevels - 1)]) /\n first(errors[2:numlevels]),\n )\n numlevels > 1 && @info \"Convergence rates\\n\" * join(\n [\n \"rate for levels $l → $(l + 1) = $(rates[l])\"\n for l in 1:(numlevels - 1)\n ],\n \"\\n\",\n )\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n split_explicit_implicit,\n FT,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = IsentropicVortexReferenceState{FT}(setup),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n\n linear_model = AtmosAcousticLinearModel(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n dg_linear = DGModel(\n linear_model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n state_auxiliary = dg.state_auxiliary,\n )\n\n if split_explicit_implicit\n dg_nonlinear = remainder_DGModel(dg, (dg_linear,))\n end\n\n timeend = FT(2 * setup.domain_halflength / setup.translation_speed)\n\n # determine the time step\n elementsize = minimum(step.(brickrange))\n dt = elementsize / soundspeed_air(param_set, setup.T∞) / polynomialorder^2\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dg, FT(0))\n\n linearsolver = GeneralizedMinimalResidual(Q; M = 10, rtol = 1e-10)\n ode_solver = ARK2GiraldoKellyConstantinescu(\n split_explicit_implicit ? dg_nonlinear : dg,\n dg_linear,\n LinearBackwardEulerSolver(linearsolver; isadjustable = true),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = split_explicit_implicit,\n paperversion = true,\n )\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(ode_solver) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex_imex\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)_level$(level)\" *\n \"_$(split_explicit_implicit)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, ode_solver; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex_imex\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_lmars.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init(parse_clargs = true)\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 4\n\n expected_error = Dict()\n\n # just to make it shorter and aligning\n LMARS = LMARSNumericalFlux()\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,), dims in (2, 3), polynomialorder in (4,)\n for NumericalFlux in (LMARS,)\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n FT = %s\n NumericalFlux = %s\n dims = %d\n N_poly = %d\n \"\"\" ArrayType \"$FT\" \"$NumericalFlux\" dims polynomialorder\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n )\n\n @test isapprox(errors[level], FT(1.0); rtol = 1e-5)\n\n end\n\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n NumericalFlux,\n setup,\n FT,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n if NumericalFlux isa RoeNumericalFluxMoist\n moisture = EquilMoist()\n else\n moisture = DryModel()\n end\n\n if NumericalFlux isa LMARSNumericalFlux\n ref_state = NoReferenceState()\n end\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = moisture,\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n\n dg = DGModel(\n model,\n grid,\n NumericalFlux,\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n timeend = FT(2 * setup.domain_halflength / 10 / setup.translation_speed)\n\n # determine the time step\n elementsize = minimum(step.(brickrange))\n dt = elementsize / soundspeed_air(param_set, setup.T∞) / polynomialorder^2\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dg, FT(0))\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n polyorder = %d\n numelems = (%d, %d, %d)\n dt = %.16e\n norm(Q₀) = %.16e\n FT = %s\n \"\"\" level polynomialorder numelems... dt eng0 FT\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(lsrk) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex\" *\n \"$(typeof(NumericalFlux))\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend / 10\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk), setup)\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend, setup)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n engf / eng0\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex_lmars\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_mrigark.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n expected_error[Float64, MRIGARKERK33aSandu, 1] = 2.3357934866477940e+01\n expected_error[Float64, MRIGARKERK33aSandu, 2] = 5.3328129440361121e+00\n expected_error[Float64, MRIGARKERK33aSandu, 3] = 1.2991232057877919e-01\n expected_error[Float64, MRIGARKERK33aSandu, 4] = 6.1056067013518876e-03\n expected_error[Float64, MRIGARKERK45aSandu, 1] = 2.3207510164213900e+01\n expected_error[Float64, MRIGARKERK45aSandu, 2] = 5.2787446598866872e+00\n expected_error[Float64, MRIGARKERK45aSandu, 3] = 1.2151170640665301e-01\n expected_error[Float64, MRIGARKERK45aSandu, 4] = 2.1001271191583956e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,), dims in 2\n for mrigark_method in (MRIGARKERK33aSandu, MRIGARKERK45aSandu)\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n mrigark_method = %s\n FT = %s\n dims = %d\n \"\"\" ArrayType \"$mrigark_method\" \"$FT\" dims\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n FT,\n mrigark_method,\n dims,\n level,\n )\n\n @test errors[level] ≈\n expected_error[FT, mrigark_method, level]\n end\n\n rates = @. log2(\n first(errors[1:(numlevels - 1)]) /\n first(errors[2:numlevels]),\n )\n numlevels > 1 && @info \"Convergence rates\\n\" * join(\n [\n \"rate for levels $l → $(l + 1) = $(rates[l])\"\n for l in 1:(numlevels - 1)\n ],\n \"\\n\",\n )\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n FT,\n mrigark_method,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = IsentropicVortexReferenceState{FT}(setup),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n # The linear model has the fast time scales\n fast_model = AtmosAcousticLinearModel(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n fast_dg = DGModel(\n fast_model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n state_auxiliary = dg.state_auxiliary,\n )\n slow_dg = remainder_DGModel(dg, (fast_dg,))\n\n timeend = FT(2 * setup.domain_halflength / setup.translation_speed)\n # determine the slow time step\n elementsize = minimum(step.(brickrange))\n slow_dt =\n 2 * elementsize / soundspeed_air(param_set, setup.T∞) /\n polynomialorder^2\n nsteps = ceil(Int, timeend / slow_dt)\n slow_dt = timeend / nsteps\n\n # arbitrary and not needed for stability, just for testing\n fast_dt = slow_dt / 3\n\n Q = init_ode_state(dg, FT(0), setup)\n\n fastsolver = LSRK144NiegemannDiehlBusch(fast_dg, Q; dt = fast_dt)\n\n ode_solver = mrigark_method(slow_dg, fastsolver, Q, dt = slow_dt)\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n slow_dt = %.16e\n fast_dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... slow_dt fast_dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(ode_solver) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex_mrigark\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)\" *\n \"_$(FastMethod)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / slow_dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, ode_solver; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex_mrigark\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_mrigark_implicit.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n expected_error[Float64, MRIGARKIRK21aSandu, 1] = 2.3236071337679274e+01\n expected_error[Float64, MRIGARKIRK21aSandu, 2] = 5.2652585224989430e+00\n expected_error[Float64, MRIGARKIRK21aSandu, 3] = 1.2100430848052603e-01\n expected_error[Float64, MRIGARKIRK21aSandu, 4] = 2.1974838909870273e-03\n\n expected_error[Float64, MRIGARKESDIRK34aSandu, 1] = 2.3235626679098608e+01\n expected_error[Float64, MRIGARKESDIRK34aSandu, 2] = 5.2672845223341218e+00\n expected_error[Float64, MRIGARKESDIRK34aSandu, 3] = 1.2097276468825705e-01\n expected_error[Float64, MRIGARKESDIRK34aSandu, 4] = 2.0920468129065205e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,), dims in 2\n for mrigark_method in (MRIGARKIRK21aSandu, MRIGARKESDIRK34aSandu)\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n mrigark_method = %s\n FT = %s\n dims = %d\n \"\"\" ArrayType \"$mrigark_method\" \"$FT\" dims\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n mrigark_method,\n FT,\n dims,\n level,\n )\n\n @test errors[level] ≈\n expected_error[FT, mrigark_method, level]\n end\n\n rates = @. log2(\n first(errors[1:(numlevels - 1)]) /\n first(errors[2:numlevels]),\n )\n numlevels > 1 && @info \"Convergence rates\\n\" * join(\n [\n \"rate for levels $l → $(l + 1) = $(rates[l])\"\n for l in 1:(numlevels - 1)\n ],\n \"\\n\",\n )\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n mrigark_method,\n FT,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = IsentropicVortexReferenceState{FT}(setup),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n # This is a bad idea; this test is just testing how\n # implicit GARK composes with explicit methods\n # The linear model has the fast time scales but will be\n # treated implicitly (outer solver)\n slow_model = AtmosAcousticLinearModel(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n slow_dg = DGModel(\n slow_model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n state_auxiliary = dg.state_auxiliary,\n )\n fast_dg = remainder_DGModel(dg, (slow_dg,))\n\n timeend = FT(2 * setup.domain_halflength / setup.translation_speed)\n\n # determine the time step\n elementsize = minimum(step.(brickrange))\n dt =\n elementsize / soundspeed_air(param_set, setup.T∞) / polynomialorder^2 /\n 5\n nsteps = ceil(Int, timeend / dt)\n dt = timeend / nsteps\n\n Q = init_ode_state(dg, FT(0))\n\n fastsolver = LSRK54CarpenterKennedy(fast_dg, Q; dt = dt)\n\n linearsolver = GeneralizedMinimalResidual(Q; M = 50, rtol = 1e-10)\n\n ode_solver = mrigark_method(\n slow_dg,\n LinearBackwardEulerSolver(linearsolver; isadjustable = true),\n fastsolver,\n Q;\n dt = dt,\n t0 = 0,\n )\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(ode_solver) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex_mrigark\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)\" *\n \"_$(FastMethod)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, ode_solver; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex_mrigark\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_multirate.jl", "content": "using ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing ClimateMachine.SystemSolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\n\ninclude(\"isentropicvortex_setup.jl\")\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output_vtk = false\n\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n numlevels = integration_testing ? 4 : 1\n\n expected_error = Dict()\n expected_error[Float64, SSPRK33ShuOsher, 1] = 2.3222373077778794e+01\n expected_error[Float64, SSPRK33ShuOsher, 2] = 5.2782503174265516e+00\n expected_error[Float64, SSPRK33ShuOsher, 3] = 1.2281763287878383e-01\n expected_error[Float64, SSPRK33ShuOsher, 4] = 2.3761870907666096e-03\n\n expected_error[Float64, ARK2GiraldoKellyConstantinescu, 1] =\n 2.3245216640111998e+01\n expected_error[Float64, ARK2GiraldoKellyConstantinescu, 2] =\n 5.2626584944153949e+00\n expected_error[Float64, ARK2GiraldoKellyConstantinescu, 3] =\n 1.2324230746483673e-01\n expected_error[Float64, ARK2GiraldoKellyConstantinescu, 4] =\n 3.8777995619211627e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,), dims in 2\n for FastMethod in (SSPRK33ShuOsher, ARK2GiraldoKellyConstantinescu)\n @info @sprintf \"\"\"Configuration\n ArrayType = %s\n FastMethod = %s\n FT = %s\n dims = %d\n \"\"\" ArrayType \"$FastMethod\" \"$FT\" dims\n\n setup = IsentropicVortexSetup{FT}()\n errors = Vector{FT}(undef, numlevels)\n\n for level in 1:numlevels\n numelems =\n ntuple(dim -> dim == 3 ? 1 : 2^(level - 1) * 5, dims)\n errors[level] = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n FT,\n FastMethod,\n dims,\n level,\n )\n\n @test errors[level] ≈ expected_error[FT, FastMethod, level]\n end\n\n rates = @. log2(\n first(errors[1:(numlevels - 1)]) /\n first(errors[2:numlevels]),\n )\n numlevels > 1 && @info \"Convergence rates\\n\" * join(\n [\n \"rate for levels $l → $(l + 1) = $(rates[l])\"\n for l in 1:(numlevels - 1)\n ],\n \"\\n\",\n )\n end\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n numelems,\n setup,\n FT,\n FastMethod,\n dims,\n level,\n)\n brickrange = ntuple(dims) do dim\n range(\n -setup.domain_halflength;\n length = numelems[dim] + 1,\n stop = setup.domain_halflength,\n )\n end\n\n topology = BrickTopology(\n mpicomm,\n brickrange;\n periodicity = ntuple(_ -> true, dims),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n problem =\n AtmosProblem(boundaryconditions = (), init_state_prognostic = setup)\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = IsentropicVortexReferenceState{FT}(setup),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (),\n )\n # The linear model has the fast time scales\n fast_model = AtmosAcousticLinearModel(model)\n # The nonlinear model has the slow time scales\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n fast_dg = DGModel(\n fast_model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n state_auxiliary = dg.state_auxiliary,\n )\n slow_dg = remainder_DGModel(dg, (fast_dg,))\n\n timeend = FT(2 * setup.domain_halflength / setup.translation_speed)\n # determine the slow time step\n elementsize = minimum(step.(brickrange))\n slow_dt =\n 8 * elementsize / soundspeed_air(param_set, setup.T∞) /\n polynomialorder^2\n nsteps = ceil(Int, timeend / slow_dt)\n slow_dt = timeend / nsteps\n\n # arbitrary and not needed for stability, just for testing\n fast_dt = slow_dt / 3\n\n Q = init_ode_state(dg, FT(0), setup)\n\n slow_ode_solver = LSRK144NiegemannDiehlBusch(slow_dg, Q; dt = slow_dt)\n\n # check if FastMethod is ARK, is there a better way ?\n if FastMethod == ARK2GiraldoKellyConstantinescu\n\n linearsolver = GeneralizedMinimalResidual(Q; M = 10, rtol = 1e-10)\n # splitting the fast part into full and linear but the fast part\n # is already linear so full_dg == linear_dg == fast_dg\n fast_ode_solver = FastMethod(\n fast_dg,\n fast_dg,\n LinearBackwardEulerSolver(linearsolver; isadjustable = true),\n Q;\n dt = fast_dt,\n paperversion = true,\n )\n else\n fast_ode_solver = FastMethod(fast_dg, Q; dt = fast_dt)\n end\n\n ode_solver = MultirateRungeKutta((slow_ode_solver, fast_ode_solver))\n\n eng0 = norm(Q)\n dims == 2 && (numelems = (numelems..., 0))\n @info @sprintf \"\"\"Starting refinement level %d\n numelems = (%d, %d, %d)\n slow_dt = %.16e\n fast_dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" level numelems... slow_dt fast_dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(ode_solver) runtime energy\n end\n end\n callbacks = (cbinfo,)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_isentropicvortex_multirate\" *\n \"_poly$(polynomialorder)_dims$(dims)_$(ArrayType)_$(FT)\" *\n \"_$(FastMethod)_level$(level)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model)\n\n # setup the output callback\n outputtime = timeend\n cbvtk = EveryXSimulationSteps(floor(outputtime / slow_dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver), setup)\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, ode_solver; timeend = timeend, callbacks = callbacks)\n\n # final statistics\n Qe = init_ode_state(dg, timeend, setup)\n engf = norm(Q)\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished refinement level %d\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" level engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n testname = \"isentropicvortex_mutirate\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/Euler/isentropicvortex_setup.jl", "content": "import ClimateMachine.Atmos: atmos_init_aux!, vars_state\n\nBase.@kwdef struct IsentropicVortexSetup{FT}\n p∞::FT = 10^5\n T∞::FT = 300\n ρ∞::FT = air_density(param_set, FT(T∞), FT(p∞))\n translation_speed::FT = 150\n translation_angle::FT = pi / 4\n vortex_speed::FT = 50\n vortex_radius::FT = 1 // 200\n domain_halflength::FT = 1 // 20\nend\n\nfunction (setup::IsentropicVortexSetup)(\n problem,\n bl,\n state,\n aux,\n localgeo,\n t,\n args...,\n)\n FT = eltype(state)\n x = MVector(localgeo.coord)\n param_set = parameter_set(bl)\n\n ρ∞ = setup.ρ∞\n p∞ = setup.p∞\n T∞ = setup.T∞\n translation_speed = setup.translation_speed\n α = setup.translation_angle\n vortex_speed = setup.vortex_speed\n R = setup.vortex_radius\n L = setup.domain_halflength\n\n u∞ = SVector(translation_speed * cos(α), translation_speed * sin(α), 0)\n\n x .-= u∞ * t\n # make the function periodic\n x .-= floor.((x .+ L) / 2L) * 2L\n\n @inbounds begin\n r = sqrt(x[1]^2 + x[2]^2)\n δu_x = -vortex_speed * x[2] / R * exp(-(r / R)^2 / 2)\n δu_y = vortex_speed * x[1] / R * exp(-(r / R)^2 / 2)\n end\n u = u∞ .+ SVector(δu_x, δu_y, 0)\n\n _kappa_d::FT = kappa_d(param_set)\n T = T∞ * (1 - _kappa_d * vortex_speed^2 / 2 * ρ∞ / p∞ * exp(-(r / R)^2))\n # adiabatic/isentropic relation\n p = p∞ * (T / T∞)^(FT(1) / _kappa_d)\n ρ = air_density(param_set, T, p)\n\n state.ρ = ρ\n state.ρu = ρ * u\n e_kin = u' * u / 2\n state.energy.ρe = ρ * total_energy(param_set, e_kin, FT(0), T)\n if !(moisture_model(bl) isa DryModel)\n state.moisture.ρq_tot = FT(0)\n end\nend\n\nstruct IsentropicVortexReferenceState{FT} <: ReferenceState\n setup::IsentropicVortexSetup{FT}\nend\nvars_state(::IsentropicVortexReferenceState, ::Auxiliary, FT) =\n @vars(ρ::FT, ρe::FT, p::FT, T::FT)\nfunction atmos_init_aux!(\n atmos::AtmosModel,\n m::IsentropicVortexReferenceState,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n atmos,\n (args...) -> init_vortex_ref_state!(m, args...),\n state_auxiliary,\n grid,\n direction,\n )\nend\nfunction init_vortex_ref_state!(\n m::IsentropicVortexReferenceState,\n atmos::AtmosModel,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n setup = m.setup\n ρ∞ = setup.ρ∞\n p∞ = setup.p∞\n T∞ = setup.T∞\n param_set = parameter_set(atmos)\n aux.ref_state.ρ = ρ∞\n aux.ref_state.p = p∞\n aux.ref_state.T = T∞\n aux.ref_state.ρe = ρ∞ * internal_energy(param_set, T∞)\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model.jl", "content": "using StaticArrays\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.BalanceLaws:\n BalanceLaw,\n Prognostic,\n Auxiliary,\n Gradient,\n GradientFlux,\n GradientLaplacian,\n Hyperdiffusive\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n number_states,\n flux_first_order!,\n flux_second_order!,\n source!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n nodal_init_state_auxiliary!,\n update_auxiliary_state!,\n init_state_prognostic!,\n boundary_conditions,\n boundary_state!,\n wavespeed,\n transform_post_gradient_laplacian!\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.DGMethods: SpaceDiscretization\nusing ClimateMachine.DGMethods.NumericalFluxes:\n NumericalFluxFirstOrder,\n NumericalFluxSecondOrder,\n NumericalFluxGradient,\n CentralNumericalFluxDivergence,\n CentralNumericalFluxHigherOrder\n\nimport ClimateMachine.DGMethods.NumericalFluxes:\n numerical_flux_first_order!, boundary_flux_second_order!\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nstruct Advection{N} <: BalanceLaw end\nstruct NoAdvection <: BalanceLaw end\n\nstruct Diffusion{N} <: BalanceLaw end\nstruct NoDiffusion <: BalanceLaw end\n\nstruct HyperDiffusion{N} <: BalanceLaw end\nstruct NoHyperDiffusion <: BalanceLaw end\n\nabstract type AdvectionDiffusionProblem end\n\n# Boundary condition types\n\n# boundary condition for operator of order O\n# O = 0 -> state BC (Dirichlet)\n# O = 1 -> gradient BC (Neumann)\n# O = 2 -> laplacian BC\n# O = 3 -> gradient laplacian BC\nabstract type AbstractBC{O} end\nstruct HomogeneousBC{O} <: AbstractBC{O} end\nstruct InhomogeneousBC{O} <: AbstractBC{O} end\n\nany_isa(bcs::AbstractBC, bc) = bcs isa bc\nany_isa(bcs::Tuple, bc) = mapreduce(x -> x isa bc, |, bcs)\n\n\"\"\"\n AdvectionDiffusion{N} <: BalanceLaw\n\nA balance law describing a system of `N` advection-diffusion-hyperdiffusion\nequations:\n\n```\n∂ρ\n-- = - ∇ • (u ρ - σ + η)\n∂t\n\nσ = D ∇ ρ\nη = H ∇ Δρ\n```\nWhere\n\n - `ρ` is the solution vector\n - `u` is the advection velocity\n - `σ` is the DG diffusion auxiliary variable\n - `D` is the diffusion tensor\n - `η` is the DG hyperdiffusion auxiliary variable\n - `H` is the hyperdiffusion tensor\n\"\"\"\nstruct AdvectionDiffusion{N, dim, P, fluxBC, A, D, HD, BC} <: BalanceLaw\n problem::P\n advection::A\n diffusion::D\n hyperdiffusion::HD\n boundary_conditions::BC\n\n function AdvectionDiffusion{dim}(\n problem::P,\n boundary_conditions::BC = ();\n num_equations = 1,\n flux_bc = false,\n advection::Bool = true,\n diffusion::Bool = true,\n hyperdiffusion::Bool = false,\n ) where {dim, P <: AdvectionDiffusionProblem, BC}\n N = num_equations\n adv = advection ? Advection{N}() : NoAdvection()\n A = typeof(adv)\n diff = diffusion ? Diffusion{N}() : NoDiffusion()\n D = typeof(diff)\n hyperdiff = hyperdiffusion ? HyperDiffusion{N}() : NoHyperDiffusion()\n HD = typeof(hyperdiff)\n new{N, dim, P, flux_bc, A, D, HD, BC}(\n problem,\n adv,\n diff,\n hyperdiff,\n boundary_conditions,\n )\n end\nend\n\n# Auxiliary variables, always store\n# `coord` coordinate points (needed for BCs)\nfunction vars_state(m::AdvectionDiffusion, st::Auxiliary, FT)\n @vars begin\n coord::SVector{3, FT}\n advection::vars_state(m.advection, st, FT)\n diffusion::vars_state(m.diffusion, st, FT)\n hyperdiffusion::vars_state(m.hyperdiffusion, st, FT)\n end\nend\n\n# `u` advection velocity\nvars_state(::Advection{1}, ::Auxiliary, FT) = @vars(u::SVector{3, FT})\nvars_state(::Advection{N}, ::Auxiliary, FT) where {N} =\n @vars(u::SMatrix{3, N, FT, 3N})\n# `D` diffusion tensor\nvars_state(::Diffusion{1}, ::Auxiliary, FT) = @vars(D::SMatrix{3, 3, FT, 9})\nvars_state(::Diffusion{N}, ::Auxiliary, FT) where {N} =\n @vars(D::SArray{Tuple{3, 3, N}, FT, 3, 9N})\n# `H` hyperdiffusion tensor\nvars_state(::HyperDiffusion{1}, ::Auxiliary, FT) =\n @vars(H::SMatrix{3, 3, FT, 9})\nvars_state(::HyperDiffusion{N}, ::Auxiliary, FT) where {N} =\n @vars(H::SArray{Tuple{3, 3, N}, FT, 3, 9N})\n\n# Density `ρ` is the only state\nvars_state(::AdvectionDiffusion{1}, ::Prognostic, FT) = @vars(ρ::FT)\nvars_state(::AdvectionDiffusion{N}, ::Prognostic, FT) where {N} =\n @vars(ρ::SVector{N, FT})\n\nfunction vars_state(m::AdvectionDiffusion{N}, ::Gradient, FT) where {N}\n # For pure advection we don't need gradients\n if m.diffusion isa NoDiffusion && m.hyperdiffusion isa NoHyperDiffusion\n return @vars()\n else # Take the gradient of density\n return N == 1 ? @vars(ρ::FT) : @vars(ρ::SVector{N, FT})\n end\nend\n\n# Take the gradient of laplacian of density ρ\nvars_state(::HyperDiffusion{1}, ::GradientLaplacian, FT) = @vars(ρ::FT)\nvars_state(::HyperDiffusion{N}, ::GradientLaplacian, FT) where {N} =\n @vars(ρ::SVector{N, FT})\nvars_state(m::AdvectionDiffusion, st::GradientLaplacian, FT) =\n vars_state(m.hyperdiffusion, st, FT)\n\n# The DG diffusion auxiliary variable: σ = D ∇ρ\nvars_state(::Diffusion{1}, ::GradientFlux, FT) = @vars(σ::SVector{3, FT})\nvars_state(::Diffusion{N}, ::GradientFlux, FT) where {N} =\n @vars(σ::SMatrix{3, N, FT, 3N})\nvars_state(m::AdvectionDiffusion, st::GradientFlux, FT) =\n vars_state(m.diffusion, st, FT)\n\n# The DG hyperdiffusion auxiliary variable: η = H ∇ Δρ\nvars_state(::HyperDiffusion{1}, ::Hyperdiffusive, FT) = @vars(η::SVector{3, FT})\nvars_state(::HyperDiffusion{N}, ::Hyperdiffusive, FT) where {N} =\n @vars(η::SMatrix{3, N, FT, 3N})\nvars_state(m::AdvectionDiffusion, st::Hyperdiffusive, FT) =\n vars_state(m.hyperdiffusion, st, FT)\n\n\"\"\"\n flux_first_order!(::Advection, flux::Grad, state::Vars, aux::Vars)\n\nComputes non-diffusive flux `F_adv = u ρ` where\n\n - `u` is the advection velocity\n - `ρ` is the advected quantity\n\"\"\"\nfunction flux_first_order!(\n ::Advection{N},\n flux::Grad,\n state::Vars,\n aux::Vars,\n) where {N}\n ρ = state.ρ\n u = aux.advection.u\n flux.ρ += u .* ρ'\nend\nflux_first_order!(::NoAdvection, flux::Grad, state::Vars, aux::Vars) = nothing\nflux_first_order!(\n m::AdvectionDiffusion,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n) = flux_first_order!(m.advection, flux, state, aux)\n\n\"\"\"\n flux_second_order!(::Diffusion, flux::Grad, auxDG::Vars)\n\nComputes diffusive flux `F_diff = -σ` where:\n\n - `σ` is DG diffusion auxiliary variable (`σ = D ∇ ρ`\n with `D` being the diffusion tensor)\n\"\"\"\nfunction flux_second_order!(::Diffusion, flux::Grad, auxDG::Vars)\n σ = auxDG.σ\n flux.ρ += -σ\nend\nflux_second_order!(::NoDiffusion, flux::Grad, auxDG::Vars) = nothing\n\n\"\"\"\n flux_second_order!(::HyperDiffusion, flux::Grad, auxHDG::Vars)\n\nComputes hyperdiffusive flux `F_hyperdiff = η` where:\n\n - `η` is DG hyperdiffusion auxiliary variable (`η = H ∇ Δρ`\n with `H` being the hyperdiffusion tensor)\n\"\"\"\nfunction flux_second_order!(::HyperDiffusion, flux::Grad, auxHDG::Vars)\n η = auxHDG.η\n flux.ρ += η\nend\nflux_second_order!(::NoHyperDiffusion, flux::Grad, auxHDG::Vars) = nothing\n\nfunction flux_second_order!(\n m::AdvectionDiffusion,\n flux::Grad,\n state::Vars,\n auxDG::Vars,\n auxHDG::Vars,\n aux::Vars,\n t::Real,\n)\n flux_second_order!(m.diffusion, flux, auxDG)\n flux_second_order!(m.hyperdiffusion, flux, auxHDG)\nend\n\n\n\"\"\"\n compute_gradient_argument!(m::AdvectionDiffusion, transform::Vars, state::Vars,\n aux::Vars, t::Real)\n\nSet the variable to take the gradient of (`ρ` in this case)\n\"\"\"\nfunction compute_gradient_argument!(\n m::AdvectionDiffusion,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρ = state.ρ\nend\n\n\"\"\"\n compute_gradient_flux!(::Diffusion, auxDG::Vars, gradvars::Grad, aux::Vars)\n\nComputes the DG diffusion auxiliary variable `σ = D ∇ ρ` where `D` is\nthe diffusion tensor.\n\"\"\"\nfunction compute_gradient_flux!(\n ::Diffusion{N},\n auxDG::Vars,\n gradvars::Grad,\n aux::Vars,\n) where {N}\n ∇ρ = gradvars.ρ\n D = aux.diffusion.D\n if N == 1\n auxDG.σ = D * ∇ρ\n else\n auxDG.σ = hcat(ntuple(n -> D[:, :, n] * ∇ρ[:, n], Val(N))...)\n end\nend\ncompute_gradient_flux!(::NoDiffusion, auxDG::Vars, gradvars::Grad, aux::Vars) =\n nothing\ncompute_gradient_flux!(\n m::AdvectionDiffusion,\n auxDG::Vars,\n gradvars::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) = compute_gradient_flux!(m.diffusion, auxDG, gradvars, aux)\n\n\n\"\"\"\n transform_post_gradient_laplacian!(::AdvectionDiffusion, auxHDG::Vars,\n gradvars::Grad, state::Vars, aux::Vars, t::Real)\n\nComputes the DG hyperdiffusion auxiliary variable `η = H ∇ Δρ` where `H` is\nthe hyperdiffusion tensor.\n\"\"\"\nfunction transform_post_gradient_laplacian!(\n m::AdvectionDiffusion{N},\n auxHDG::Vars,\n gradvars::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {N}\n ∇Δρ = gradvars.ρ\n H = aux.hyperdiffusion.H\n if N == 1\n auxHDG.η = H * ∇Δρ\n else\n auxHDG.η = hcat(ntuple(n -> H[:, :, n] * ∇Δρ[:, n], Val(N))...)\n end\nend\n\n\"\"\"\n source!(m::AdvectionDiffusion, _...)\n\nThere is no source in the advection-diffusion model\n\"\"\"\nsource!(m::AdvectionDiffusion, _...) = nothing\n\n\"\"\"\n wavespeed(m::AdvectionDiffusion, nM, state::Vars, aux::Vars, t::Real)\n\nWavespeed with respect to vector `nM`\n\"\"\"\nwavespeed(\n m::AdvectionDiffusion,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n) = wavespeed(m.advection, nM, aux)\nfunction wavespeed(::Advection{N}, nM, aux::Vars) where {N}\n u = aux.advection.u\n if N == 1\n abs(nM' * u)\n else\n SVector(ntuple(n -> abs(nM' * u[:, n]), Val(N)))\n end\nend\nwavespeed(::NoAdvection, nM, aux::Vars) = 0\n\nfunction nodal_init_state_auxiliary!(\n m::AdvectionDiffusion,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.coord = geom.coord\n init_velocity_diffusion!(m.problem, aux, geom)\nend\n\nhas_variable_coefficients(::AdvectionDiffusionProblem) = false\nfunction update_auxiliary_state!(\n spacedisc::SpaceDiscretization,\n m::AdvectionDiffusion,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n if has_variable_coefficients(m.problem)\n update_auxiliary_state!(spacedisc, m, Q, t, elems) do m, state, aux, t\n update_velocity_diffusion!(m.problem, m, state, aux, t)\n end\n return true\n end\n return false\nend\n\nfunction init_state_prognostic!(\n m::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n initial_condition!(m.problem, state, aux, localgeo, t)\nend\n\n\"\"\"\n inhomogeneous_data!(::Val{O}, problem, data, aux, x, t)\n\nPrescribes `problem` boundary condition data for an operator of order `O`\n\"\"\"\nfunction inhomogeneous_data! end\n\nboundary_conditions(m::AdvectionDiffusion) = m.boundary_conditions\nfunction boundary_state!(\n nf,\n bcs,\n m::AdvectionDiffusion{N},\n stateP::Vars,\n auxP::Vars,\n nM,\n stateM::Vars,\n auxM::Vars,\n t,\n _...,\n) where {N}\n if any_isa(bcs, InhomogeneousBC{0}) # Dirichlet\n inhomogeneous_data!(\n Val(0),\n m.problem,\n stateP,\n auxP,\n (coord = auxP.coord,),\n t,\n )\n elseif any_isa(bcs, AbstractBC{1}) # Neumann\n stateP.ρ = stateM.ρ\n elseif any_isa(bcs, HomogeneousBC{0}) # zero Dirichlet\n stateP.ρ = N == 1 ? 0 : zeros(typeof(stateP.ρ))\n end\nend\n\nfunction boundary_state!(\n nf::CentralNumericalFluxSecondOrder,\n bcs,\n m::AdvectionDiffusion,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻::SVector,\n state⁻::Vars,\n diff⁻::Vars,\n hyperdiff⁻::Vars,\n aux⁻::Vars,\n t,\n _...,\n)\n if m.diffusion isa NoDiffusion && m.hyperdiffusion isa NoHyperDiffusion\n return nothing\n end\n\n if m.diffusion isa Diffusion\n if any_isa(bcs, AbstractBC{0}) # Dirchlet\n # Just use the minus side values since Dirchlet\n diff⁺.σ = diff⁻.σ\n elseif any_isa(bcs, InhomogeneousBC{1}) # Neumann with data\n FT = eltype(diff⁺)\n ngrad = number_states(m, Gradient())\n ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(\n parent(diff⁺),\n Size(3, ngrad),\n ))\n # Get analytic gradient\n inhomogeneous_data!(Val(1), m.problem, ∇state, aux⁻, aux⁻.coord, t)\n compute_gradient_flux!(m.diffusion, diff⁺, ∇state, aux⁻)\n # compute the diffusive flux using the boundary state\n elseif any_isa(bcs, HomogeneousBC{1}) # zero Neumann\n FT = eltype(diff⁺)\n ngrad = number_states(m, Gradient())\n ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(\n parent(diff⁺),\n Size(3, ngrad),\n ))\n # Get analytic gradient\n ∇state.ρ = zeros(typeof(∇state.ρ))\n # convert to auxDG variables\n compute_gradient_flux!(m.diffusion, diff⁺, ∇state, aux⁻)\n end\n end\n\n if m.hyperdiffusion isa HyperDiffusion\n if any_isa(bcs, InhomogeneousBC{3})\n FT = eltype(hyperdiff⁺)\n ngradlap = number_states(m, GradientLaplacian())\n ∇Δstate = Grad{vars_state(m, GradientLaplacian(), FT)}(similar(\n parent(hyperdiff⁺),\n Size(3, ngradlap),\n ))\n # Get analytic gradient of laplacian\n inhomogeneous_data!(Val(3), m.problem, ∇Δstate, aux⁻, aux⁻.coord, t)\n transform_post_gradient_laplacian!(\n m,\n hyperdiff⁺,\n ∇Δstate,\n state⁻,\n aux⁻,\n t,\n )\n elseif any_isa(bcs, HomogeneousBC{3})\n FT = eltype(hyperdiff⁺)\n ngradlap = number_states(m, GradientLaplacian())\n ∇Δstate =\n Grad{vars_state(m, GradientLaplacian(), FT)}(zeros(SMatrix{\n 3,\n ngradlap,\n FT,\n }))\n transform_post_gradient_laplacian!(\n m,\n hyperdiff⁺,\n ∇Δstate,\n state⁻,\n aux⁻,\n t,\n )\n end\n end\n nothing\nend\n\nfunction boundary_flux_second_order!(\n nf::CentralNumericalFluxSecondOrder,\n bcs,\n m::AdvectionDiffusion{N, dim, P, true},\n F,\n state⁺,\n diff⁺,\n hyperdiff⁺,\n aux⁺,\n n⁻,\n state⁻,\n diff⁻,\n hyperdiff⁻,\n aux⁻,\n t,\n _...,\n) where {N, dim, P}\n if m.diffusion isa NoDiffusion && m.hyperdiffusion isa NoHyperDiffusion\n return nothing\n end\n\n # Default initialize flux to minus side\n if any_isa(bcs, AbstractBC{0}) # Dirchlet\n # Just use the minus side values since Dirchlet\n flux_second_order!(m, F, state⁻, diff⁻, hyperdiff⁻, aux⁻, t)\n elseif any_isa(bcs, InhomogeneousBC{1}) # Neumann data\n FT = eltype(diff⁺)\n ngrad = number_states(m, Gradient())\n ∇state = Grad{vars_state(m, Gradient(), FT)}(similar(\n parent(diff⁺),\n Size(3, ngrad),\n ))\n # Get analytic gradient\n inhomogeneous_data!(Val(1), m.problem, ∇state, aux⁻, aux⁻.coord, t)\n # get the diffusion coefficient\n D = aux⁻.diffusion.D\n # exact the exact data\n ∇ρ = ∇state.ρ\n # set the flux\n if N == 1\n F.ρ = -D * ∇ρ\n else\n F.ρ = hcat(ntuple(n -> -D[:, :, n] * ∇ρ[:, n], Val(N))...)\n end\n elseif any_isa(bcs, HomogeneousBC{1}) # Zero Neumann\n F.ρ = zeros(typeof(F.ρ))\n end\n nothing\nend\n\nfunction boundary_state!(\n nf::CentralNumericalFluxDivergence,\n bcs,\n m::AdvectionDiffusion,\n grad⁺::Grad,\n aux⁺::Vars,\n n⁻::SVector,\n grad⁻::Grad,\n aux⁻::Vars,\n t,\n)\n if m.hyperdiffusion isa NoHyperDiffusion\n return nothing\n end\n\n if any_isa(bcs, InhomogeneousBC{1})\n # Get analytic gradient\n inhomogeneous_data!(Val(1), m.problem, grad⁺, aux⁻, aux⁻.coord, t)\n elseif any_isa(bcs, HomogeneousBC{1})\n grad⁺.ρ = zeros(typeof(grad⁺.ρ))\n end\n nothing\nend\n\nfunction boundary_state!(\n ::CentralNumericalFluxHigherOrder,\n bcs,\n m::AdvectionDiffusion{N},\n state⁺::Vars,\n aux⁺::Vars,\n lap⁺::Vars,\n n⁻::SVector,\n state⁻::Vars,\n aux⁻::Vars,\n lap⁻::Vars,\n t,\n) where {N}\n if m.hyperdiffusion isa NoHyperDiffusion\n return nothing\n end\n\n if any_isa(bcs, InhomogeneousBC{2})\n # Get analytic laplacian\n inhomogeneous_data!(Val(2), m.problem, lap⁺, aux⁻, aux⁻.coord, t)\n elseif any_isa(bcs, HomogeneousBC{2})\n lap⁺.ρ = N == 1 ? 0 : zeros(typeof(lap⁺.ρ))\n end\n nothing\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bgmres.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing Test\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SystemSolvers: BatchedGeneralizedMinimalResidual\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n if length(ARGS) > 0\n const integration_testing = parse(Bool, ARGS[1])\n else\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\n end\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # Diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (\n InhomogeneousBC{0}(),\n InhomogeneousBC{1}(),\n HomogeneousBC{0}(),\n HomogeneousBC{1}(),\n )\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = EveryDirection(),\n )\n\n vdg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n\n Q = init_ode_state(dg, FT(0))\n\n linearsolver = BatchedGeneralizedMinimalResidual(\n dg,\n Q;\n atol = sqrt(eps(FT)) * 0.01,\n rtol = sqrt(eps(FT)) * 0.01,\n )\n\n ode_solver = ARK548L2SA2KennedyCarpenter(\n dg,\n vdg,\n LinearBackwardEulerSolver(linearsolver; isadjustable = true),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(ode_solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(\n Q,\n ode_solver;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64] = 7.2801198255507391e-02\n expected_result[2, 2, Float64] = 6.8160295851506783e-03\n expected_result[2, 3, Float64] = 1.4439137164205592e-04\n expected_result[2, 4, Float64] = 2.4260727323386998e-06\n expected_result[3, 1, Float64] = 1.0462203776357534e-01\n expected_result[3, 2, Float64] = 1.0280535683502070e-02\n expected_result[3, 3, Float64] = 2.0631857053908848e-04\n expected_result[3, 4, Float64] = 3.3460492914169325e-06\n expected_result[2, 1, Float32] = 7.2801239788532257e-02\n expected_result[2, 2, Float32] = 6.8159680813550949e-03\n expected_result[2, 3, Float32] = 1.4439738879445940e-04\n # This is near roundoff so we will not check it\n # expected_result[2, 4, Float32] = 2.6432753656990826e-06\n expected_result[3, 1, Float32] = 1.0462204366922379e-01\n expected_result[3, 2, Float32] = 1.0280583053827286e-02\n expected_result[3, 3, Float32] = 2.0646647317335010e-04\n expected_result[3, 4, Float32] = 2.0226731066941284e-05\n\n numlevels = integration_testing ? 4 : 1\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n result = zeros(FT, numlevels)\n for dim in 2:3\n for fluxBC in (true, false)\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n connectivity = dim == 2 ? :face : :full\n linearsolvertype = \"Batched GMRES\"\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(FT(-1); length = Ne + 1, stop = 1),\n dim - 1,\n )...,\n range(FT(-5); length = 5Ne + 1, stop = 5),\n )\n\n periodicity = ntuple(j -> false, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = (\n ntuple(j -> (1, 2), dim - 1)...,\n (3, 4),\n ),\n connectivity = connectivity,\n )\n dt = (α / 4) / (Ne * polynomialorder^2)\n\n outputtime = 0.01\n timeend = 0.5\n\n @info (ArrayType, FT, dim, linearsolvertype, l, fluxBC)\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_$(linearsolvertype)_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n )\n # Test the errors significantly larger than floating point epsilon\n if !(dim == 2 && l == 4 && FT == Float32)\n @test result[l] ≈ FT(expected_result[dim, l, FT])\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/advection_diffusion_model_1dimex_bjfnks.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing Test\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n if length(ARGS) > 0\n const integration_testing = parse(Bool, ARGS[1])\n else\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\n end\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (\n InhomogeneousBC{0}(),\n InhomogeneousBC{1}(),\n HomogeneousBC{0}(),\n HomogeneousBC{1}(),\n )\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = EveryDirection(),\n )\n\n vdg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n linvdg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n # linearsolver = GeneralizedMinimalResidual(Q; M = 30, rtol = 1e-5)\n linearsolver =\n BatchedGeneralizedMinimalResidual(dg, Q; atol = -1.0, rtol = 1e-5)\n\n nonlinearsolver =\n JacobianFreeNewtonKrylovSolver(Q, linearsolver; tol = 1e-4)\n\n ode_solver = ARK548L2SA2KennedyCarpenter(\n dg,\n vdg,\n NonLinearBackwardEulerSolver(\n nonlinearsolver;\n isadjustable = true,\n preconditioner_update_freq = 1000,\n ),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(ode_solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(\n Q,\n ode_solver;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64] = 7.2801198255507391e-02\n expected_result[2, 2, Float64] = 6.8160295851506783e-03\n expected_result[2, 3, Float64] = 1.4439137164205592e-04\n expected_result[2, 4, Float64] = 2.4260727323386998e-06\n expected_result[3, 1, Float64] = 1.0462203776357534e-01\n expected_result[3, 2, Float64] = 1.0280535683502070e-02\n expected_result[3, 3, Float64] = 2.0631857053908848e-04\n expected_result[3, 4, Float64] = 3.3460492914169325e-06\n expected_result[2, 1, Float32] = 7.2801239788532257e-02\n expected_result[2, 2, Float32] = 6.8159680813550949e-03\n expected_result[2, 3, Float32] = 1.4439738879445940e-04\n # This is near roundoff so we will not check it\n # expected_result[2, 4, Float32] = 2.6432753656990826e-06\n expected_result[3, 1, Float32] = 1.0462204366922379e-01\n expected_result[3, 2, Float32] = 1.0280583053827286e-02\n expected_result[3, 3, Float32] = 2.0646647317335010e-04\n expected_result[3, 4, Float32] = 2.0226731066941284e-05\n\n numlevels = integration_testing ? 4 : 1\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n result = zeros(FT, numlevels)\n for dim in 2:3\n for fluxBC in (true, false)\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n\n linearsolvertype = \"Batched GMRES\"\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(FT(-1); length = Ne + 1, stop = 1),\n dim - 1,\n )...,\n range(FT(-5); length = 5Ne + 1, stop = 5),\n )\n\n periodicity = ntuple(j -> false, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = (\n ntuple(j -> (1, 2), dim - 1)...,\n (3, 4),\n ),\n )\n dt = (α / 4) / (Ne * polynomialorder^2)\n\n outputtime = 0.01\n timeend = 0.5\n\n @info (ArrayType, FT, dim, linearsolvertype, l, fluxBC)\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_$(linearsolvertype)_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n )\n # test the errors significantly larger than floating point epsilon\n if !(dim == 2 && l == 4 && FT == Float32)\n @test result[l] ≈ FT(expected_result[dim, l, FT])\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/advection_sphere.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.BalanceLaws: update_auxiliary_state!\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Atmos: SphericalOrientation, latitude, longitude\nusing ClimateMachine.Orientations\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\nimport ClimateMachine.BalanceLaws: boundary_state!\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\n# This is a setup similar to the one presented in [Williamson1992](@cite)\nstruct SolidBodyRotation <: AdvectionDiffusionProblem end\nfunction init_velocity_diffusion!(\n ::SolidBodyRotation,\n aux::Vars,\n geom::LocalGeometry,\n)\n FT = eltype(aux)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n r = norm(geom.coord)\n\n uλ = 2 * FT(π) * cos(φ) * r\n uφ = 0\n aux.advection.u = SVector(\n -uλ * sin(λ) - uφ * cos(λ) * sin(φ),\n +uλ * cos(λ) - uφ * sin(λ) * sin(φ),\n +uφ * cos(φ),\n )\nend\nfunction initial_condition!(::SolidBodyRotation, state, aux, localgeo, t)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n state.ρ = exp(-((3λ)^2 + (3φ)^2))\nend\nfinaltime(::SolidBodyRotation) = 1\nu_scale(::SolidBodyRotation) = 2π\n\n# This is a setup similar to the one presented in [Lauritzen2012](@cite)\nstruct ReversingDeformationalFlow <: AdvectionDiffusionProblem end\ninit_velocity_diffusion!(\n ::ReversingDeformationalFlow,\n aux::Vars,\n geom::LocalGeometry,\n) = nothing\nfunction initial_condition!(::ReversingDeformationalFlow, state, aux, coord, t)\n x, y, z = aux.coord\n r = norm(aux.coord)\n h_max = 0.95\n b = 5\n state.ρ = 0\n for (λ, φ) in ((5π / 6, 0), (7π / 6, 0))\n xi = r * cos(φ) * cos(λ)\n yi = r * cos(φ) * sin(λ)\n zi = r * sin(φ)\n state.ρ += h_max * exp(-b * ((x - xi)^2 + (y - yi)^2 + (z - zi)^2))\n end\nend\nhas_variable_coefficients(::ReversingDeformationalFlow) = true\nfunction update_velocity_diffusion!(\n ::ReversingDeformationalFlow,\n ::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n FT = eltype(aux)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n r = norm(aux.coord)\n T = FT(5)\n λp = λ - FT(2π) * t / T\n uλ =\n 10 * r / T * sin(λp)^2 * sin(2φ) * cos(FT(π) * t / T) +\n FT(2π) * r / T * cos(φ)\n uφ = 10 * r / T * sin(2λp) * cos(φ) * cos(FT(π) * t / T)\n aux.advection.u = SVector(\n -uλ * sin(λ) - uφ * cos(λ) * sin(φ),\n +uλ * cos(λ) - uφ * sin(λ) * sin(φ),\n +uφ * cos(φ),\n )\nend\nu_scale(::ReversingDeformationalFlow) = 2.9\nfinaltime(::ReversingDeformationalFlow) = 5\n\nfunction advective_courant(\n m::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n direction,\n)\n return Δt * norm(aux.advection.u) / Δx\nend\n\nstruct NoFlowBC end\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n ::NoFlowBC,\n ::AdvectionDiffusion,\n stateP::Vars,\n auxP::Vars,\n nM,\n stateM::Vars,\n auxM::Vars,\n t,\n _...,\n)\n auxP.advection.u = -auxM.advection.u\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n topl,\n problem,\n explicit_method,\n cfl,\n N,\n timeend,\n FT,\n vtkdir,\n outputtime,\n)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n dx = min_node_distance(grid, HorizontalDirection())\n dt = FT(cfl * dx / u_scale(problem()))\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n bcs = (NoFlowBC(),)\n model = AdvectionDiffusion{3}(problem(), bcs, diffusion = false)\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n odesolver = explicit_method(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n problem = %s\n method = %s\n time step = %.16e\n norm(Q₀) = %.16e\"\"\" problem explicit_method dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(odesolver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n cbcfl = EveryXSimulationSteps(10) do\n dt = ODESolvers.getdt(odesolver)\n cfl = DGMethods.courant(\n advective_courant,\n dg,\n model,\n Q,\n dt,\n HorizontalDirection(),\n )\n @info @sprintf(\n \"\"\"Courant number\n simtime = %.16e\n courant = %.16e\"\"\",\n gettime(odesolver),\n cfl\n )\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model, \"advection_sphere\")\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(odesolver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_sphere\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n Δmass = abs(weightedsum(Q) - weightedsum(Qe)) / weightedsum(Qe)\n @info @sprintf \"\"\"Finished\n Δmass = %.16e\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" Δmass engf engf / eng0 engf - eng0 errf errf / engfe\n return errf, Δmass\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 2\n\n max_cfl = Dict(\n LSRK144NiegemannDiehlBusch => 5.0,\n SSPRK33ShuOsher => 1.0,\n SSPRK34SpiteriRuuth => 1.5,\n )\n\n expected_result = Dict()\n\n expected_result[SolidBodyRotation, LSRK144NiegemannDiehlBusch, 1] =\n 1.3199024557832748e-01\n expected_result[SolidBodyRotation, LSRK144NiegemannDiehlBusch, 2] =\n 1.9868931633120656e-02\n expected_result[SolidBodyRotation, LSRK144NiegemannDiehlBusch, 3] =\n 1.4052110916915061e-03\n expected_result[SolidBodyRotation, LSRK144NiegemannDiehlBusch, 4] =\n 9.0193766298676310e-05\n\n expected_result[SolidBodyRotation, SSPRK33ShuOsher, 1] =\n 1.1055145388897809e-01\n expected_result[SolidBodyRotation, SSPRK33ShuOsher, 2] =\n 1.5510740467668628e-02\n expected_result[SolidBodyRotation, SSPRK33ShuOsher, 3] =\n 1.8629481690361454e-03\n expected_result[SolidBodyRotation, SSPRK33ShuOsher, 4] =\n 2.3567040048588889e-04\n\n expected_result[SolidBodyRotation, SSPRK34SpiteriRuuth, 1] =\n 1.2641959922001456e-01\n expected_result[SolidBodyRotation, SSPRK34SpiteriRuuth, 2] =\n 2.2780375948714751e-02\n expected_result[SolidBodyRotation, SSPRK34SpiteriRuuth, 3] =\n 3.1274951764826459e-03\n expected_result[SolidBodyRotation, SSPRK34SpiteriRuuth, 4] =\n 3.9734060514021565e-04\n\n expected_result[ReversingDeformationalFlow, LSRK144NiegemannDiehlBusch, 1] =\n 5.5387951598735408e-01\n expected_result[ReversingDeformationalFlow, LSRK144NiegemannDiehlBusch, 2] =\n 3.7610388138383732e-01\n expected_result[ReversingDeformationalFlow, LSRK144NiegemannDiehlBusch, 3] =\n 1.7823508719111605e-01\n expected_result[ReversingDeformationalFlow, LSRK144NiegemannDiehlBusch, 4] =\n 3.8639493470255713e-02\n\n expected_result[ReversingDeformationalFlow, SSPRK33ShuOsher, 1] =\n 5.5353962032596349e-01\n expected_result[ReversingDeformationalFlow, SSPRK33ShuOsher, 2] =\n 3.7645487928038762e-01\n expected_result[ReversingDeformationalFlow, SSPRK33ShuOsher, 3] =\n 1.7823263736245307e-01\n expected_result[ReversingDeformationalFlow, SSPRK33ShuOsher, 4] =\n 3.8605903366230925e-02\n\n expected_result[ReversingDeformationalFlow, SSPRK34SpiteriRuuth, 1] =\n 5.5404045660832824e-01\n expected_result[ReversingDeformationalFlow, SSPRK34SpiteriRuuth, 2] =\n 3.7788858038003154e-01\n expected_result[ReversingDeformationalFlow, SSPRK34SpiteriRuuth, 3] =\n 1.8007113931230376e-01\n expected_result[ReversingDeformationalFlow, SSPRK34SpiteriRuuth, 4] =\n 3.9941331660544775e-02\n\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ? 4 :\n 1\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,)\n for problem in (SolidBodyRotation, ReversingDeformationalFlow)\n for explicit_method in (\n LSRK144NiegemannDiehlBusch,\n SSPRK33ShuOsher,\n SSPRK34SpiteriRuuth,\n )\n cfl = max_cfl[explicit_method]\n result = zeros(FT, numlevels)\n for l in 1:numlevels\n numelems_horizontal = 2^(l - 1) * base_num_elem\n numelems_vertical = 1\n\n topl = StackedCubedSphereTopology(\n mpicomm,\n numelems_horizontal,\n range(\n FT(1),\n stop = 2,\n length = numelems_vertical + 1,\n ),\n )\n\n timeend = finaltime(problem())\n outputtime = timeend\n\n @info (ArrayType, FT)\n vtkdir =\n output ?\n \"vtk_advection_sphere\" *\n \"_$problem\" *\n \"_$explicit_method\" *\n \"_poly$(polynomialorder)\" *\n \"_$(ArrayType)_$(FT)\" *\n \"_level$(l)\" :\n nothing\n\n result[l], Δmass = test_run(\n mpicomm,\n ArrayType,\n topl,\n problem,\n explicit_method,\n cfl,\n polynomialorder,\n timeend,\n FT,\n vtkdir,\n outputtime,\n )\n @test result[l] ≈\n FT(expected_result[problem, explicit_method, l])\n @test Δmass <= FT(5e-14)\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/diffusion_hyperdiffusion_sphere.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.Mesh.Grids: min_node_distance\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct DiffusionSphere <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n problem::DiffusionSphere,\n aux::Vars,\n geom::LocalGeometry,\n)\n FT = eltype(aux)\n IM = SMatrix{3, 3, FT}(I)\n ZM = zeros(SMatrix{3, 3, FT})\n μ = FT(1 / 10000)\n aux.diffusion.D = μ * hcat(IM, ZM)\n aux.hyperdiffusion.H = μ * hcat(ZM, IM)\nend\n\nfunction initial_condition!(problem::DiffusionSphere, state, aux, localgeo, t)\n coord = localgeo.coord\n x, y, z = coord\n r = norm(coord)\n θ = atan(sqrt(x^2 + y^2), z)\n φ = atan(y, x)\n @inbounds begin\n # m = 1 l = 2 spherical harmonic\n ρ₀ = cos(φ) * sin(θ) * cos(θ)\n l = 2\n c = l * (l + 1) / r^2\n μ = aux.diffusion.D[1]\n state.ρ = (ρ₀ * exp(-c * μ * t), ρ₀ * exp(-c^2 * μ * t))\n end\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction run(mpicomm, ArrayType, topl, N, timeend, FT, vtkdir, outputtime)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n dx = min_node_distance(grid)\n dt = 300 * dx^4\n @info \"time step\" dt\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n model = AdvectionDiffusion{3}(\n DiffusionSphere(),\n advection = false,\n diffusion = true,\n hyperdiffusion = true,\n num_equations = 2,\n )\n dg = DGModel(\n model,\n grid,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n diffusion_direction = HorizontalDirection(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q, dims = (1, 3))\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0[1]\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model, \"diffusion_sphere\")\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"diffusion_sphere\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q, dims = (1, 3))\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe, dims = (1, 3))\n errf = norm(Q .- Qe, dims = (1, 3))\n\n metrics = @. (engf, engf / eng0, engf - eng0, errf, errf / engfe)\n\n @info @sprintf \"\"\"Finished\n Diffusion:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n HyperDiffusion:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" first.(metrics)... last.(metrics)...\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 3\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[Diffusion, 1] = 2.6002775334282785e-06\n expected_result[Diffusion, 2] = 4.1602462623838931e-07\n expected_result[Diffusion, 3] = 5.8997889725858304e-08\n\n expected_result[HyperDiffusion, 1] = 9.3032674316702424e-05\n expected_result[HyperDiffusion, 2] = 1.0149377619104328e-05\n expected_result[HyperDiffusion, 3] = 1.2985297333025857e-06\n\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ? 3 :\n 1\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,)\n result = Dict()\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n vert_range = grid1d(1, 2, nelem = 1)\n topl = StackedCubedSphereTopology(\n mpicomm,\n Ne,\n vert_range,\n boundary = (0, 0),\n )\n\n timeend = FT(1)\n outputtime = FT(2)\n\n @info (ArrayType, FT)\n vtkdir =\n output ?\n \"vtk_diffusion_sphere\" *\n \"_poly$(polynomialorder)\" *\n \"_$(ArrayType)_$(FT)\" *\n \"_level$(l)\" :\n nothing\n result[l] = run(\n mpicomm,\n ArrayType,\n topl,\n polynomialorder,\n timeend,\n FT,\n vtkdir,\n outputtime,\n )\n @test result[l][1] ≈ expected_result[Diffusion, l]\n @test result[l][2] ≈ expected_result[HyperDiffusion, l]\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = @. log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rates for level %d Diffusion = %e\",\n l,\n rate[1]\n )\n msg *= @sprintf(\", HyperDiffusion = %e\\n\", rate[2])\n end\n msg\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/direction_splitting_advection_diffusion.jl", "content": "using Test\nusing MPI\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing LinearAlgebra\nusing Printf\nusing Random\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Box{dim} end\nstruct Sphere end\n\nvertical_unit_vector(::Box{2}, ::SVector{3}) = SVector(0, 1, 0)\nvertical_unit_vector(::Box{3}, ::SVector{3}) = SVector(0, 0, 1)\nvertical_unit_vector(::Sphere, coord::SVector{3}) = coord / norm(coord)\n\nprojection(::EveryDirection, ::SVector{3}) = I\nprojection(::VerticalDirection, k::SVector{3}) = k * k'\nprojection(::HorizontalDirection, k::SVector{3}) = I - k * k'\n\nstruct TestProblem{adv, diff, dir, topo} <: AdvectionDiffusionProblem end\n\ninitial_ρ(::Box, x) = prod(sin.(π * x))\nfunction initial_ρ(::Sphere, x)\n r = norm(x)\n φ = atan(x[2], x[1])\n θ = atan(sqrt(x[1]^2 + x[2]^2), x[3])\n return sin(π * (r - 1)) * sin(φ) * sin(θ)\nend\n\nvelocity(::Box, x) = sin.(π * x)\nfunction velocity(::Sphere, x)\n r = norm(x)\n φ = atan(x[2], x[1])\n θ = atan(sqrt(x[1]^2 + x[2]^2), x[3])\n return sin(π * (r - 1)) .*\n SVector(cos(φ) * cos(θ), sin(φ) * cos(θ), cos(φ) * sin(θ))\nend\n\nfunction init_velocity_diffusion!(\n ::TestProblem{adv, diff, dir, topo},\n aux::Vars,\n geom::LocalGeometry,\n) where {adv, diff, dir, topo}\n k = vertical_unit_vector(topo, geom.coord)\n P = projection(dir, k)\n aux.advection.u =\n !isnothing(adv) ? P * velocity(topo, geom.coord) : zeros(SVector{3})\n aux.diffusion.D =\n !isnothing(diff) ? SMatrix{3, 3}(P) / 200 : zeros(SMatrix{3, 3})\nend\n\nfunction initial_condition!(\n ::TestProblem{adv, diff, dir, topo},\n state,\n aux,\n localgeo,\n t,\n) where {adv, diff, dir, topo}\n state.ρ = initial_ρ(topo, localgeo.coord)\nend\n\nfunction create_topology(::Box{dim}, mpicomm, Ne, FT) where {dim}\n brickrange = ntuple(j -> range(FT(0); length = Ne + 1, stop = 1), dim)\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 1), dim)\n connectivity = dim == 3 ? :full : :face\n StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\nend\n\nfunction create_topology(::Sphere, mpicomm, Ne, FT)\n vert_range = grid1d(FT(1), FT(2), nelem = Ne)\n StackedCubedSphereTopology(mpicomm, Ne, vert_range, boundary = (1, 1))\nend\n\ncreate_dg(model, grid, direction) = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = direction,\n)\n\nfunction test_run(\n adv,\n diff,\n topo,\n mpicomm,\n ArrayType,\n FT,\n polynomialorder,\n Ne,\n level,\n)\n topology = create_topology(topo(), mpicomm, Ne, FT)\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = topo == Sphere ? equiangular_cubed_sphere_warp :\n (x...) -> identity(x),\n )\n\n problems = (\n p = TestProblem{adv, diff, EveryDirection(), topo()}(),\n vp = TestProblem{adv, diff, VerticalDirection(), topo()}(),\n hp = TestProblem{adv, diff, HorizontalDirection(), topo()}(),\n )\n\n bcs = (HomogeneousBC{0}(),)\n models = map(p -> AdvectionDiffusion{3}(p, bcs), problems)\n\n dgmodels = map(models) do m\n (\n dg = create_dg(m, grid, EveryDirection()),\n vdg = create_dg(m, grid, VerticalDirection()),\n hdg = create_dg(m, grid, HorizontalDirection()),\n )\n end\n\n Q = init_ode_state(dgmodels.p.dg, FT(0), init_on_cpu = true)\n # do one Euler step to trigger numerical fluxes in subsequent evaluations\n let\n dt = 1e-3\n dQ = similar(Q)\n dgmodels.p.dg(dQ, Q, nothing, FT(0))\n Q .+= dt .* dQ\n end\n\n # evaluate all combinations\n dQ = map(\n x -> map(dg -> (dQ = similar(Q); dg(dQ, Q, nothing, FT(0)); dQ), x),\n dgmodels,\n )\n\n # set up tolerances\n atolm = 6e-13\n atolv = 3e-4 / 5^(level - 1)\n atolh = 0.0003 / 5^(level - 1)\n\n @testset \"total\" begin\n atol = topo <: Box || diff == nothing ? atolm : atolh\n @test isapprox(norm(dQ.p.dg .- dQ.p.vdg .- dQ.p.hdg), 0, atol = atol)\n @test isapprox(norm(dQ.vp.dg .- dQ.vp.vdg .- dQ.vp.hdg), 0, atol = atol)\n @test isapprox(norm(dQ.hp.dg .- dQ.hp.vdg .- dQ.hp.hdg), 0, atol = atol)\n end\n\n @testset \"vertical\" begin\n atol = topo <: Box ? atolm : atolv\n @test isapprox(norm(dQ.vp.dg .- dQ.vp.vdg), 0, atol = atol)\n @test isapprox(norm(dQ.vp.hdg), 0, atol = atol)\n @test isapprox(norm(dQ.vp.dg .- dQ.p.vdg), 0, atol = atol)\n end\n\n @testset \"horizontal\" begin\n atol = topo <: Box ? atolm : atolh\n @test isapprox(norm(dQ.hp.dg .- dQ.hp.hdg), 0, atol = atol)\n @test isapprox(norm(dQ.hp.vdg), 0, atol = atol)\n @test isapprox(norm(dQ.hp.dg - dQ.p.hdg), 0, atol = atol)\n end\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n FT = Float64\n numlevels = 2\n base_num_elem = 4\n\n # This test doesn't do any heavy computational work,\n # but compiles a lot of model/discretization combinations.\n # Compilation times on the GPU are longer, so we only run one\n # variable-degree case on the GPU\n if ArrayType == Array\n polynomialorders = ((4, 4), (4, 2))\n else\n polynomialorders = ((4, 2),)\n end\n\n @info @sprintf \"\"\"Test parameters:\n ArrayType = %s\n FloatType = %s\n Polynomial orders = %s\n \"\"\" ArrayType FT polynomialorders\n\n @testset \"$(@__FILE__)\" begin\n @testset for polyorders in polynomialorders\n @testset for topo in (Box{2}, Box{3}, Sphere)\n @testset for (adv, diff) in (\n (Advection, nothing),\n (nothing, Diffusion),\n (Advection, Diffusion),\n )\n @testset for level in 1:numlevels\n Ne = 2^(level - 1) * base_num_elem\n test_run(\n adv,\n diff,\n topo,\n mpicomm,\n ArrayType,\n FT,\n polyorders,\n Ne,\n level,\n )\n end\n end\n end\n end\n end\nend\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_advection.jl", "content": "import Printf: @sprintf\nimport LinearAlgebra: dot, norm\nimport Dates\nimport MPI\n\nimport ClimateMachine\nimport ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nimport ClimateMachine.DGMethods.NumericalFluxes:\n RusanovNumericalFlux,\n CentralNumericalFluxSecondOrder,\n CentralNumericalFluxGradient\nimport ClimateMachine.DGMethods: DGFVModel, init_ode_state\nimport ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nimport ClimateMachine.MPIStateArrays: MPIStateArray, euclidean_distance\nimport ClimateMachine.Mesh.Grids:\n DiscontinuousSpectralElementGrid, EveryDirection\nimport ClimateMachine.Mesh.Topologies: StackedBrickTopology\nimport ClimateMachine.ODESolvers: LSRK54CarpenterKennedy, solve!, gettime\nimport ClimateMachine.VTK: writevtk, writepvtu\n\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α},\n state,\n aux,\n localgeo,\n t,\n) where {n, α}\n ξn = dot(n, localgeo.coord)\n state.ρ = sin((ξn - α * t) * pi)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n vtkdir,\n outputtime,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (InhomogeneousBC{0}(),)\n model = AdvectionDiffusion{dim}(Pseudo1D{n, α}(), bcs, diffusion = false)\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = EveryDirection(),\n )\n\n Q = init_ode_state(dgfvm, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dgfvm, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(Dates.now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = Dates.now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(lsrk))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dgfvm, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64, FVConstant()] = 1.0404261715459338e-01\n expected_result[2, 2, Float64, FVConstant()] = 5.5995868545685376e-02\n expected_result[2, 3, Float64, FVConstant()] = 2.9383695610072275e-02\n expected_result[2, 4, Float64, FVConstant()] = 1.5171779426843507e-02\n\n expected_result[2, 1, Float64, FVLinear()] = 8.3196657944903635e-02\n expected_result[2, 2, Float64, FVLinear()] = 3.9277132273521774e-02\n expected_result[2, 3, Float64, FVLinear()] = 1.7155773433218020e-02\n expected_result[2, 4, Float64, FVLinear()] = 7.6525022056819231e-03\n\n expected_result[3, 1, Float64, FVConstant()] = 9.6785620234054362e-02\n expected_result[3, 2, Float64, FVConstant()] = 5.3406412788842651e-02\n expected_result[3, 3, Float64, FVConstant()] = 2.8471535157235807e-02\n expected_result[3, 4, Float64, FVConstant()] = 1.4846239937398318e-02\n\n expected_result[3, 1, Float64, FVLinear()] = 8.5860120005258181e-02\n expected_result[3, 2, Float64, FVLinear()] = 4.2844889694123235e-02\n expected_result[3, 3, Float64, FVLinear()] = 1.9302295207100174e-02\n expected_result[3, 4, Float64, FVLinear()] = 8.6084633401356733e-03\n\n expected_result[2, 1, Float32, FVConstant()] = 1.0404255986213684e-01\n expected_result[2, 2, Float32, FVConstant()] = 5.5995877832174301e-02\n expected_result[2, 3, Float32, FVConstant()] = 2.9383875429630280e-02\n expected_result[2, 4, Float32, FVConstant()] = 1.5171864069998264e-02\n\n expected_result[2, 1, Float32, FVLinear()] = 8.3196602761745453e-02\n expected_result[2, 2, Float32, FVLinear()] = 3.9277125149965286e-02\n expected_result[2, 3, Float32, FVLinear()] = 1.7155680805444717e-02\n expected_result[2, 4, Float32, FVLinear()] = 7.6521718874573708e-03\n\n expected_result[3, 1, Float32, FVConstant()] = 9.6785508096218109e-02\n expected_result[3, 2, Float32, FVConstant()] = 5.3406376391649246e-02\n expected_result[3, 3, Float32, FVConstant()] = 2.8471505269408226e-02\n expected_result[3, 4, Float32, FVConstant()] = 1.4849635772407055e-02\n\n expected_result[3, 1, Float32, FVLinear()] = 8.5860058665275574e-02\n expected_result[3, 2, Float32, FVLinear()] = 4.2844854295253754e-02\n expected_result[3, 3, Float32, FVLinear()] = 1.9302234053611755e-02\n expected_result[3, 4, Float32, FVLinear()] = 8.6139924824237823e-03\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n numlevels =\n integration_testing ||\n ClimateMachine.Settings.integration_testing ?\n 4 : 1\n result = zeros(FT, numlevels)\n for dim in 2:3\n N = (ntuple(j -> 4, dim - 1)..., 0)\n n =\n dim == 2 ? SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0) :\n SVector{3, FT}(1 / sqrt(3), 1 / sqrt(3), 1 / sqrt(3))\n α = FT(1)\n connectivity = dim == 2 ? :face : :full\n\n for fvmethod in (FVConstant(), FVLinear())\n @info @sprintf \"\"\"Configuration\n FT = %s\n ArrayType = %s\n FV Reconstruction = %s\n dims = %d\n \"\"\" FT ArrayType fvmethod dim\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(FT(-1); length = Ne + 1, stop = 1),\n dim - 1,\n )...,\n range(FT(-1); length = N[1] * Ne + 1, stop = 1),\n )\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 1), dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n dt = (α / 4) / (Ne * max(1, maximum(N))^2)\n\n timeend = FT(1 // 4)\n outputtime = timeend\n\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(N)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n vtkdir,\n outputtime,\n )\n @test result[l] ≈\n FT(expected_result[dim, l, FT, fvmethod])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Test\nimport ClimateMachine.VTK: writevtk, writepvtu\nimport ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing Dates\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{ns, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{ns, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {ns, α, β}\n # Direction of flow is ns[i] with magnitude α\n aux.advection.u = hcat(ntuple(i -> α * ns[i], Val(length(ns)))...)\n\n # Diffusion of strength β in the ns[i] direction\n aux.diffusion.D = hcat(ntuple(i -> β * ns[i] * ns[i]', Val(length(ns)))...)\nend\n\nfunction gaussian(x, t, α, β, μ, δ)\n exp(-(x - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\nfunction ∇gaussian(n, x, t, α, β, μ, δ)\n -2n * (x - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(x - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\nfunction initial_condition!(\n ::Pseudo1D{ns, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {ns, α, β, μ, δ}\n ρ = ntuple(Val(length(ns))) do i\n ξn = dot(ns[i], localgeo.coord)\n gaussian(ξn, t, α, β, μ, δ)\n end\n state.ρ = length(ns) == 1 ? ρ[1] : ρ\nend\n\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\n\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{ns, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {ns, α, β, μ, δ}\n ∇state.ρ = hcat(ntuple(Val(length(ns))) do i\n ξn = dot(ns[i], x)\n ∇gaussian(ns[i], ξn, t, α, β, μ, δ)\n end...)\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n dim,\n fvmethod,\n polynomialorders,\n level,\n ArrayType,\n FT,\n vtkdir,\n direction,\n)\n\n n_hd =\n dim == 2 ? SVector{3, FT}(1, 0, 0) :\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0)\n\n n_vd = dim == 2 ? SVector{3, FT}(0, 1, 0) : SVector{3, FT}(0, 0, 1)\n\n n_dg =\n dim == 2 ? SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0) :\n SVector{3, FT}(1 / sqrt(3), 1 / sqrt(3), 1 / sqrt(3))\n connectivity = dim == 2 ? :face : :full\n if direction isa EveryDirection\n ns = (n_hd, n_vd, n_dg)\n elseif direction isa HorizontalDirection\n ns = (n_hd,)\n elseif direction isa VerticalDirection\n ns = (n_vd,)\n end\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n\n # Grid/topology information\n base_num_elem = 4\n Ne = 2^(level - 1) * base_num_elem\n N = polynomialorders\n L = ntuple(j -> FT(j == dim ? 1 : N[1]) / 4, dim)\n brickrange = ntuple(j -> range(-L[j]; length = Ne + 1, stop = L[j]), dim)\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 2), dim)\n\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n\n dt = (α / 4) * L[1] / (Ne * polynomialorders[1]^2)\n timeend = 1\n outputtime = timeend / 10\n @info \"time step\" dt\n\n @info @sprintf \"\"\"Test parameters:\n FVM Reconstruction = %s\n ArrayType = %s\n FloatType = %s\n Dimension = %s\n Direction = %s\n Horizontal polynomial order = %s\n Vertical polynomial order = %s\n \"\"\" fvmethod ArrayType FT dim direction polynomialorders[1] polynomialorders[end]\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorders,\n )\n\n bcs = (InhomogeneousBC{0}(), InhomogeneousBC{1}())\n # Model being tested\n model = AdvectionDiffusion{dim}(\n Pseudo1D{ns, α, β, μ, δ}(),\n bcs,\n num_equations = length(ns),\n )\n\n # Main DG discretization\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = direction,\n )\n\n # Initialize all relevant state arrays and create solvers\n Q = init_ode_state(dgfvm, FT(0))\n\n eng0 = norm(Q, dims = (1, 3))\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0[1]\n\n solver = LSRK54CarpenterKennedy(dgfvm, Q; dt = dt, t0 = 0)\n\n # Set up the information callback\n starttime = Ref(Dates.now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = Dates.now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, solver; timeend = timeend, callbacks = callbacks)\n\n # Reference solution\n engf = norm(Q, dims = (1, 3))\n Q_ref = init_ode_state(dgfvm, FT(timeend))\n\n engfe = norm(Q_ref, dims = (1, 3))\n errf = norm(Q_ref .- Q, dims = (1, 3))\n\n metrics = @. (engf, engf / eng0, engf - eng0, errf, errf / engfe)\n\n @info begin\n j = 1\n msg = \"Finished\\n\"\n if direction isa HorizontalDirection || direction isa EveryDirection\n msg *= @sprintf \"\"\"\n Horizontal field:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" ntuple(f -> metrics[f][j], 5)...\n j += 1\n end\n\n if direction isa VerticalDirection || direction isa EveryDirection\n msg *= @sprintf \"\"\"\n Vertical field:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" ntuple(f -> metrics[f][j], 5)...\n j += 1\n end\n\n if direction isa EveryDirection\n msg *= @sprintf \"\"\"\n Diagonal field:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" ntuple(f -> metrics[f][j], 5)...\n end\n msg\n end\n\n return Tuple(errf)\nend\n\n\"\"\"\n main()\n\nRun this test problem\n\"\"\"\nfunction main()\n\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n # Dictionary keys: dim, level, and FT\n expected_result = Dict()\n expected_result[2, 1, Float32, FVConstant()] =\n (2.3391991853713989e-02, 5.0738707184791565e-02, 4.1857458651065826e-02)\n expected_result[2, 2, Float32, FVConstant()] =\n (2.0331495907157660e-03, 3.3669386059045792e-02, 2.3904174566268921e-02)\n expected_result[2, 3, Float32, FVConstant()] =\n (2.6557327146292664e-05, 2.0002065226435661e-02, 1.2790882028639317e-02)\n\n expected_result[2, 1, Float32, FVLinear()] =\n (2.3391995579004288e-02, 4.9536284059286118e-02, 3.5911198705434799e-02)\n expected_result[2, 2, Float32, FVLinear()] =\n (2.0331430714577436e-03, 2.2621221840381622e-02, 1.5531544573605061e-02)\n expected_result[2, 3, Float32, FVLinear()] =\n (2.6568108296487480e-05, 7.5304862111806870e-03, 6.4526572823524475e-03)\n\n expected_result[3, 1, Float32, FVConstant()] =\n (8.7377587333321571e-03, 7.1755334734916687e-02, 4.3733172118663788e-02)\n expected_result[3, 2, Float32, FVConstant()] =\n (6.2517996411770582e-04, 4.7615684568881989e-02, 2.3986011743545532e-02)\n expected_result[3, 3, Float32, FVConstant()] =\n (3.5004395613214001e-05, 2.8287241235375404e-02, 1.2639496475458145e-02)\n\n expected_result[3, 1, Float32, FVLinear()] =\n (8.7377615272998810e-03, 7.0054858922958374e-02, 3.7571556866168976e-02)\n expected_result[3, 2, Float32, FVLinear()] =\n (6.2518107006326318e-04, 3.1991228461265564e-02, 1.6405479982495308e-02)\n expected_result[3, 3, Float32, FVLinear()] =\n (3.5004966775886714e-05, 1.0649790056049824e-02, 7.1499347686767578e-03)\n\n expected_result[2, 1, Float64, FVConstant()] =\n (2.3391871809628567e-02, 5.0738761087541523e-02, 4.1857480220018319e-02)\n expected_result[2, 2, Float64, FVConstant()] =\n (2.0332783913617892e-03, 3.3669399006499491e-02, 2.3904204301839819e-02)\n expected_result[2, 3, Float64, FVConstant()] =\n (2.6572168347086839e-05, 2.0002110124502395e-02, 1.2790993871268752e-02)\n expected_result[2, 4, Float64, FVConstant()] =\n (1.9890000550154039e-07, 1.0929144069594809e-02, 6.5110938897763263e-03)\n\n expected_result[2, 1, Float64, FVLinear()] =\n (2.3391871809628550e-02, 4.9536356061135857e-02, 3.5911213637569099e-02)\n expected_result[2, 2, Float64, FVLinear()] =\n (2.0332783913618278e-03, 2.2621252826152356e-02, 1.5531565558700888e-02)\n expected_result[2, 3, Float64, FVLinear()] =\n (2.6572168347111800e-05, 7.5305291439555161e-03, 6.4526740563561544e-03)\n expected_result[2, 4, Float64, FVLinear()] =\n (1.9890000549995218e-07, 2.5302226025389427e-03, 2.7792905025059711e-03)\n\n expected_result[3, 1, Float64, FVConstant()] =\n (8.7378337431297422e-03, 7.1755444068009461e-02, 4.3733196512722658e-02)\n expected_result[3, 2, Float64, FVConstant()] =\n (6.2510740807095622e-04, 4.7615720711942776e-02, 2.3986017606198881e-02)\n expected_result[3, 3, Float64, FVConstant()] =\n (3.4995405318038341e-05, 2.8287255414151377e-02, 1.2639742577376042e-02)\n expected_result[3, 4, Float64, FVConstant()] =\n (1.4362091045094841e-06, 1.5456143768350493e-02, 6.3677406803847331e-03)\n\n expected_result[3, 1, Float64, FVLinear()] =\n (8.7378337431297439e-03, 7.0054986572200995e-02, 3.7571599264936015e-02)\n expected_result[3, 2, Float64, FVLinear()] =\n (6.2510740807095286e-04, 3.1991282544615307e-02, 1.6405489038457288e-02)\n expected_result[3, 3, Float64, FVLinear()] =\n (3.4995405318035868e-05, 1.0649776447227678e-02, 7.1502921502446283e-03)\n expected_result[3, 4, Float64, FVLinear()] =\n (1.4362091045110612e-06, 3.5782751203335653e-03, 3.1186151510318319e-03)\n\n @testset \"Variable degree DG: advection diffusion model\" begin\n for FT in (Float32, Float64)\n for direction in\n (EveryDirection(), HorizontalDirection(), VerticalDirection())\n numlevels =\n integration_testing ||\n ClimateMachine.Settings.integration_testing ?\n (FT == Float64 ? 4 : 3) : 1\n for dim in 2:3\n for fvmethod in (FVConstant(), FVLinear(), FVLinear{3}())\n polynomialorders = (4, 0)\n result = Dict()\n for level in 1:numlevels\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorders)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_fvmethod$(fvmethod)\" *\n \"_level$(level)\" :\n nothing\n result[level] = test_run(\n mpicomm,\n dim,\n fvmethod,\n polynomialorders,\n level,\n ArrayType,\n FT,\n vtkdir,\n direction,\n )\n fv_key =\n fvmethod isa FVLinear ? FVLinear() : fvmethod\n if direction isa EveryDirection\n @test all(\n result[level] .≈\n FT.(expected_result[\n dim,\n level,\n FT,\n fv_key,\n ]),\n )\n elseif direction isa HorizontalDirection\n @test result[level][1] .≈ FT.(expected_result[\n dim,\n level,\n FT,\n fv_key,\n ])[1]\n elseif direction isa VerticalDirection\n @test result[level][1] .≈ FT.(expected_result[\n dim,\n level,\n FT,\n fv_key,\n ])[2]\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = @. log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\"\\n rates for level %d\", l)\n j = 1\n if direction isa HorizontalDirection ||\n direction isa EveryDirection\n msg *=\n @sprintf(\", Horizontal = %e\", rate[j])\n j += 1\n end\n if direction isa VerticalDirection ||\n direction isa EveryDirection\n msg *= @sprintf(\", Vertical = %e\", rate[j])\n j += 1\n end\n if direction isa EveryDirection\n msg *= @sprintf(\", Diagonal = %e\", rate[j])\n end\n msg *= \"\\n\"\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion_model_1dimex_bjfnks.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing Test\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n if length(ARGS) > 0\n const integration_testing = parse(Bool, ARGS[1])\n else\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\n end\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # Diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n fvmethod,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = (N, 0),\n )\n\n bcs = (\n InhomogeneousBC{0}(),\n InhomogeneousBC{1}(),\n HomogeneousBC{0}(),\n HomogeneousBC{1}(),\n )\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = EveryDirection(),\n )\n\n vdgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dgfvm.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n\n Q = init_ode_state(dgfvm, FT(0))\n\n linearsolver = BatchedGeneralizedMinimalResidual(\n dgfvm,\n Q;\n max_subspace_size = 5,\n atol = sqrt(eps(FT)) * 0.01,\n rtol = sqrt(eps(FT)) * 0.01,\n )\n\n nonlinearsolver =\n JacobianFreeNewtonKrylovSolver(Q, linearsolver; tol = 1e-4)\n\n ode_solver = ARK2ImplicitExplicitMidpoint(\n dgfvm,\n vdgfvm,\n NonLinearBackwardEulerSolver(\n nonlinearsolver;\n isadjustable = true,\n preconditioner_update_freq = 1000,\n ),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(ode_solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(ode_solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(\n Q,\n ode_solver;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dgfvm, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n # dim, refinement level, FT, vertical scheme\n expected_result[2, 1, Float64, FVConstant()] = 1.4890228213182394e-01\n expected_result[2, 2, Float64, FVConstant()] = 1.1396608915201276e-01\n expected_result[2, 3, Float64, FVConstant()] = 7.9360293305962212e-02\n\n expected_result[2, 1, Float64, FVLinear()] = 1.0755278583097065e-01\n expected_result[2, 2, Float64, FVLinear()] = 4.5924219102504410e-02\n expected_result[2, 3, Float64, FVLinear()] = 1.0775256661783415e-02\n\n expected_result[3, 1, Float64, FVConstant()] = 2.1167998484526718e-01\n expected_result[3, 2, Float64, FVConstant()] = 1.5936623436529485e-01\n expected_result[3, 3, Float64, FVConstant()] = 1.1069147173311458e-01\n\n expected_result[3, 1, Float64, FVLinear()] = 1.5361757743888277e-01\n expected_result[3, 2, Float64, FVLinear()] = 6.3622331921472902e-02\n expected_result[3, 3, Float64, FVLinear()] = 1.4836612761110498e-02\n\n\n numlevels = integration_testing ? 3 : 1\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,)\n result = zeros(FT, numlevels)\n for dim in 2:3\n connectivity = dim == 3 ? :full : :face\n for fvmethod in (FVConstant(), FVLinear())\n\n\n for fluxBC in (true, false)\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n\n solvertype = \"HEVI_Nolinearsolver\"\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(\n FT(-1);\n length = Ne + 1,\n stop = 1,\n ),\n dim - 1,\n )...,\n range(\n FT(-5);\n length = 5Ne * polynomialorder + 1,\n stop = 5,\n ),\n )\n\n periodicity = ntuple(j -> false, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = (\n ntuple(j -> (1, 2), dim - 1)...,\n (3, 4),\n ),\n connectivity = connectivity,\n )\n dt = 2 * (α / 4) / (Ne * polynomialorder^2)\n\n outputtime = 0.01\n timeend = 0.5\n\n @info (\n ArrayType,\n FT,\n dim,\n solvertype,\n l,\n polynomialorder,\n fvmethod,\n fluxBC,\n )\n vtkdir =\n output ?\n \"vtk_fvm_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_fvmethod$(fvmethod)\" *\n \"_$(solvertype)_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n fvmethod,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n fluxBC,\n )\n # Test the errors significantly larger than floating point epsilon\n if !(dim == 2 && l == 4 && FT == Float32)\n @test result[l] ≈\n FT(expected_result[dim, l, FT, fvmethod])\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_advection_diffusion_periodic.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nimport ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Test\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{u, v, ν} <: AdvectionDiffusionProblem end\n\n# This test has two 2D equations : the first one is an advection equation with\n# the Gaussian initial condition the second one is an advection diffusion\n# equation with the Gaussian initial condition\nfunction init_velocity_diffusion!(\n ::Pseudo1D{u, v, ν},\n aux::Vars,\n geom::LocalGeometry,\n) where {u, v, ν}\n # Advection velocity of the flow is [u, v]\n uvec = SVector(u, v, 0)\n aux.advection.u = hcat(uvec, uvec)\n\n # Diffusion of the flow is νI (isentropic diffusivity)\n I3 = @SMatrix [1 0 0; 0 1 0; 0 0 1]\n aux.diffusion.D = hcat(0 * I3, ν * I3)\n\nend\n\nfunction initial_condition!(\n ::Pseudo1D{u, v, ν},\n state,\n aux,\n localgeo,\n t,\n) where {u, v, ν}\n FT = typeof(u)\n # The computational domain is [-1.5 1.5]×[-1.5 1.5]\n Lx, Ly = 3, 3\n x, y, _ = localgeo.coord\n\n μx, μz, σ = 0, 0, Lx / 10\n ρ1 = exp.(-(((x - μx) / σ)^2 + ((y - μz) / σ)^2) / 2) / (σ * sqrt(2 * pi))\n ρ2 = exp.(-(((x - μx) / σ)^2 + ((y - μz) / σ)^2) / 2) / (σ * sqrt(2 * pi))\n\n state.ρ = (ρ1, ρ2)\nend\n\ninhomogenous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\n\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n vtkdir,\n fvmethod,\n polynomialorders,\n level,\n ArrayType,\n FT,\n)\n\n dim = 2\n Lx, Ly = FT(3), FT(3)\n u, v = FT(1), FT(1)\n ν = FT(1 // 100)\n\n\n # Grid/topology information\n base_num_elem = 4\n Ne = 2^(level - 1) * base_num_elem\n\n # Match number of points\n N_dg_point, N_fvm_point = Ne + 1, Ne * polynomialorders[1] + 1\n\n\n brickrange = (\n range(-Lx / 2; length = N_dg_point, stop = Lx / 2),\n range(-Ly / 2; length = N_fvm_point, stop = Ly / 2),\n )\n\n periodicity = ntuple(j -> true, dim)\n bc = ntuple(j -> (1, 2), dim)\n connectivity = dim == 3 ? :full : :face\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n\n # One period\n timeend = Lx / u\n # dt ≤ CFL (Δx / Np²)/u u = √2 CFL = 1/√2\n dt = Lx / (Ne * polynomialorders[1]^2)\n\n Nt = (Ne * polynomialorders[1]^2)\n\n outputtime = Ne * dt\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorders,\n )\n\n # Model being tested\n model = AdvectionDiffusion{dim}(Pseudo1D{u, v, ν}(), num_equations = 2)\n\n # Main DG discretization\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = EveryDirection(),\n )\n\n # Initialize all relevant state arrays and create solvers\n Q = init_ode_state(dgfvm, FT(0))\n\n solver = LSRK54CarpenterKennedy(dgfvm, Q; dt = dt, t0 = 0)\n\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Q,\n model,\n \"advection_diffusion_periodic\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n \"advection_diffusion_periodic\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(Q, solver; timeend = timeend, callbacks = callbacks)\n\n engf = norm(Q, dims = (1, 3))\n\n # Reference solution\n Q_ref = init_ode_state(dgfvm, FT(0))\n engfe = norm(Q_ref, dims = (1, 3))\n\n errf = norm(Q_ref .- Q, dims = (1, 3))\n\n metrics = [engf; engfe; errf]\n\n @info @sprintf \"\"\"Finished\n Advection equation:\n norm(Q) = %.16e\n norm(Qe) = %.16e\n norm(Q - Qe) = %.16e\n Advection diffusion equation:\n norm(Q) = %.16e\n norm(Qe) = %.16e\n norm(Q - Qe) = %.16e\n \"\"\" metrics[:]...\n\n return errf\nend\n\n# Run this test problem\nfunction main()\n\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n # Dictionary keys: dim, level, polynomial order, FT, and direction\n expected_result = Dict()\n\n # Dim 2, degree 4 in the horizontal, FV order 1, refinement level, Float64, equation number\n expected_result[2, 4, FVConstant(), 1, Float64, 1] = 4.5196354911392578e-01\n expected_result[2, 4, FVConstant(), 1, Float64, 2] = 4.8622171647472179e-01\n expected_result[2, 4, FVConstant(), 2, Float64, 1] = 3.5051983693457450e-01\n expected_result[2, 4, FVConstant(), 2, Float64, 2] = 4.1168975732563706e-01\n expected_result[2, 4, FVConstant(), 3, Float64, 1] = 2.5130141068332995e-01\n expected_result[2, 4, FVConstant(), 3, Float64, 2] = 3.4635415661628755e-01\n expected_result[2, 4, FVConstant(), 4, Float64, 1] = 1.6320856055402777e-01\n expected_result[2, 4, FVConstant(), 4, Float64, 2] = 2.9647989774687805e-01\n expected_result[2, 4, FVConstant(), 5, Float64, 1] = 9.6641259241910610e-02\n expected_result[2, 4, FVConstant(), 5, Float64, 2] = 2.6372336617960646e-01\n\n expected_result[2, 4, FVConstant(), 1, Float32, 1] = 4.5196333527565002e-01\n expected_result[2, 4, FVConstant(), 1, Float32, 2] = 4.8622152209281921e-01\n expected_result[2, 4, FVConstant(), 2, Float32, 1] = 3.5051953792572021e-01\n expected_result[2, 4, FVConstant(), 2, Float32, 2] = 4.1168937087059021e-01\n expected_result[2, 4, FVConstant(), 3, Float32, 1] = 2.5130018591880798e-01\n expected_result[2, 4, FVConstant(), 3, Float32, 2] = 3.4635293483734131e-01\n expected_result[2, 4, FVConstant(), 4, Float32, 1] = 1.6320574283599854e-01\n expected_result[2, 4, FVConstant(), 4, Float32, 2] = 2.9647585749626160e-01\n expected_result[2, 4, FVConstant(), 5, Float32, 1] = 9.6632070839405060e-02\n expected_result[2, 4, FVConstant(), 5, Float32, 2] = 2.6370745897293091e-01\n\n\n # Dim 2, degree 4 in the horizontal, FV order 1, refinement level, Float64, equation number\n\n expected_result[2, 4, FVLinear(), 1, Float64, 1] = 2.2783152269907422e-01\n expected_result[2, 4, FVLinear(), 1, Float64, 2] = 3.1823753821941969e-01\n expected_result[2, 4, FVLinear(), 2, Float64, 1] = 9.2628269590376469e-02\n expected_result[2, 4, FVLinear(), 2, Float64, 2] = 2.4517823755458742e-01\n expected_result[2, 4, FVLinear(), 3, Float64, 1] = 3.1401247771425542e-02\n expected_result[2, 4, FVLinear(), 3, Float64, 2] = 2.2608977452273774e-01\n expected_result[2, 4, FVLinear(), 4, Float64, 1] = 9.8710350648653390e-03\n expected_result[2, 4, FVLinear(), 4, Float64, 2] = 2.2377315397364847e-01\n expected_result[2, 4, FVLinear(), 5, Float64, 1] = 2.9619222883137744e-03\n expected_result[2, 4, FVLinear(), 5, Float64, 2] = 2.2359698753564772e-01\n\n\n expected_result[2, 4, FVLinear(), 1, Float32, 1] = 2.2783152269907422e-01\n expected_result[2, 4, FVLinear(), 1, Float32, 2] = 3.1823753821941969e-01\n expected_result[2, 4, FVLinear(), 2, Float32, 1] = 9.2628269590376469e-02\n expected_result[2, 4, FVLinear(), 2, Float32, 2] = 2.4517823755458742e-01\n expected_result[2, 4, FVLinear(), 3, Float32, 1] = 3.1401247771425542e-02\n expected_result[2, 4, FVLinear(), 3, Float32, 2] = 2.2608977452273774e-01\n expected_result[2, 4, FVLinear(), 4, Float32, 1] = 9.8710350648653390e-03\n expected_result[2, 4, FVLinear(), 4, Float32, 2] = 2.2377315397364847e-01\n expected_result[2, 4, FVLinear(), 5, Float32, 1] = 2.9614968225359917e-03\n expected_result[2, 4, FVLinear(), 5, Float32, 2] = 2.2358775138854980e-01\n\n polynomialorders = (4, 0)\n numlevels = integration_testing ? 5 : 1\n\n # Dictionary keys: dim, level, polynomial order, FT, and direction\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n result = Dict()\n for fvmethod in (FVConstant(), FVLinear(), FVLinear{3}())\n @info @sprintf \"\"\"Test parameters:\n ArrayType = %s\n FloatType = %s\n FV Reconstruction = %s\n Dimension = %s\n Horizontal polynomial order = %s\n Vertical polynomial order = %s\n \"\"\" ArrayType FT fvmethod 2 polynomialorders[1] polynomialorders[end]\n for level in 1:numlevels\n result[level] = test_run(\n mpicomm,\n output ?\n \"vtk_advection_diffusion_2d\" *\n \"_poly$(polynomialorders)\" *\n \"_$(ArrayType)_$(FT)\" *\n \"_level$(level)\" :\n nothing,\n fvmethod,\n polynomialorders,\n level,\n ArrayType,\n FT,\n )\n\n\n fv_key = fvmethod isa FVLinear ? FVLinear() : fvmethod\n @test result[level][1] ≈ FT(expected_result[\n 2,\n polynomialorders[1],\n fv_key,\n level,\n FT,\n 1,\n ])\n @test result[level][2] ≈ FT(expected_result[\n 2,\n polynomialorders[1],\n fv_key,\n level,\n FT,\n 2,\n ])\n end\n\n @info begin\n msg = \"advection equation\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l][1]) - log2(result[l + 1][1])\n msg *= @sprintf(\"\\n rate for level %d = %e\", l, rate)\n end\n msg\n end\n # Not an analytic solution, convergence doesn't make sense\n # @info begin\n # msg = \"advection-diffusion equation\"\n # for l in 1:(numlevels - 1)\n # rate = log2(result[l][2]) - log2(result[l + 1][2])\n # msg *= @sprintf(\"\\n rate for level %d = %e\", l, rate)\n # end\n # msg\n # end\n end\n end\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_advection_sphere.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.BalanceLaws: update_auxiliary_state!\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Atmos: SphericalOrientation, latitude, longitude\nusing ClimateMachine.Orientations\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\nimport ClimateMachine.BalanceLaws: boundary_state!\nusing CLIMAParameters.Planet: planet_radius\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\n\n# This is a setup similar to the one presented in [Williamson1992](@cite)\nstruct SolidBodyRotation <: AdvectionDiffusionProblem end\nfunction init_velocity_diffusion!(\n ::SolidBodyRotation,\n aux::Vars,\n geom::LocalGeometry,\n)\n FT = eltype(aux)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n r = norm(geom.coord)\n\n uλ = 2 * FT(π) * cos(φ) * r\n uφ = 0\n aux.advection.u = SVector(\n -uλ * sin(λ) - uφ * cos(λ) * sin(φ),\n +uλ * cos(λ) - uφ * sin(λ) * sin(φ),\n +uφ * cos(φ),\n )\nend\nfunction initial_condition!(::SolidBodyRotation, state, aux, localgeo, t)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n state.ρ = exp(-((3λ)^2 + (3φ)^2))\nend\nfinaltime(::SolidBodyRotation) = 1\nu_scale(::SolidBodyRotation) = 2π\n\n\"\"\"\nThis is the Divergent flow with smooth initial condition test case, the Case 4 in\n\n@article{nair2010class,\n title={A class of deformational flow test cases for linear transport problems on the sphere},\n author={Nair, Ramachandran D and Lauritzen, Peter H},\n journal={Journal of Computational Physics},\n volume={229},\n number={23},\n pages={8868--8887},\n year={2010},\n publisher={Elsevier}\n}\n\n\"\"\"\nstruct ReversingDeformationalFlow <: AdvectionDiffusionProblem end\ninit_velocity_diffusion!(\n ::ReversingDeformationalFlow,\n aux::Vars,\n geom::LocalGeometry,\n) = nothing\nfunction initial_condition!(::ReversingDeformationalFlow, state, aux, coord, t)\n x, y, z = aux.coord\n r = norm(aux.coord)\n h_max = 0.95\n b = 5\n state.ρ = 0\n for (λ, φ) in ((5π / 6, 0), (7π / 6, 0))\n xi = r * cos(φ) * cos(λ)\n yi = r * cos(φ) * sin(λ)\n zi = r * sin(φ)\n state.ρ +=\n h_max * exp(-b * ((x - xi)^2 + (y - yi)^2 + (z - zi)^2) / r^2)\n end\nend\nhas_variable_coefficients(::ReversingDeformationalFlow) = true\nfunction update_velocity_diffusion!(\n ::ReversingDeformationalFlow,\n ::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n FT = eltype(aux)\n λ = longitude(SphericalOrientation(), aux)\n φ = latitude(SphericalOrientation(), aux)\n r = norm(aux.coord)\n T = FT(5)\n λp = λ - FT(2π) * t / T\n uλ =\n 10 * r / T * sin(λp)^2 * sin(2φ) * cos(FT(π) * t / T) +\n FT(2π) * r / T * cos(φ)\n uφ = 10 * r / T * sin(2λp) * cos(φ) * cos(FT(π) * t / T)\n aux.advection.u = SVector(\n -uλ * sin(λ) - uφ * cos(λ) * sin(φ),\n +uλ * cos(λ) - uφ * sin(λ) * sin(φ),\n +uφ * cos(φ),\n )\nend\nu_scale(::ReversingDeformationalFlow) = 2.9\nfinaltime(::ReversingDeformationalFlow) = 5\n\nfunction advective_courant(\n m::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n direction,\n)\n return Δt * norm(aux.advection.u) / Δx\nend\n\nstruct NoFlowBC end\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n ::NoFlowBC,\n ::AdvectionDiffusion,\n stateP::Vars,\n auxP::Vars,\n nM,\n stateM::Vars,\n auxM::Vars,\n t,\n _...,\n)\n auxP.advection.u = -auxM.advection.u\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dgfvm, Q, Qe, model, testname)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dgfvm, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n vert_range,\n topl,\n problem,\n explicit_method,\n fvmethod,\n cfl,\n N,\n timeend,\n FT,\n vtkdir,\n outputtime,\n)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = (N, 0),\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n dx = min_node_distance(grid, HorizontalDirection())\n dt = FT(cfl * dx / (vert_range[2] * u_scale(problem())) / N)\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n bcs = (NoFlowBC(),)\n model = AdvectionDiffusion{3}(problem(), bcs, diffusion = false)\n dgfvm = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dgfvm, FT(0))\n\n odesolver = explicit_method(dgfvm, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n problem = %s\n ArrayType = %s\n FV Reconstruction = %s\n method = %s\n time step = %.16e\n norm(Q₀) = %.16e\"\"\" problem ArrayType fvmethod explicit_method dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(odesolver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n cbcfl = EveryXSimulationSteps(10) do\n dt = ODESolvers.getdt(odesolver)\n cfl = DGMethods.courant(\n advective_courant,\n dgfvm,\n model,\n Q,\n dt,\n HorizontalDirection(),\n )\n @info @sprintf(\n \"\"\"Courant number\n simtime = %.16e\n courant = %.16e\"\"\",\n gettime(odesolver),\n cfl\n )\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Q,\n model,\n \"fvm_advection_sphere\",\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dgfvm, gettime(odesolver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dgfvm,\n Q,\n Qe,\n model,\n \"fvm_advection_sphere\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dgfvm, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n Δmass = abs(weightedsum(Q) - weightedsum(Qe)) / weightedsum(Qe)\n @info @sprintf \"\"\"Finished\n Δmass = %.16e\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" Δmass engf engf / eng0 engf - eng0 errf errf / engfe\n return errf, Δmass\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 3\n\n max_cfl = Dict(LSRK144NiegemannDiehlBusch => 5.0)\n\n expected_result = Dict()\n\n expected_result[SolidBodyRotation, 1, FVConstant] = 1.6249678501611961e+07\n expected_result[SolidBodyRotation, 2, FVConstant] = 7.2020207047738554e+05\n expected_result[SolidBodyRotation, 3, FVConstant] = 5.2452627634607365e+04\n expected_result[SolidBodyRotation, 4, FVConstant] = 3.1132403618328990e+03\n\n expected_result[SolidBodyRotation, 1, FVLinear] = 1.6649963041466445e+07\n expected_result[SolidBodyRotation, 2, FVLinear] = 7.2518593691652094e+05\n expected_result[SolidBodyRotation, 3, FVLinear] = 5.2473203187596591e+04\n expected_result[SolidBodyRotation, 4, FVLinear] = 3.1133127473480104e+03\n\n expected_result[ReversingDeformationalFlow, 1, FVConstant] =\n 2.0097347028034222e+08\n expected_result[ReversingDeformationalFlow, 2, FVConstant] =\n 7.0092390520693690e+07\n expected_result[ReversingDeformationalFlow, 3, FVConstant] =\n 7.7527847763998993e+06\n expected_result[ReversingDeformationalFlow, 4, FVConstant] =\n 1.4209138735343830e+05\n\n expected_result[ReversingDeformationalFlow, 1, FVLinear] =\n 2.0156864805489743e+08\n expected_result[ReversingDeformationalFlow, 2, FVLinear] =\n 7.0092714938572392e+07\n expected_result[ReversingDeformationalFlow, 3, FVLinear] =\n 7.7527849036761308e+06\n expected_result[ReversingDeformationalFlow, 4, FVLinear] =\n 1.4209138776027536e+05\n\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ? 4 :\n 1\n explicit_method = LSRK144NiegemannDiehlBusch\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64,)\n for problem in (SolidBodyRotation, ReversingDeformationalFlow)\n for fvmethod in (FVConstant, FVLinear)\n cfl = max_cfl[explicit_method]\n result = zeros(FT, numlevels)\n for l in 1:numlevels\n numelems_horizontal = 2^(l - 1) * base_num_elem\n numelems_vertical = 3\n\n _planet_radius = FT(planet_radius(param_set))\n domain_height = FT(10e3)\n vert_range =\n (_planet_radius, _planet_radius + domain_height)\n\n topl = StackedCubedSphereTopology(\n mpicomm,\n numelems_horizontal,\n range(\n vert_range[1],\n stop = vert_range[2],\n length = numelems_vertical + 1,\n ),\n )\n\n timeend = finaltime(problem())\n outputtime = timeend\n\n @info (ArrayType, FT)\n vtkdir =\n output ?\n \"vtk_fvm_advection_sphere\" *\n \"_$problem\" *\n \"_$explicit_method\" *\n \"_poly$(polynomialorder)\" *\n \"_$(ArrayType)_$(FT)\" *\n \"_level$(l)\" :\n nothing\n\n result[l], Δmass = test_run(\n mpicomm,\n ArrayType,\n vert_range,\n topl,\n problem,\n explicit_method,\n fvmethod(),\n cfl,\n polynomialorder,\n timeend,\n FT,\n vtkdir,\n outputtime,\n )\n @test result[l] ≈\n FT(expected_result[problem, l, fvmethod])\n @test Δmass <= FT(1e-12)\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/fvm_swirl.jl", "content": "# This tutorial uses the TMAR Filter from [Light2016](@cite)\n#\n# to reproduce the tutorial in section 4b. It is a shear swirling\n# flow deformation of a transported quantity from LeVeque 1996. The exact\n# solution at the final time is the same as the initial condition.\n\nusing MPI\nusing Test\nusing ClimateMachine\nClimateMachine.init()\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)))\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nBase.@kwdef struct SwirlingFlow{FT} <: AdvectionDiffusionProblem\n period::FT = 5\nend\n\ninit_velocity_diffusion!(::SwirlingFlow, aux::Vars, geom::LocalGeometry) =\n nothing\n\nfunction initial_condition!(::SwirlingFlow, state, aux, localgeo, t)\n FT = eltype(state)\n x, y, _ = aux.coord\n x0, y0 = FT(1 // 3), FT(1 // 3)\n τ = 6hypot(x - x0, y - y0)\n state.ρ = exp(-τ^2)\nend\n\nhas_variable_coefficients(::SwirlingFlow) = true\nfunction update_velocity_diffusion!(\n problem::SwirlingFlow,\n ::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n x, y, _ = aux.coord\n sx, cx = sinpi(x), cospi(x)\n sy, cy = sinpi(y), cospi(y)\n ct = cospi(t / problem.period)\n\n u = 2 * sx^2 * sy * cy * ct\n v = -2 * sy^2 * sx * cx * ct\n aux.advection.u = SVector(u, v, 0)\nend\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n testname;\n number_sample_points = 0,\n)\n ## Name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n\n writevtk(\n filename,\n Q,\n dg,\n statenames;\n number_sample_points = number_sample_points,\n )\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## Name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## Name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(pvtuprefix, prefixes, (statenames...,), eltype(Q))\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n topl,\n problem,\n dt,\n N,\n timeend,\n FT,\n vtkdir,\n outputtime,\n)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n bcs = (HomogeneousBC{0}(),)\n model = AdvectionDiffusion{2}(problem, bcs, diffusion = false)\n\n dg = DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n ## We integrate so that the final solution is equal to the initial solution\n Qe = copy(Q)\n odesolver = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n # Set up the information callback\n starttime = Ref(Dates.now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = Dates.now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(odesolver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # Create vtk dir\n if MPI.Comm_rank(mpicomm) == 0\n mkpath(vtkdir)\n end\n MPI.Barrier(mpicomm)\n\n vtkstep = 0\n # Output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n \"swirling_flow\";\n number_sample_points = N[1] + 1,\n )\n\n # Setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(odesolver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n \"swirling_flow\";\n number_sample_points = N[1] + 1,\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)\n\n error = euclidean_distance(Q, Qe)\n\n # Print some end of the simulation information\n eng0 = norm(Qe)\n engf = norm(Q)\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n expected_result = Dict()\n expected_result[1, FVConstant()] = 1.2437314516997458e-01\n expected_result[2, FVConstant()] = 9.8829388561567838e-02\n expected_result[3, FVConstant()] = 7.2096312912198937e-02\n expected_result[4, FVConstant()] = 4.7720226730379636e-02\n expected_result[1, FVLinear()] = 4.3807967239640234e-02\n expected_result[2, FVLinear()] = 1.4913083518239653e-02\n expected_result[3, FVLinear()] = 4.3666055848495802e-03\n expected_result[4, FVLinear()] = 1.2749240022458719e-03\n\n @testset \"$(@__FILE__)\" begin\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ?\n 4 : 1\n FT = Float64\n for fvmethod in (FVConstant(), FVLinear())\n result = zeros(FT, numlevels)\n for l in 1:numlevels\n Ne = 20 * 2^(l - 1)\n polynomialorder = (4, 0)\n\n problem = SwirlingFlow()\n\n brickrange = (\n range(FT(0); length = Ne + 1, stop = 1),\n range(\n FT(0);\n length = polynomialorder[1] * Ne + 1,\n stop = 1,\n ),\n )\n\n topology = StackedBrickTopology(\n mpicomm,\n brickrange,\n boundary = ((1, 1), (1, 1)),\n connectivity = :face,\n )\n\n maxvelocity = 2\n elementsize = 1 / Ne\n dx = elementsize / polynomialorder[1]^2\n CFL = 1\n dt = CFL * dx / maxvelocity\n\n vtkdir = abspath(joinpath(\n ClimateMachine.Settings.output_dir,\n \"fvm_swirl_lvl_$l\",\n ))\n\n timeend = problem.period\n\n outputtime = timeend / 10\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n @info @sprintf \"\"\"Starting\n FT = %s\n ArrayType = %s\n FV Reconstuction = %s\n dim = %d\n Ne = %d\n polynomial order = %d\n final time = %.16e\n time step = %.16e\n \"\"\" FT ArrayType fvmethod 2 Ne polynomialorder[1] timeend dt\n\n result[l] = test_run(\n mpicomm,\n ArrayType,\n fvmethod,\n topology,\n problem,\n dt,\n polynomialorder,\n timeend,\n FT,\n output ? vtkdir : nothing,\n outputtime,\n )\n @test result[l] ≈ FT(expected_result[l, fvmethod])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\"\\n rate for level %d = %e\\n\", l, rate)\n end\n msg\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_bc.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.Mesh.Grids: min_node_distance\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct ConstantHyperDiffusion{FT} <: AdvectionDiffusionProblem\n μ::FT\n k::SVector{3, FT}\nend\n\nfunction init_velocity_diffusion!(\n problem::ConstantHyperDiffusion,\n aux::Vars,\n geom::LocalGeometry,\n)\n FT = eltype(aux)\n aux.hyperdiffusion.H = problem.μ * SMatrix{3, 3, FT}(I)\nend\n\nfunction initial_condition!(\n problem::ConstantHyperDiffusion,\n state,\n aux,\n localgeo,\n t,\n)\n x, y, z = localgeo.coord\n k = problem.k\n k2 = sum(k .^ 2)\n @inbounds begin\n state.ρ =\n cos(k[1] * x) *\n cos(k[2] * y) *\n cos(k[3] * z) *\n exp(-k2^2 * problem.μ * t)\n end\nend\n\n# Boundary conditions data\ninhomogeneous_data!(::Val{0}, p::ConstantHyperDiffusion, x...) =\n initial_condition!(p, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n problem::ConstantHyperDiffusion,\n ∇state,\n aux,\n x,\n t,\n)\n k = problem.k\n k2 = sum(k .^ 2)\n ∇state.ρ =\n -SVector(\n k[1] * sin(k[1] * x[1]) * cos(k[2] * x[2]) * cos(k[3] * x[3]),\n k[2] * cos(k[1] * x[1]) * sin(k[2] * x[2]) * cos(k[3] * x[3]),\n k[3] * cos(k[1] * x[1]) * cos(k[2] * x[2]) * sin(k[3] * x[3]),\n ) * exp(-k2^2 * problem.μ * t)\nend\nfunction inhomogeneous_data!(\n ::Val{2},\n problem::ConstantHyperDiffusion,\n Δstate,\n aux,\n x,\n t,\n)\n k = problem.k\n k2 = sum(k .^ 2)\n Δstate.ρ =\n -k2 *\n cos(k[1] * x[1]) *\n cos(k[2] * x[2]) *\n cos(k[3] * x[3]) *\n exp(-k2^2 * problem.μ * t)\nend\nfunction inhomogeneous_data!(\n ::Val{3},\n problem::ConstantHyperDiffusion,\n ∇Δstate,\n aux,\n x,\n t,\n)\n k = problem.k\n k2 = sum(k .^ 2)\n ∇Δstate.ρ =\n k2 *\n SVector(\n k[1] * sin(k[1] * x[1]) * cos(k[2] * x[2]) * cos(k[3] * x[3]),\n k[2] * cos(k[1] * x[1]) * sin(k[2] * x[2]) * cos(k[3] * x[3]),\n k[3] * cos(k[1] * x[1]) * cos(k[2] * x[2]) * sin(k[3] * x[3]),\n ) *\n exp(-k2^2 * problem.μ * t)\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n vtkdir,\n outputtime,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n μ = 1 // 1000\n dx = min_node_distance(grid)\n dt = dx^4 / 100 / μ\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n bcx1 = (InhomogeneousBC{0}(), InhomogeneousBC{2}())\n bcx2 = (InhomogeneousBC{0}(), InhomogeneousBC{1}())\n\n bcy1 = (InhomogeneousBC{3}(), InhomogeneousBC{1}())\n bcy2 = (InhomogeneousBC{3}(), InhomogeneousBC{2}())\n\n bcz1 = (HomogeneousBC{3}(), HomogeneousBC{1}())\n bcz2 = (HomogeneousBC{3}(), HomogeneousBC{1}())\n\n k = SVector(1, 1, 0)\n bcs = (bcx1, bcx2, bcy1, bcy2, bcz1, bcz2)\n model = AdvectionDiffusion{dim}(\n ConstantHyperDiffusion{FT}(μ, k),\n bcs;\n advection = false,\n diffusion = false,\n hyperdiffusion = true,\n )\n dg = DGModel(\n model,\n grid,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n dim = %d\n dt = %.16e\n norm(Q₀) = %.16e\"\"\" dim dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model, \"hyperdiffusion\")\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"hyperdiffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64] = 1.1666787574326038e-03\n expected_result[2, 2, Float64] = 8.5948289965604964e-05\n expected_result[2, 3, Float64] = 2.5117423516568199e-06\n\n expected_result[3, 1, Float64] = 3.0867418520680958e-03\n expected_result[3, 2, Float64] = 2.2739780086168765e-04\n expected_result[3, 3, Float64] = 6.6454456228180395e-06\n\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ? 3 :\n 1\n\n for FT in (Float64,)\n result = zeros(FT, numlevels)\n for dim in (2, 3)\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n xrange = range(FT(2); length = Ne + 1, stop = FT(9))\n brickrange = ntuple(j -> xrange, dim)\n periodicity = ntuple(j -> false, dim)\n connectivity = dim == 2 ? :face : :full\n boundary = ((1, 2), (3, 4), (5, 6))[1:dim]\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity,\n boundary,\n connectivity,\n )\n timeend = 1\n outputtime = timeend\n\n @info (ArrayType, FT)\n vtkdir =\n output ?\n \"vtk_hyperdiffusion_bc\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n vtkdir,\n outputtime,\n )\n @test result[l] ≈ FT(expected_result[dim, l, FT])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\"\\n rate for level %d = %e\\n\", l, rate)\n end\n msg\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/hyperdiffusion_model.jl", "content": "using StaticArrays\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.BalanceLaws:\n BalanceLaw,\n Prognostic,\n Auxiliary,\n Gradient,\n GradientFlux,\n GradientLaplacian,\n Hyperdiffusive\nimport ClimateMachine.BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n boundary_conditions,\n boundary_state!,\n wavespeed,\n transform_post_gradient_laplacian!\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.DGMethods.NumericalFluxes:\n NumericalFluxFirstOrder, NumericalFluxSecondOrder\n\nabstract type HyperDiffusionProblem end\nstruct HyperDiffusion{dim, P} <: BalanceLaw\n problem::P\n function HyperDiffusion{dim}(\n problem::P,\n ) where {dim, P <: HyperDiffusionProblem}\n new{dim, P}(problem)\n end\nend\n\nvars_state(::HyperDiffusion, ::Auxiliary, FT) = @vars(D::SMatrix{3, 3, FT, 9})\n#\n# Density is only state\nvars_state(::HyperDiffusion, ::Prognostic, FT) = @vars(ρ::FT)\n\n# Take the gradient of density\nvars_state(::HyperDiffusion, ::Gradient, FT) = @vars(ρ::FT)\n# Take the gradient of laplacian of density\nvars_state(::HyperDiffusion, ::GradientLaplacian, FT) = @vars(ρ::FT)\n\nvars_state(::HyperDiffusion, ::GradientFlux, FT) = @vars()\n# The hyperdiffusion DG auxiliary variable: D ∇ Δρ\nvars_state(::HyperDiffusion, ::Hyperdiffusive, FT) = @vars(σ::SVector{3, FT})\n\nfunction flux_first_order!(m::HyperDiffusion, _...) end\n\n\"\"\"\n flux_second_order!(m::HyperDiffusion, flux::Grad, state::Vars,\n auxDG::Vars, auxHDG::Vars, aux::Vars, t::Real)\n\nComputes diffusive flux `F` in:\n\n```\n∂ρ\n-- = - ∇ • (σ) = - ∇ • F\n∂t\n```\nWhere\n\n - `σ` is hyperdiffusion DG auxiliary variable (`σ = D ∇ Δρ` with D being the hyperdiffusion tensor)\n\"\"\"\nfunction flux_second_order!(\n m::HyperDiffusion,\n flux::Grad,\n state::Vars,\n auxDG::Vars,\n auxHDG::Vars,\n aux::Vars,\n t::Real,\n)\n σ = auxHDG.σ\n flux.ρ += σ\nend\n\n\"\"\"\n compute_gradient_argument!(m::HyperDiffusion, transform::Vars, state::Vars,\n aux::Vars, t::Real)\n\nSet the variable to take the gradient of (`ρ` in this case)\n\"\"\"\nfunction compute_gradient_argument!(\n m::HyperDiffusion,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρ = state.ρ\nend\n\ncompute_gradient_flux!(m::HyperDiffusion, _...) = nothing\nfunction transform_post_gradient_laplacian!(\n m::HyperDiffusion,\n auxHDG::Vars,\n gradvars::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ∇Δρ = gradvars.ρ\n D = aux.D\n auxHDG.σ = D * ∇Δρ\nend\n\n\"\"\"\n source!(m::HyperDiffusion, _...)\n\nThere is no source in the hyperdiffusion model\n\"\"\"\nsource!(m::HyperDiffusion, _...) = nothing\n\nfunction init_state_prognostic!(\n m::HyperDiffusion,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n initial_condition!(m.problem, state, aux, localgeo, t)\nend\n\nboundary_conditions(::HyperDiffusion) = ()\nboundary_state!(nf, ::HyperDiffusion, _...) = nothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/periodic_3D_hyperdiffusion.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.Mesh.Grids: min_node_distance\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct ConstantHyperDiffusion{dim, dir, FT} <: AdvectionDiffusionProblem\n D::SMatrix{3, 3, FT, 9}\nend\n\nfunction init_velocity_diffusion!(\n problem::ConstantHyperDiffusion,\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n aux.hyperdiffusion.H = problem.D\nend\n\nfunction initial_condition!(\n problem::ConstantHyperDiffusion{dim, dir},\n state,\n aux,\n localgeo,\n t,\n) where {dim, dir}\n @inbounds begin\n k = SVector(1, 2, 3)\n kD = k * k' .* problem.D\n if dir === EveryDirection()\n c = sum(abs2, k[SOneTo(dim)]) * sum(kD[SOneTo(dim), SOneTo(dim)])\n elseif dir === HorizontalDirection()\n c =\n sum(abs2, k[SOneTo(dim - 1)]) *\n sum(kD[SOneTo(dim - 1), SOneTo(dim - 1)])\n elseif dir === VerticalDirection()\n c = k[dim]^2 * kD[dim, dim]\n end\n x = localgeo.coord\n state.ρ = sin(dot(k[SOneTo(dim)], x[SOneTo(dim)])) * exp(-c * t)\n end\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n direction,\n D,\n vtkdir,\n outputtime,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n dx = min_node_distance(grid)\n dt = dx^4 / 25 / sum(D)\n @info \"time step\" dt\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n model = AdvectionDiffusion{dim}(\n ConstantHyperDiffusion{dim, direction(), FT}(D);\n advection = false,\n diffusion = false,\n hyperdiffusion = true,\n )\n dg = DGModel(\n model,\n grid,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = direction(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, Q, model, \"hyperdiffusion\")\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"hyperdiffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n numlevels =\n integration_testing || ClimateMachine.Settings.integration_testing ? 3 :\n 1\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64, EveryDirection] = 7.6772960298563120e-03\n expected_result[2, 2, Float64, EveryDirection] = 2.3268371815073617e-03\n expected_result[2, 3, Float64, EveryDirection] = 4.2641957936779901e-05\n expected_result[2, 1, Float64, HorizontalDirection] = 8.4650675812606650e-04\n expected_result[2, 2, Float64, HorizontalDirection] = 4.4626814795055979e-05\n expected_result[2, 3, Float64, HorizontalDirection] = 1.2193396277823764e-06\n expected_result[2, 1, Float64, VerticalDirection] = 6.1690465335730834e-03\n expected_result[2, 2, Float64, VerticalDirection] = 2.3407593209031621e-03\n expected_result[2, 3, Float64, VerticalDirection] = 4.3775160749787010e-05\n\n expected_result[3, 1, Float64, EveryDirection] = 1.7363355506160003e-01\n expected_result[3, 2, Float64, EveryDirection] = 7.3049474767548042e-02\n expected_result[3, 3, Float64, EveryDirection] = 5.8530711333407105e-04\n expected_result[3, 1, Float64, HorizontalDirection] = 1.9244127301149615e-02\n expected_result[3, 2, Float64, HorizontalDirection] = 5.8325158696244947e-03\n expected_result[3, 3, Float64, HorizontalDirection] = 1.0688753745025491e-04\n expected_result[3, 1, Float64, VerticalDirection] = 1.4412891107361228e-01\n expected_result[3, 2, Float64, VerticalDirection] = 6.3744013545812925e-02\n expected_result[3, 3, Float64, VerticalDirection] = 9.0891011404938341e-04\n\n numlevels = integration_testing ? 3 : 1\n for FT in (Float64,)\n D =\n 1 // 100 * SMatrix{3, 3, FT}(\n 9 // 50,\n 3 // 50,\n 5 // 50,\n 3 // 50,\n 7 // 50,\n 4 // 50,\n 5 // 50,\n 4 // 50,\n 10 // 50,\n )\n\n result = zeros(FT, numlevels)\n for dim in (2, 3)\n connectivity = dim == 2 ? :face : :full\n for direction in\n (EveryDirection, HorizontalDirection, VerticalDirection)\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n xrange = range(FT(0); length = Ne + 1, stop = FT(2pi))\n brickrange = ntuple(j -> xrange, dim)\n periodicity = ntuple(j -> true, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n connectivity = connectivity,\n )\n timeend = 1\n outputtime = 1\n\n @info (ArrayType, FT, dim, direction)\n vtkdir =\n output ?\n \"vtk_hyperdiffusion\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)_$(direction)\" *\n \"_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n direction,\n D,\n vtkdir,\n outputtime,\n )\n\n @test result[l] ≈ FT(expected_result[dim, l, FT, direction])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\"\\n rate for level %d = %e\\n\", l, rate)\n end\n msg\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n direction,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (InhomogeneousBC{0}(), InhomogeneousBC{1}())\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = direction(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(lsrk))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = callbacks)\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64, EveryDirection] = 1.2357162985295326e-02\n expected_result[2, 2, Float64, EveryDirection] = 8.8403537443388974e-04\n expected_result[2, 3, Float64, EveryDirection] = 4.9490976821250842e-05\n expected_result[2, 4, Float64, EveryDirection] = 2.0311063939957157e-06\n expected_result[2, 1, Float64, HorizontalDirection] = 4.6783743619257120e-02\n expected_result[2, 2, Float64, HorizontalDirection] = 4.0665567827235134e-03\n expected_result[2, 3, Float64, HorizontalDirection] = 5.3144336694498106e-05\n expected_result[2, 4, Float64, HorizontalDirection] = 3.9780001102022647e-07\n expected_result[2, 1, Float64, VerticalDirection] = 4.6783743619257120e-02\n expected_result[2, 2, Float64, VerticalDirection] = 4.0665567827234102e-03\n expected_result[2, 3, Float64, VerticalDirection] = 5.3144336694469002e-05\n expected_result[2, 4, Float64, VerticalDirection] = 3.9780001102679913e-07\n expected_result[3, 1, Float64, EveryDirection] = 9.6252415559793265e-03\n expected_result[3, 2, Float64, EveryDirection] = 6.0160564826756122e-04\n expected_result[3, 3, Float64, EveryDirection] = 4.0359079531195022e-05\n expected_result[3, 4, Float64, EveryDirection] = 2.9987655364452885e-06\n expected_result[3, 1, Float64, HorizontalDirection] = 1.7475667486259477e-02\n expected_result[3, 2, Float64, HorizontalDirection] = 1.2502148161420126e-03\n expected_result[3, 3, Float64, HorizontalDirection] = 6.9990810635609785e-05\n expected_result[3, 4, Float64, HorizontalDirection] = 2.8724182090259972e-06\n expected_result[3, 1, Float64, VerticalDirection] = 6.6162204724938736e-02\n expected_result[3, 2, Float64, VerticalDirection] = 5.7509797542876833e-03\n expected_result[3, 3, Float64, VerticalDirection] = 7.5157441716412944e-05\n expected_result[3, 4, Float64, VerticalDirection] = 5.6257417070312578e-07\n expected_result[2, 1, Float32, EveryDirection] = 1.2357043102383614e-02\n expected_result[2, 2, Float32, EveryDirection] = 8.8409741874784231e-04\n expected_result[2, 3, Float32, EveryDirection] = 4.9193818995263427e-05\n expected_result[2, 1, Float32, HorizontalDirection] = 4.6783901751041412e-02\n expected_result[2, 2, Float32, HorizontalDirection] = 4.0663350373506546e-03\n expected_result[2, 3, Float32, HorizontalDirection] = 5.3044739615870640e-05\n expected_result[2, 1, Float32, VerticalDirection] = 4.6783857047557831e-02\n expected_result[2, 2, Float32, VerticalDirection] = 4.0663424879312515e-03\n expected_result[2, 3, Float32, VerticalDirection] = 5.3172039770288393e-05\n expected_result[3, 1, Float32, EveryDirection] = 9.6252141520380974e-03\n expected_result[3, 2, Float32, EveryDirection] = 6.0162233421579003e-04\n expected_result[3, 3, Float32, EveryDirection] = 4.0522103518014774e-05\n expected_result[3, 1, Float32, HorizontalDirection] = 1.7475549131631851e-02\n expected_result[3, 2, Float32, HorizontalDirection] = 1.2503013713285327e-03\n expected_result[3, 3, Float32, HorizontalDirection] = 7.2867427661549300e-05\n expected_result[3, 1, Float32, VerticalDirection] = 6.6162362694740295e-02\n expected_result[3, 2, Float32, VerticalDirection] = 5.7506239973008633e-03\n expected_result[3, 3, Float32, VerticalDirection] = 1.0193029447691515e-04\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n numlevels =\n integration_testing ||\n ClimateMachine.Settings.integration_testing ?\n (FT == Float64 ? 4 : 3) : 1\n result = zeros(FT, numlevels)\n for dim in 2:3\n connectivity = dim == 2 ? :face : :full\n for direction in\n (EveryDirection, HorizontalDirection, VerticalDirection)\n for fluxBC in (true, false)\n if direction <: EveryDirection\n n =\n dim == 2 ?\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0) :\n SVector{3, FT}(\n 1 / sqrt(3),\n 1 / sqrt(3),\n 1 / sqrt(3),\n )\n elseif direction <: HorizontalDirection\n n =\n dim == 2 ? SVector{3, FT}(1, 0, 0) :\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0)\n elseif direction <: VerticalDirection\n n =\n dim == 2 ? SVector{3, FT}(0, 1, 0) :\n SVector{3, FT}(0, 0, 1)\n end\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = ntuple(\n j -> range(FT(-1); length = Ne + 1, stop = 1),\n dim,\n )\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 2), dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n dt = (α / 4) / (Ne * polynomialorder^2)\n @info \"time step\" dt\n\n timeend = 1\n outputtime = 1\n\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n @info (ArrayType, FT, dim, direction, fluxBC)\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)_$(direction)\" *\n \"_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n direction,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n fluxBC,\n )\n @test result[l] ≈\n FT(expected_result[dim, l, FT, direction])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_1dimex.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing Test\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n if length(ARGS) > 0\n const integration_testing = parse(Bool, ARGS[1])\n else\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\n end\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (\n InhomogeneousBC{0}(),\n InhomogeneousBC{1}(),\n HomogeneousBC{0}(),\n HomogeneousBC{1}(),\n )\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = EveryDirection(),\n )\n\n vdg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n\n Q = init_ode_state(dg, FT(0))\n\n ode_solver = ARK548L2SA2KennedyCarpenter(\n dg,\n vdg,\n LinearBackwardEulerSolver(linearsolvertype(); isadjustable = false),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(ode_solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(\n Q,\n ode_solver;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64] = 7.2801198255507391e-02\n expected_result[2, 2, Float64] = 6.8160295851506783e-03\n expected_result[2, 3, Float64] = 1.4439137164205592e-04\n expected_result[2, 4, Float64] = 2.4260727323386998e-06\n expected_result[3, 1, Float64] = 1.0462203776357534e-01\n expected_result[3, 2, Float64] = 1.0280535683502070e-02\n expected_result[3, 3, Float64] = 2.0631857053908848e-04\n expected_result[3, 4, Float64] = 3.3460492914169325e-06\n expected_result[2, 1, Float32] = 7.2801239788532257e-02\n expected_result[2, 2, Float32] = 6.8159680813550949e-03\n expected_result[2, 3, Float32] = 1.4439738879445940e-04\n # This is near roundoff so we will not check it\n # expected_result[2, 4, Float32] = 2.6432753656990826e-06\n expected_result[3, 1, Float32] = 1.0462204366922379e-01\n expected_result[3, 2, Float32] = 1.0280583053827286e-02\n expected_result[3, 3, Float32] = 2.0646647317335010e-04\n expected_result[3, 4, Float32] = 2.0226731066941284e-05\n\n numlevels = integration_testing ? 4 : 1\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n result = zeros(FT, numlevels)\n for dim in 2:3\n connectivity = dim == 2 ? :face : :full\n for fluxBC in (true, false)\n for linearsolvertype in (SingleColumnLU, ManyColumnLU)\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(\n FT(-1);\n length = Ne + 1,\n stop = 1,\n ),\n dim - 1,\n )...,\n range(FT(-5); length = 5Ne + 1, stop = 5),\n )\n\n periodicity = ntuple(j -> false, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = (\n ntuple(j -> (1, 2), dim - 1)...,\n (3, 4),\n ),\n connectivity = connectivity,\n )\n dt = (α / 4) / (Ne * polynomialorder^2)\n\n outputtime = 0.01\n timeend = 0.5\n\n @info (\n ArrayType,\n FT,\n dim,\n linearsolvertype,\n l,\n fluxBC,\n )\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_$(linearsolvertype)_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n )\n # test the errors significantly larger than floating point epsilon\n if !(dim == 2 && l == 4 && FT == Float32)\n @test result[l] ≈\n FT(expected_result[dim, l, FT])\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/pseudo1D_advection_diffusion_mrigark_implicit.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing Test\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nif !@isdefined integration_testing\n if length(ARGS) > 0\n const integration_testing = parse(Bool, ARGS[1])\n else\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\n end\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(localgeo.coord) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, x)\n ∇state.ρ =\n -(\n 2n * (ξn - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\nend\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, Qe, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n exactnames = statenames .* \"_exact\"\n\n writevtk(filename, Q, dg, statenames, Qe, exactnames)\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., exactnames...),\n eltype(Q),\n )\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n bcs = (\n InhomogeneousBC{0}(),\n InhomogeneousBC{1}(),\n HomogeneousBC{0}(),\n HomogeneousBC{1}(),\n )\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n, α, β, μ, δ}(),\n bcs,\n flux_bc = fluxBC,\n )\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n vdg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n state_auxiliary = dg.state_auxiliary,\n direction = VerticalDirection(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n # With diffussion the Vertical + Horizontal ≠ Full because of 2nd\n # derivative mixing. So we define a custom RHS\n dQ2 = similar(Q)\n function rhs!(dQ, Q, p, time; increment)\n dg(dQ, Q, p, time; increment = increment)\n vdg(dQ2, Q, p, time; increment = false)\n dQ .= dQ .- dQ2\n end\n\n fastsolver = LSRK144NiegemannDiehlBusch(rhs!, Q; dt = dt)\n\n # We're dominated by spatial error, so we can get away with a low order time\n # integrator\n ode_solver = MRIGARKESDIRK46aSandu(\n vdg,\n LinearBackwardEulerSolver(linearsolvertype(); isadjustable = false),\n fastsolver,\n Q;\n dt = dt,\n t0 = 0,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(ode_solver),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n if ~isnothing(vtkdir)\n # create vtk dir\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Q,\n model,\n \"advection_diffusion\",\n )\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n Qe = init_ode_state(dg, gettime(ode_solver))\n do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n Qe,\n model,\n \"advection_diffusion\",\n )\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n solve!(\n Q,\n ode_solver;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64] = 7.3188989633310442e-02\n expected_result[2, 2, Float64] = 6.8155958327716232e-03\n expected_result[2, 3, Float64] = 1.4832570828897563e-04\n expected_result[2, 4, Float64] = 3.3905353801669396e-06\n expected_result[3, 1, Float64] = 1.0501469629884301e-01\n expected_result[3, 2, Float64] = 1.0341917570778314e-02\n expected_result[3, 3, Float64] = 2.1032014288411172e-04\n expected_result[3, 4, Float64] = 4.5797013335024617e-06\n expected_result[2, 1, Float32] = 7.3186099529266357e-02\n expected_result[2, 2, Float32] = 6.8112355656921864e-03\n expected_result[2, 3, Float32] = 1.4748815738130361e-04\n # This is near roundoff so we will not check it\n # expected_result[2, 4, Float32] = 2.9435863325488754e-05\n expected_result[3, 1, Float32] = 1.0500905662775040e-01\n expected_result[3, 2, Float32] = 1.0342594236135483e-02\n expected_result[3, 3, Float32] = 4.2716524330899119e-04\n # This is near roundoff so we will not check it\n # expected_result[3, 4, Float32] = 1.9564463291317225e-03\n\n numlevels = integration_testing ? 4 : 1\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n result = zeros(FT, numlevels)\n for dim in 2:3\n connectivity = dim == 3 ? :full : :face\n for fluxBC in (true, false)\n for linearsolvertype in (SingleColumnLU,)# ManyColumnLU)\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = (\n ntuple(\n j -> range(\n FT(-1);\n length = Ne + 1,\n stop = 1,\n ),\n dim - 1,\n )...,\n range(FT(-5); length = 5Ne + 1, stop = 5),\n )\n\n periodicity = ntuple(j -> false, dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = (\n ntuple(j -> (1, 2), dim - 1)...,\n (3, 4),\n ),\n connectivity = connectivity,\n )\n dt = 32 * (α / 4) / (Ne * polynomialorder^2)\n\n outputtime = 0.01\n timeend = 0.5\n\n @info (\n ArrayType,\n FT,\n dim,\n linearsolvertype,\n l,\n fluxBC,\n )\n vtkdir =\n output ?\n \"vtk_advection\" *\n \"_poly$(polynomialorder)\" *\n \"_dim$(dim)_$(ArrayType)_$(FT)\" *\n \"_$(linearsolvertype)_level$(l)\" :\n nothing\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n dt,\n n,\n α,\n β,\n μ,\n δ,\n vtkdir,\n outputtime,\n linearsolvertype,\n fluxBC,\n )\n # test the errors significantly larger than floating point epsilon\n if !(l == 4 && FT == Float32)\n @test result[l] ≈\n FT(expected_result[dim, l, FT])\n end\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/pseudo1D_heat_eqn.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nimport ClimateMachine.DGMethods.NumericalFluxes:\n normal_boundary_flux_second_order!\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct HeatEqn{n, κ, A} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::HeatEqn{n},\n aux::Vars,\n geom::LocalGeometry,\n) where {n}\n # diffusion of strength 1 in the n direction\n aux.diffusion.D = n * n'\nend\n\n# solution is such that\n# u(1, t) = 1\n# ∇u(0,t) = n\nfunction initial_condition!(\n ::HeatEqn{n, κ, A},\n state,\n aux,\n localgeo,\n t,\n) where {n, κ, A}\n ξn = dot(n, localgeo.coord)\n state.ρ = ξn + sum(A .* cos.(κ * ξn) .* exp.(-κ .^ 2 * t))\nend\ninhomogeneous_data!(::Val{0}, P::HeatEqn, x...) = initial_condition!(P, x...)\n\nfunction normal_boundary_flux_second_order!(\n ::CentralNumericalFluxSecondOrder,\n bcs,\n ::AdvectionDiffusion{1, dim, HeatEqn{nd, κ, A}},\n fluxᵀn::Vars{S},\n n⁻,\n state⁻,\n diff⁻,\n hyperdiff⁻,\n aux⁻,\n state⁺,\n diff⁺,\n hyperdiff⁺,\n aux⁺,\n t,\n _...,\n) where {S, dim, nd, κ, A}\n\n if any_isa(bcs, InhomogeneousBC{0})\n fluxᵀn.ρ = -diff⁻.σ' * n⁻\n elseif any_isa(bcs, InhomogeneousBC{1})\n # Get exact gradient of ρ\n x = aux⁻.coord\n ξn = dot(nd, x)\n ∇ρ = SVector(ntuple(\n i ->\n nd[i] * (1 - sum(A .* κ .* sin.(κ * ξn) .* exp.(-κ .^ 2 * t))),\n Val(3),\n ))\n\n # Compute flux value\n D = aux⁻.diffusion.D\n fluxᵀn.ρ = -(D * ∇ρ)' * n⁻\n end\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n N,\n timeend,\n FT,\n direction,\n dt,\n n,\n κ = 10 * FT(π) / 2,\n A = 1,\n)\n\n numberofsteps = convert(Int64, cld(timeend, dt))\n dt = timeend / numberofsteps\n @info \"time step\" dt numberofsteps dt * numberofsteps timeend\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n bcs = (InhomogeneousBC{1}(), InhomogeneousBC{0}())\n model = AdvectionDiffusion{dim}(HeatEqn{n, κ, A}(), bcs; advection = false)\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = direction(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK144NiegemannDiehlBusch(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n callbacks = (cbinfo,)\n\n solve!(\n Q,\n lsrk;\n numberofsteps = numberofsteps,\n callbacks = callbacks,\n adjustfinalstep = false,\n )\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = Dict()\n expected_result[2, 1, Float64, EveryDirection] = 0.005157483268127576\n expected_result[2, 2, Float64, EveryDirection] = 6.5687731035717e-5\n expected_result[2, 3, Float64, EveryDirection] = 1.6644861275185443e-6\n expected_result[2, 1, Float64, HorizontalDirection] = 0.020515449798977983\n expected_result[2, 2, Float64, HorizontalDirection] = 0.0005686256942296802\n expected_result[2, 3, Float64, HorizontalDirection] = 1.0132022682547854e-5\n expected_result[2, 1, Float64, VerticalDirection] = 0.02051544979897792\n expected_result[2, 2, Float64, VerticalDirection] = 0.0005686256942296017\n expected_result[2, 3, Float64, VerticalDirection] = 1.0132022682754848e-5\n expected_result[3, 1, Float64, EveryDirection] = 0.001260581018671256\n expected_result[3, 2, Float64, EveryDirection] = 2.214908522198975e-5\n expected_result[3, 3, Float64, EveryDirection] = 5.931735594156876e-7\n expected_result[3, 1, Float64, HorizontalDirection] = 0.005157483268127569\n expected_result[3, 2, Float64, HorizontalDirection] = 6.568773103570526e-5\n expected_result[3, 3, Float64, HorizontalDirection] = 1.6644861273865866e-6\n expected_result[3, 1, Float64, VerticalDirection] = 0.020515449798978087\n expected_result[3, 2, Float64, VerticalDirection] = 0.0005686256942297547\n expected_result[3, 3, Float64, VerticalDirection] = 1.0132022682817856e-5\n\n expected_result[2, 1, Float32, EveryDirection] = 0.005157135\n expected_result[2, 2, Float32, EveryDirection] = 6.5721644e-5\n expected_result[2, 3, Float32, EveryDirection] = 3.280845e-6\n expected_result[2, 1, Float32, HorizontalDirection] = 0.020514594\n expected_result[2, 2, Float32, HorizontalDirection] = 0.0005684704\n expected_result[2, 3, Float32, HorizontalDirection] = 1.02350195e-5\n expected_result[2, 1, Float32, VerticalDirection] = 0.020514673\n expected_result[2, 2, Float32, VerticalDirection] = 0.0005684843\n expected_result[2, 3, Float32, VerticalDirection] = 1.0227403e-5\n expected_result[3, 1, Float32, EveryDirection] = 0.0012602004\n expected_result[3, 2, Float32, EveryDirection] = 2.2415780e-5\n expected_result[3, 3, Float32, EveryDirection] = 1.1309192e-5\n expected_result[3, 1, Float32, HorizontalDirection] = 0.005157044\n expected_result[3, 2, Float32, HorizontalDirection] = 6.66792e-5\n expected_result[3, 3, Float32, HorizontalDirection] = 9.930429e-5\n expected_result[3, 1, Float32, VerticalDirection] = 0.020514654\n expected_result[3, 2, Float32, VerticalDirection] = 0.0005684157\n expected_result[3, 3, Float32, VerticalDirection] = 3.224683e-5\n\n @testset \"$(@__FILE__)\" begin\n for FT in (Float64, Float32)\n numlevels =\n integration_testing ||\n ClimateMachine.Settings.integration_testing ?\n 3 : 1\n result = zeros(FT, numlevels)\n for dim in 2:3\n connectivity = dim == 2 ? :face : :full\n for direction in\n (EveryDirection, HorizontalDirection, VerticalDirection)\n if direction <: EveryDirection\n n =\n dim == 2 ?\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0) :\n SVector{3, FT}(\n 1 / sqrt(3),\n 1 / sqrt(3),\n 1 / sqrt(3),\n )\n elseif direction <: HorizontalDirection\n n =\n dim == 2 ? SVector{3, FT}(1, 0, 0) :\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0)\n elseif direction <: VerticalDirection\n n =\n dim == 2 ? SVector{3, FT}(0, 1, 0) :\n SVector{3, FT}(0, 0, 1)\n end\n for l in 1:numlevels\n Ne = 2^(l - 1) * base_num_elem\n brickrange = ntuple(\n j -> range(FT(0); length = Ne + 1, stop = 1),\n dim,\n )\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 2), dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n dt = 1 / (Ne * polynomialorder^2)^2\n\n timeend = 0.01\n\n @info (ArrayType, FT, dim, direction)\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n polynomialorder,\n timeend,\n FT,\n direction,\n dt,\n n,\n )\n\n\n @test (\n result[l] ≈\n FT(expected_result[dim, l, FT, direction]) ||\n result[l] <\n FT(expected_result[dim, l, FT, direction])\n )\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/advection_diffusion/variable_degree_advection_diffusion.jl", "content": "using MPI\nusing ClimateMachine\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Test\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nconst output = parse(Bool, lowercase(get(ENV, \"JULIA_CLIMA_OUTPUT\", \"false\")))\n\ninclude(\"advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n1, n2, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n1, n2, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n1, n2, α, β}\n # Direction of flow is n1 (resp n2) with magnitude α\n aux.advection.u = hcat(α * n1, α * n2)\n\n # diffusion of strength β in the n1 and n2 directions\n aux.diffusion.D = hcat(β * n1 * n1', β * n2 * n2')\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n1, n2, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n1, n2, α, β, μ, δ}\n ξn1 = dot(n1, localgeo.coord)\n ξn2 = dot(n2, localgeo.coord)\n ρ1 = exp(-(ξn1 - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n ρ2 = exp(-(ξn2 - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n state.ρ = (ρ1, ρ2)\nend\n\ninhomogeneous_data!(::Val{0}, P::Pseudo1D, x...) = initial_condition!(P, x...)\n\nfunction inhomogeneous_data!(\n ::Val{1},\n ::Pseudo1D{n1, n2, α, β, μ, δ},\n ∇state,\n aux,\n x,\n t,\n) where {n1, n2, α, β, μ, δ}\n ξn1 = dot(n1, x)\n ξn2 = dot(n2, x)\n ∇ρ1 =\n -(\n 2n1 * (ξn1 - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn1 - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\n ∇ρ2 =\n -(\n 2n2 * (ξn2 - μ - α * t) / (4 * β * (δ + t)) *\n exp(-(ξn2 - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\n )\n ∇state.ρ = hcat(∇ρ1, ∇ρ2)\nend\n\nfunction test_run(mpicomm, dim, polynomialorders, level, ArrayType, FT)\n\n n_hd =\n dim == 2 ? SVector{3, FT}(1, 0, 0) :\n SVector{3, FT}(1 / sqrt(2), 1 / sqrt(2), 0)\n\n n_vd = dim == 2 ? SVector{3, FT}(0, 1, 0) : SVector{3, FT}(0, 0, 1)\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n\n # Grid/topology information\n base_num_elem = 4\n Ne = 2^(level - 1) * base_num_elem\n brickrange = ntuple(j -> range(FT(-1); length = Ne + 1, stop = 1), dim)\n periodicity = ntuple(j -> false, dim)\n bc = ntuple(j -> (1, 2), dim)\n connectivity = dim == 3 ? :full : :face\n topl = StackedBrickTopology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = bc,\n connectivity = connectivity,\n )\n\n dt = (α / 4) / (Ne * maximum(polynomialorders)^2)\n timeend = 1\n @info \"time step\" dt\n\n @info @sprintf \"\"\"Test parameters:\n ArrayType = %s\n FloatType = %s\n Dimension = %s\n Horizontal polynomial order = %s\n Vertical polynomial order = %s\n \"\"\" ArrayType FT dim polynomialorders[1] polynomialorders[end]\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorders,\n )\n\n bcs = (InhomogeneousBC{0}(), InhomogeneousBC{1}())\n # Model being tested\n model = AdvectionDiffusion{dim}(\n Pseudo1D{n_hd, n_vd, α, β, μ, δ}(),\n bcs,\n num_equations = 2,\n )\n\n # Main DG discretization\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n direction = EveryDirection(),\n )\n\n # Initialize all relevant state arrays and create solvers\n Q = init_ode_state(dg, FT(0))\n\n eng0 = norm(Q, dims = (1, 3))\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0[1]\n\n solver = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n solve!(Q, solver; timeend = timeend)\n\n # Reference solution\n engf = norm(Q, dims = (1, 3))\n Q_ref = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Q_ref, dims = (1, 3))\n errf = norm(Q_ref .- Q, dims = (1, 3))\n\n metrics = @. (engf, engf / eng0, engf - eng0, errf, errf / engfe)\n\n @info @sprintf \"\"\"Finished\n Horizontal field:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n Vertical field:\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" first.(metrics)... last.(metrics)...\n\n return errf\nend\n\n\"\"\"\n main()\n\nRun this test problem\n\"\"\"\nfunction main()\n\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n # Dictionary keys: dim, level, polynomial order, FT, and direction\n expected_result = Dict()\n\n # Dim 2, degree 4 in the horizontal, Float64\n expected_result[2, 1, 4, Float64, HorizontalDirection] = 0.0467837436192571\n expected_result[2, 2, 4, Float64, HorizontalDirection] =\n 0.004066556782723549\n expected_result[2, 3, 4, Float64, HorizontalDirection] =\n 5.3144336694234015e-5\n expected_result[2, 4, 4, Float64, HorizontalDirection] =\n 3.978000110046181e-7\n\n # Dim 2, degree 2 in the vertical, Float64\n expected_result[2, 1, 2, Float64, VerticalDirection] = 0.15362016594121006\n expected_result[2, 2, 2, Float64, VerticalDirection] = 0.04935353328794371\n expected_result[2, 3, 2, Float64, VerticalDirection] = 0.015530511948609192\n expected_result[2, 4, 2, Float64, VerticalDirection] = 0.0006275095484456197\n\n # Dim 2, degree 2 in the horizontal, Float64\n expected_result[2, 1, 2, Float64, HorizontalDirection] = 0.15362016594121003\n expected_result[2, 2, 2, Float64, HorizontalDirection] = 0.04935353328794369\n expected_result[2, 3, 2, Float64, HorizontalDirection] =\n 0.015530511948609204\n expected_result[2, 4, 2, Float64, HorizontalDirection] =\n 0.0006275095484455967\n\n # Dim 2, degree 4 in the vertical, Float64\n expected_result[2, 1, 4, Float64, VerticalDirection] = 0.04678374361925714\n expected_result[2, 2, 4, Float64, VerticalDirection] = 0.0040665567827235\n expected_result[2, 3, 4, Float64, VerticalDirection] = 5.3144336694109365e-5\n expected_result[2, 4, 4, Float64, VerticalDirection] = 3.978000109805811e-7\n\n # Dim 3, degree 4 in the horizontal, Float64\n expected_result[3, 1, 4, Float64, HorizontalDirection] =\n 0.017475667486259432\n expected_result[3, 2, 4, Float64, HorizontalDirection] =\n 0.0012502148161420109\n expected_result[3, 3, 4, Float64, HorizontalDirection] =\n 6.999081063570052e-5\n expected_result[3, 4, 4, Float64, HorizontalDirection] =\n 2.8724182090419642e-6\n\n # Dim 3, degree 2 in the vertical, Float64\n expected_result[3, 1, 2, Float64, VerticalDirection] = 0.2172517221280645\n expected_result[3, 2, 2, Float64, VerticalDirection] = 0.06979643612684193\n expected_result[3, 3, 2, Float64, VerticalDirection] = 0.02196346062832051\n expected_result[3, 4, 2, Float64, VerticalDirection] = 0.0008874325139302493\n\n # Dim 3, degree 2 in the horizontal, Float64\n expected_result[3, 1, 2, Float64, HorizontalDirection] = 0.10343354980172516\n expected_result[3, 2, 2, Float64, HorizontalDirection] = 0.03415137756593495\n expected_result[3, 3, 2, Float64, HorizontalDirection] =\n 0.0035959803480493553\n expected_result[3, 4, 2, Float64, HorizontalDirection] =\n 0.0002714157844893719\n\n # Dim 3, degree 4 in the vertical, Float64\n expected_result[3, 1, 4, Float64, VerticalDirection] = 0.06616220472493903\n expected_result[3, 2, 4, Float64, VerticalDirection] = 0.005750979754288175\n expected_result[3, 3, 4, Float64, VerticalDirection] = 7.515744171591452e-5\n expected_result[3, 4, 4, Float64, VerticalDirection] = 5.625741705890895e-7\n\n # Dim 2, degree 4 in the horizontal, Float32\n expected_result[2, 1, 4, Float32, HorizontalDirection] = 0.046783954f0\n expected_result[2, 2, 4, Float32, HorizontalDirection] = 0.004066328f0\n expected_result[2, 3, 4, Float32, HorizontalDirection] = 5.327546f-5\n\n # Dim 2, degree 2 in the vertical, Float32\n expected_result[2, 1, 2, Float32, VerticalDirection] = 0.15362015f0\n expected_result[2, 2, 2, Float32, VerticalDirection] = 0.04935346f0\n expected_result[2, 3, 2, Float32, VerticalDirection] = 0.015530386f0\n\n # Dim 2, degree 2 in the horizontal, Float32\n expected_result[2, 1, 2, Float32, HorizontalDirection] = 0.1536202f0\n expected_result[2, 2, 2, Float32, HorizontalDirection] = 0.04935346f0\n expected_result[2, 3, 2, Float32, HorizontalDirection] = 0.015530357f0\n\n # Dim 2, degree 4 in the vertical, Float32\n expected_result[2, 1, 4, Float32, VerticalDirection] = 0.04678398f0\n expected_result[2, 2, 4, Float32, VerticalDirection] = 0.0040662177f0\n expected_result[2, 3, 4, Float32, VerticalDirection] = 5.3401447f-5\n\n # Dim 3, degree 4 in the horizontal, Float32\n expected_result[3, 1, 4, Float32, HorizontalDirection] = 0.01747554f0\n expected_result[3, 2, 4, Float32, HorizontalDirection] = 0.0012502924f0\n expected_result[3, 3, 4, Float32, HorizontalDirection] = 7.00218f-5\n\n # Dim 3, degree 2 in the vertical, Float32\n expected_result[3, 1, 2, Float32, VerticalDirection] = 0.21725166f0\n expected_result[3, 2, 2, Float32, VerticalDirection] = 0.06979626f0\n expected_result[3, 3, 2, Float32, VerticalDirection] = 0.021963252f0\n\n # Dim 3, degree 2 in the horizontal, Float32\n expected_result[3, 1, 2, Float32, HorizontalDirection] = 0.10343349f0\n expected_result[3, 2, 2, Float32, HorizontalDirection] = 0.034151305f0\n expected_result[3, 3, 2, Float32, HorizontalDirection] = 0.0035958516f0\n\n # Dim 3, degree 4 in the vertical, Float32\n expected_result[3, 1, 4, Float32, VerticalDirection] = 0.06616244f0\n expected_result[3, 2, 4, Float32, VerticalDirection] = 0.005750495f0\n expected_result[3, 3, 4, Float32, VerticalDirection] = 7.538217f-5\n\n @testset \"Variable degree DG: advection diffusion model\" begin\n for FT in (Float32, Float64)\n numlevels =\n integration_testing ||\n ClimateMachine.Settings.integration_testing ?\n (FT == Float64 ? 4 : 3) : 1\n for dim in 2:3\n for polynomialorders in ((4, 2), (2, 4))\n result = Dict()\n for level in 1:numlevels\n result[level] = test_run(\n mpicomm,\n dim,\n polynomialorders,\n level,\n ArrayType,\n FT,\n )\n horiz_poly = polynomialorders[1]\n vert_poly = polynomialorders[2]\n @test result[level][1] ≈ FT(expected_result[\n dim,\n level,\n horiz_poly,\n FT,\n HorizontalDirection,\n ])\n @test result[level][2] ≈ FT(expected_result[\n dim,\n level,\n vert_poly,\n FT,\n VerticalDirection,\n ])\n end\n @info begin\n msg = \"\"\n for l in 1:(numlevels - 1)\n rate = @. log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rates for level %d Horizontal = %e\",\n l,\n rate[1]\n )\n msg *= @sprintf(\", Vertical = %e\\n\", rate[2])\n end\n msg\n end\n end\n end\n end\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/density_current_model.jl", "content": "using Test\nusing Dates\nusing LinearAlgebra\nusing MPI\nusing Printf\nusing Random\nusing StaticArrays\n\nusing ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.VariableTemplates\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\n# -------------- Problem constants ------------------- #\nconst dim = 3\nconst (xmin, xmax) = (0, 12800)\nconst (ymin, ymax) = (0, 400)\nconst (zmin, zmax) = (0, 6400)\nconst Ne = (100, 2, 50)\nconst polynomialorder = 4\nconst dt = 0.01\nconst timeend = 10dt\n\n# ------------- Initial condition function ----------- #\n\"\"\"\n# Reference\n\nSee [Straka1993](@cite)\n\"\"\"\nfunction Initialise_Density_Current!(\n problem,\n bl,\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n)\n (x1, x2, x3) = localgeo.coord\n param_set = parameter_set(bl)\n FT = eltype(state)\n _R_d::FT = R_d(param_set)\n _grav::FT = grav(param_set)\n _cp_d::FT = cp_d(param_set)\n _cv_d::FT = cv_d(param_set)\n _MSLP::FT = MSLP(param_set)\n # initialise with dry domain\n q_tot::FT = 0\n q_liq::FT = 0\n q_ice::FT = 0\n # perturbation parameters for rising bubble\n rx = 4000\n rz = 2000\n xc = 0\n zc = 3000\n r = sqrt((x1 - xc)^2 / rx^2 + (x3 - zc)^2 / rz^2)\n θ_ref::FT = 300\n θ_c::FT = -15\n Δθ::FT = 0\n if r <= 1\n Δθ = θ_c * (1 + cospi(r)) / 2\n end\n qvar = PhasePartition(q_tot)\n θ = θ_ref + Δθ # potential temperature\n π_exner = FT(1) - _grav / (_cp_d * θ) * x3 # exner pressure\n ρ = _MSLP / (_R_d * θ) * (π_exner)^(_cv_d / _R_d) # density\n\n ts = PhaseEquil_ρθq(param_set, ρ, θ, q_tot)\n q_pt = PhasePartition(ts)\n\n U, V, W = FT(0), FT(0), FT(0) # momentum components\n # energy definitions\n e_kin = (U^2 + V^2 + W^2) / (2 * ρ) / ρ\n e_pot = gravitational_potential(bl, aux)\n e_int = internal_energy(ts)\n E = ρ * (e_int + e_kin + e_pot) #* total_energy(e_kin, e_pot, T, q_tot, q_liq, q_ice)\n state.ρ = ρ\n state.ρu = SVector(U, V, W)\n state.energy.ρe = E\n state.moisture.ρq_tot = ρ * q_pt.tot\nend\n# --------------- Driver definition ------------------ #\nfunction test_run(\n mpicomm,\n ArrayType,\n topl,\n dim,\n Ne,\n polynomialorder,\n timeend,\n FT,\n dt,\n)\n # -------------- Define grid ----------------------------------- #\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n # -------------- Define model ---------------------------------- #\n T_profile = DryAdiabaticProfile{FT}(param_set)\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(T_profile),\n turbulence = AnisoMinDiss{FT}(1),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = Initialise_Density_Current!,\n source = (Gravity(),),\n )\n # -------------- Define DGModel --------------------------- #\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\n FloatType = %s\"\"\" eng0 ArrayType FT\n\n # Set up the information callback (output field dump is via vtk callback: see cbinfo)\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(10, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n vtkstep = [0]\n cbvtk = GenericCallbacks.EveryXSimulationSteps(3000) do (init = false)\n mkpath(\"./vtk-dc/\")\n outprefix = @sprintf(\n \"./vtk-dc/DC_%dD_mpirank%04d_step%04d\",\n dim,\n MPI.Comm_rank(mpicomm),\n vtkstep[1]\n )\n @debug \"doing VTK output\" outprefix\n writevtk(\n outprefix,\n Q,\n dg,\n flattenednames(vars_state(model, Prognostic(), FT)),\n dg.state_auxiliary,\n flattenednames(vars_state(model, Auxiliary(), FT)),\n )\n vtkstep[1] += 1\n nothing\n end\n\n\n solve!(Q, lsrk; timeend = timeend, callbacks = (cbinfo, cbvtk))\n # End of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n engf / eng0\nend\n# --------------- Test block / Loggers ------------------ #\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n for FT in (Float32, Float64)\n brickrange = (\n range(FT(xmin); length = Ne[1] + 1, stop = xmax),\n range(FT(ymin); length = Ne[2] + 1, stop = ymax),\n range(FT(zmin); length = Ne[3] + 1, stop = zmax),\n )\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, true, false),\n )\n engf_eng0 = test_run(\n mpicomm,\n ArrayType,\n topl,\n dim,\n Ne,\n polynomialorder,\n timeend,\n FT,\n dt,\n )\n @test engf_eng0 ≈ FT(9.9999970927037096e-01)\n end\nend\n\n\n#nothing\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_atmos.jl", "content": "using Test\nusing Dates\nusing LinearAlgebra\nusing MPI\nusing Printf\nusing StaticArrays\nusing UnPack\n\nusing ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.Orientations\nusing ClimateMachine.VariableTemplates\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VTK\n\nimport ClimateMachine.BalanceLaws: source, prognostic_vars\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport CLIMAParameters\n# Assume zero reference temperature\nCLIMAParameters.Planet.T_0(::EarthParameterSet) = 0\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\ninclude(\"mms_solution_generated.jl\")\n\nimport Thermodynamics: total_specific_enthalpy\nusing ClimateMachine.Atmos\n\ntotal_specific_enthalpy(ts::PhaseDry{FT}, e_tot::FT) where {FT <: Real} =\n zero(FT)\n\nfunction mms2_init_state!(problem, bl, state::Vars, aux::Vars, localgeo, t)\n (x1, x2, x3) = localgeo.coord\n state.ρ = ρ_g(t, x1, x2, x3, Val(2))\n state.ρu = SVector(\n U_g(t, x1, x2, x3, Val(2)),\n V_g(t, x1, x2, x3, Val(2)),\n W_g(t, x1, x2, x3, Val(2)),\n )\n state.energy.ρe = E_g(t, x1, x2, x3, Val(2))\nend\n\nstruct MMSSource{N} <: TendencyDef{Source} end\n\nprognostic_vars(::MMSSource{N}) where {N} = (Mass(), Momentum(), Energy())\n\nfunction source(::Mass, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return Sρ_g(t, x1, x2, x3, Val(N))\nend\nfunction source(::Momentum, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return SVector(\n SU_g(t, x1, x2, x3, Val(N)),\n SV_g(t, x1, x2, x3, Val(N)),\n SW_g(t, x1, x2, x3, Val(N)),\n )\nend\nfunction source(::Energy, s::MMSSource{N}, m, args) where {N}\n @unpack aux, t = args\n x1, x2, x3 = aux.coord\n return SE_g(t, x1, x2, x3, Val(N))\nend\n\nfunction mms3_init_state!(problem, bl, state::Vars, aux::Vars, localgeo, t)\n (x1, x2, x3) = localgeo.coord\n state.ρ = ρ_g(t, x1, x2, x3, Val(3))\n state.ρu = SVector(\n U_g(t, x1, x2, x3, Val(3)),\n V_g(t, x1, x2, x3, Val(3)),\n W_g(t, x1, x2, x3, Val(3)),\n )\n state.energy.ρe = E_g(t, x1, x2, x3, Val(3))\nend\n\n# initial condition\n\nfunction test_run(mpicomm, ArrayType, dim, topl, warpfun, N, timeend, FT, dt)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(μ_exact), WithDivergence()),\n moisture = DryModel(),\n )\n if dim == 2\n problem = AtmosProblem(\n boundaryconditions = (InitStateBC(),),\n init_state_prognostic = mms2_init_state!,\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (MMSSource{2}(),),\n )\n else\n problem = AtmosProblem(\n boundaryconditions = (InitStateBC(),),\n init_state_prognostic = mms3_init_state!,\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n orientation = NoOrientation(),\n source = (MMSSource{3}(),),\n )\n end\n show_tendencies(model)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n Qcpu = init_ode_state(dg, FT(0); init_on_cpu = true)\n @test euclidean_distance(Q, Qcpu) < sqrt(eps(FT))\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = (cbinfo,))\n # solve!(Q, lsrk; timeend=timeend, callbacks=(cbinfo, cbvtk))\n\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = [\n 1.6931876910307017e-01 5.4603193051929394e-03 2.3307776694542282e-04\n 3.3983777728925593e-02 1.7808380837573065e-03 9.176181458773599e-5\n ]\n lvls = integration_testing ? size(expected_result, 2) : 1\n\n @testset \"mms_bc_atmos\" begin\n for FT in (Float64,) #Float32)\n result = zeros(FT, lvls)\n for dim in 2:3\n for l in 1:lvls\n if dim == 2\n Ne = (\n 2^(l - 1) * base_num_elem,\n 2^(l - 1) * base_num_elem,\n )\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false),\n )\n dt = 1e-2 / Ne[1]\n warpfun =\n (x1, x2, _) -> begin\n (x1 + sin(x1 * x2), x2 + sin(2 * x1 * x2), 0)\n end\n\n elseif dim == 3\n Ne = (\n 2^(l - 1) * base_num_elem,\n 2^(l - 1) * base_num_elem,\n )\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false, false),\n )\n dt = 5e-3 / Ne[1]\n warpfun =\n (x1, x2, x3) -> begin\n (\n x1 +\n (x1 - 1 / 2) * cos(2 * π * x2 * x3) / 4,\n x2 + exp(sin(2π * (x1 * x2 + x3))) / 20,\n x3 + x1 / 4 + x2^2 / 2 + sin(x1 * x2 * x3),\n )\n end\n end\n timeend = 1\n nsteps = ceil(Int64, timeend / dt)\n dt = timeend / nsteps\n\n @info (ArrayType, FT, dim, nsteps, dt)\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n warpfun,\n polynomialorder,\n timeend,\n FT,\n dt,\n )\n @test result[l] ≈ FT(expected_result[dim - 1, l])\n end\n if integration_testing\n @info begin\n msg = \"\"\n for l in 1:(lvls - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/mms_bc_dgmodel.jl", "content": "using MPI\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\nusing ClimateMachine.VTK\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\ninclude(\"mms_solution_generated.jl\")\ninclude(\"mms_model.jl\")\n\n\n# initial condition\n\nfunction test_run(mpicomm, ArrayType, dim, topl, warpfun, N, timeend, FT, dt)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n dg = DGModel(\n MMSModel{dim}(),\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\"\"\" eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(lsrk),\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = (cbinfo,))\n # solve!(Q, lsrk; timeend=timeend, callbacks=(cbinfo, cbvtk))\n\n\n # Print some end of the simulation information\n engf = norm(Q)\n Qe = init_ode_state(dg, FT(timeend))\n\n engfe = norm(Qe)\n errf = euclidean_distance(Q, Qe)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n norm(Q - Qe) = %.16e\n norm(Q - Qe) / norm(Qe) = %.16e\n \"\"\" engf engf / eng0 engf - eng0 errf errf / engfe\n errf\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n\n expected_result = [\n 1.6694721292986181e-01 5.4178750150416337e-03 2.3066867400713085e-04\n 3.3672443923201158e-02 1.7603832251132654e-03 9.1108401774885506e-05\n ]\n lvls = integration_testing ? size(expected_result, 2) : 1\n\n @testset \"mms_bc_dgmodel\" begin\n for FT in (Float64,) #Float32)\n result = zeros(FT, lvls)\n for dim in 2:3\n for l in 1:lvls\n if dim == 2\n Ne = (\n 2^(l - 1) * base_num_elem,\n 2^(l - 1) * base_num_elem,\n )\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false),\n )\n dt = 1e-2 / Ne[1]\n warpfun =\n (x1, x2, _) -> begin\n (x1 + sin(x1 * x2), x2 + sin(2 * x1 * x2), 0)\n end\n\n elseif dim == 3\n Ne = (\n 2^(l - 1) * base_num_elem,\n 2^(l - 1) * base_num_elem,\n )\n brickrange = (\n range(FT(0); length = Ne[1] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n range(FT(0); length = Ne[2] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false, false),\n )\n dt = 5e-3 / Ne[1]\n warpfun =\n (x1, x2, x3) -> begin\n (\n x1 +\n (x1 - 1 / 2) * cos(2 * π * x2 * x3) / 4,\n x2 + exp(sin(2π * (x1 * x2 + x3))) / 20,\n x3 + x1 / 4 + x2^2 / 2 + sin(x1 * x2 * x3),\n )\n end\n end\n timeend = 1\n nsteps = ceil(Int64, timeend / dt)\n dt = timeend / nsteps\n\n @info (ArrayType, FT, dim)\n result[l] = test_run(\n mpicomm,\n ArrayType,\n dim,\n topl,\n warpfun,\n polynomialorder,\n timeend,\n FT,\n dt,\n )\n @test result[l] ≈ FT(expected_result[dim - 1, l])\n end\n if integration_testing\n @info begin\n msg = \"\"\n for l in 1:(lvls - 1)\n rate = log2(result[l]) - log2(result[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/mms_model.jl", "content": "using ClimateMachine.VariableTemplates\n\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux\n\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_conditions,\n boundary_state!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!\n\nimport ClimateMachine.DGMethods: init_ode_state\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\n\nstruct MMSModel{dim} <: BalanceLaw end\n\nvars_state(::MMSModel, ::Auxiliary, T) = @vars(x1::T, x2::T, x3::T)\nvars_state(::MMSModel, ::Prognostic, T) =\n @vars(ρ::T, ρu::T, ρv::T, ρw::T, ρe::T)\nvars_state(::MMSModel, ::Gradient, T) = @vars(u::T, v::T, w::T)\nvars_state(::MMSModel, ::GradientFlux, T) =\n @vars(τ11::T, τ22::T, τ33::T, τ12::T, τ13::T, τ23::T)\n\nfunction flux_first_order!(\n ::MMSModel,\n flux::Grad,\n state::Vars,\n auxstate::Vars,\n t::Real,\n direction,\n)\n # preflux\n T = eltype(flux)\n γ = T(γ_exact)\n ρinv = 1 / state.ρ\n u, v, w = ρinv * state.ρu, ρinv * state.ρv, ρinv * state.ρw\n P = (γ - 1) * (state.ρe - ρinv * (state.ρu^2 + state.ρv^2 + state.ρw^2) / 2)\n\n # invisc terms\n flux.ρ = SVector(state.ρu, state.ρv, state.ρw)\n flux.ρu = SVector(u * state.ρu + P, v * state.ρu, w * state.ρu)\n flux.ρv = SVector(u * state.ρv, v * state.ρv + P, w * state.ρv)\n flux.ρw = SVector(u * state.ρw, v * state.ρw, w * state.ρw + P)\n flux.ρe =\n SVector(u * (state.ρe + P), v * (state.ρe + P), w * (state.ρe + P))\nend\n\nfunction flux_second_order!(\n ::MMSModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n auxstate::Vars,\n t::Real,\n)\n ρinv = 1 / state.ρ\n u, v, w = ρinv * state.ρu, ρinv * state.ρv, ρinv * state.ρw\n\n # viscous terms\n flux.ρu -= SVector(diffusive.τ11, diffusive.τ12, diffusive.τ13)\n flux.ρv -= SVector(diffusive.τ12, diffusive.τ22, diffusive.τ23)\n flux.ρw -= SVector(diffusive.τ13, diffusive.τ23, diffusive.τ33)\n\n flux.ρe -= SVector(\n u * diffusive.τ11 + v * diffusive.τ12 + w * diffusive.τ13,\n u * diffusive.τ12 + v * diffusive.τ22 + w * diffusive.τ23,\n u * diffusive.τ13 + v * diffusive.τ23 + w * diffusive.τ33,\n )\nend\n\nfunction compute_gradient_argument!(\n ::MMSModel,\n transformstate::Vars,\n state::Vars,\n auxstate::Vars,\n t::Real,\n)\n ρinv = 1 / state.ρ\n transformstate.u = ρinv * state.ρu\n transformstate.v = ρinv * state.ρv\n transformstate.w = ρinv * state.ρw\nend\n\nfunction compute_gradient_flux!(\n ::MMSModel,\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n auxstate::Vars,\n t::Real,\n)\n T = eltype(diffusive)\n μ = T(μ_exact)\n\n dudx, dudy, dudz = ∇transform.u\n dvdx, dvdy, dvdz = ∇transform.v\n dwdx, dwdy, dwdz = ∇transform.w\n\n # strains\n ϵ11 = dudx\n ϵ22 = dvdy\n ϵ33 = dwdz\n ϵ12 = (dudy + dvdx) / 2\n ϵ13 = (dudz + dwdx) / 2\n ϵ23 = (dvdz + dwdy) / 2\n\n # deviatoric stresses\n diffusive.τ11 = 2μ * (ϵ11 - (ϵ11 + ϵ22 + ϵ33) / 3)\n diffusive.τ22 = 2μ * (ϵ22 - (ϵ11 + ϵ22 + ϵ33) / 3)\n diffusive.τ33 = 2μ * (ϵ33 - (ϵ11 + ϵ22 + ϵ33) / 3)\n diffusive.τ12 = 2μ * ϵ12\n diffusive.τ13 = 2μ * ϵ13\n diffusive.τ23 = 2μ * ϵ23\nend\n\nfunction source!(\n ::MMSModel{dim},\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n t::Real,\n direction,\n) where {dim}\n source.ρ = Sρ_g(t, aux.x1, aux.x2, aux.x3, Val(dim))\n source.ρu = SU_g(t, aux.x1, aux.x2, aux.x3, Val(dim))\n source.ρv = SV_g(t, aux.x1, aux.x2, aux.x3, Val(dim))\n source.ρw = SW_g(t, aux.x1, aux.x2, aux.x3, Val(dim))\n source.ρe = SE_g(t, aux.x1, aux.x2, aux.x3, Val(dim))\nend\n\nfunction wavespeed(::MMSModel, nM, state::Vars, aux::Vars, t::Real, direction)\n T = eltype(state)\n γ = T(γ_exact)\n ρinv = 1 / state.ρ\n u, v, w = ρinv * state.ρu, ρinv * state.ρv, ρinv * state.ρw\n P = (γ - 1) * (state.ρe - ρinv * (state.ρu^2 + state.ρv^2 + state.ρw^2) / 2)\n return abs(nM[1] * u + nM[2] * v + nM[3] * w) + sqrt(ρinv * γ * P)\nend\n\nboundary_conditions(::MMSModel) = (nothing,)\n\nfunction boundary_state!(\n ::RusanovNumericalFlux,\n bctype,\n bl::MMSModel,\n stateP::Vars,\n auxP::Vars,\n nM,\n stateM::Vars,\n auxM::Vars,\n t,\n _...,\n)\n init_state_prognostic!(\n bl,\n stateP,\n auxP,\n (coord = (auxM.x1, auxM.x2, auxM.x3),),\n t,\n )\nend\n\n# FIXME: This is probably not right....\nboundary_state!(::CentralNumericalFluxGradient, bc, bl::MMSModel, _...) =\n nothing\n\nfunction boundary_state!(\n ::CentralNumericalFluxSecondOrder,\n bctype,\n bl::MMSModel,\n stateP::Vars,\n diffP::Vars,\n hyperdiffP::Vars,\n auxP::Vars,\n nM,\n stateM::Vars,\n diffM::Vars,\n hyperdiffM::Vars,\n auxM::Vars,\n t,\n _...,\n)\n init_state_prognostic!(\n bl,\n stateP,\n auxP,\n (coord = (auxM.x1, auxM.x2, auxM.x3),),\n t,\n )\nend\n\nfunction nodal_init_state_auxiliary!(\n ::MMSModel,\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n)\n x1, x2, x3 = g.coord\n aux.x1 = x1\n aux.x2 = x2\n aux.x3 = x3\nend\n\nfunction init_state_prognostic!(\n bl::MMSModel{dim},\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n) where {dim}\n (x1, x2, x3) = localgeo.coord\n state.ρ = ρ_g(t, x1, x2, x3, Val(dim))\n state.ρu = U_g(t, x1, x2, x3, Val(dim))\n state.ρv = V_g(t, x1, x2, x3, Val(dim))\n state.ρw = W_g(t, x1, x2, x3, Val(dim))\n state.ρe = E_g(t, x1, x2, x3, Val(dim))\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/mms_solution.jl", "content": "# This file generates the solution used in method of manufactured solutions\nusing LinearAlgebra, SymPy, Printf\nusing CLIMAParameters\nusing CLIMAParameters.Planet\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n@syms x y z t real = true\nμ = 1 // 100\nγ = cp_d(param_set) / cv_d(param_set)\n\noutput = open(\"mms_solution_generated.jl\", \"w\")\n\n@printf output \"const γ_exact = %s\\n\" γ\n@printf output \"const μ_exact = %s\\n\" μ\n\nfor dim in 2:3\n if dim == 3\n ρ = cos(π * t) * sin(π * x) * cos(π * y) * cos(π * z) + 3\n U = cos(π * t) * ρ * sin(π * x) * cos(π * y) * cos(π * z)\n V = cos(π * t) * ρ * sin(π * x) * cos(π * y) * cos(π * z)\n W = cos(π * t) * ρ * sin(π * x) * cos(π * y) * sin(π * z)\n E = cos(π * t) * sin(π * x) * cos(π * y) * cos(π * z) + 100\n else\n ρ = cos(π * t) * sin(π * x) * cos(π * y) + 3\n U = cos(π * t) * ρ * sin(π * x) * cos(π * y)\n V = cos(π * t) * ρ * sin(π * x) * cos(π * y)\n W = cos(π * t) * 0\n E = cos(π * t) * sin(π * x) * cos(π * y) + 100\n end\n\n\n P = (γ - 1) * (E - (U^2 + V^2 + W^2) / 2ρ)\n\n u, v, w = U / ρ, V / ρ, W / ρ\n\n dudx, dudy, dudz = diff(u, x), diff(u, y), diff(u, z)\n dvdx, dvdy, dvdz = diff(v, x), diff(v, y), diff(v, z)\n dwdx, dwdy, dwdz = diff(w, x), diff(w, y), diff(w, z)\n\n ϵ11 = dudx\n ϵ22 = dvdy\n ϵ33 = dwdz\n ϵ12 = (dudy + dvdx) / 2\n ϵ13 = (dudz + dwdx) / 2\n ϵ23 = (dvdz + dwdy) / 2\n\n τ11 = 2μ * (ϵ11 - (ϵ11 + ϵ22 + ϵ33) / 3)\n τ22 = 2μ * (ϵ22 - (ϵ11 + ϵ22 + ϵ33) / 3)\n τ33 = 2μ * (ϵ33 - (ϵ11 + ϵ22 + ϵ33) / 3)\n τ12 = τ21 = 2μ * ϵ12\n τ13 = τ31 = 2μ * ϵ13\n τ23 = τ32 = 2μ * ϵ23\n\n Fx_x =\n diff.(\n [\n U\n u * U + P - τ11\n u * V - τ12\n u * W - τ13\n u * (E + P) - u * τ11 - v * τ12 - w * τ13\n ],\n x,\n )\n Fy_y =\n diff.(\n [\n V\n v * U - τ21\n v * V + P - τ22\n v * W - τ23\n v * (E + P) - u * τ21 - v * τ22 - w * τ23\n ],\n y,\n )\n Fz_z =\n diff.(\n [\n W\n w * U - τ31\n w * V - τ32\n w * W + P - τ33\n w * (E + P) - u * τ31 - v * τ32 - w * τ33\n ],\n z,\n )\n\n dρdt = simplify(Fx_x[1] + Fy_y[1] + Fz_z[1] + diff(ρ, t))\n dUdt = simplify(Fx_x[2] + Fy_y[2] + Fz_z[2] + diff(U, t))\n dVdt = simplify(Fx_x[3] + Fy_y[3] + Fz_z[3] + diff(V, t))\n dWdt = simplify(Fx_x[4] + Fy_y[4] + Fz_z[4] + diff(W, t))\n dEdt = simplify(Fx_x[5] + Fy_y[5] + Fz_z[5] + diff(E, t))\n\n\n @printf output \"@noinline ρ_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim ρ\n @printf output \"@noinline U_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim U\n @printf output \"@noinline V_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim V\n @printf output \"@noinline W_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim W\n @printf output \"@noinline E_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim E\n @printf output \"@noinline Sρ_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim dρdt\n @printf output \"@noinline SU_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim dUdt\n @printf output \"@noinline SV_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim dVdt\n @printf output \"@noinline SW_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim dWdt\n @printf output \"@noinline SE_g(t, x, y, z, ::Val{%d}) = %s\\n\" dim dEdt\nend\n\nclose(output)\n" }, { "path": "test/Numerics/DGMethods/compressible_Navier_Stokes/mms_solution_generated.jl", "content": "const γ_exact = 1.4\nconst μ_exact = 1 // 100\n@noinline ρ_g(t, x, y, z, ::Val{2}) =\n sin(pi * x) * cos(pi * t) * cos(pi * y) + 3\n@noinline U_g(t, x, y, z, ::Val{2}) =\n (sin(pi * x) * cos(pi * t) * cos(pi * y) + 3) *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y)\n@noinline V_g(t, x, y, z, ::Val{2}) =\n (sin(pi * x) * cos(pi * t) * cos(pi * y) + 3) *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y)\n@noinline W_g(t, x, y, z, ::Val{2}) = 0\n@noinline E_g(t, x, y, z, ::Val{2}) =\n sin(pi * x) * cos(pi * t) * cos(pi * y) + 100\n@noinline Sρ_g(t, x, y, z, ::Val{2}) =\n pi * (\n 2 * sin(2 * pi * x) - 2 * sin(2 * pi * y) - sin(pi * (2 * t - 2 * x)) +\n sin(pi * (2 * t + 2 * x)) +\n sin(pi * (2 * t - 2 * y)) - sin(pi * (2 * t + 2 * y)) +\n 2 * sin(pi * (2 * x + 2 * y)) +\n sin(pi * (-2 * t + 2 * x + 2 * y)) +\n sin(pi * (2 * t + 2 * x + 2 * y)) +\n 10 * cos(pi * (-t + x + y)) - 2 * cos(pi * (t - x + y)) +\n 2 * cos(pi * (t + x - y)) +\n 14 * cos(pi * (t + x + y))\n ) / 8\n@noinline SU_g(t, x, y, z, ::Val{2}) =\n pi * (\n -2.0 * sin(pi * t) * sin(pi * x)^2 * cos(pi * t) * cos(pi * y)^2 -\n 3.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) -\n 3.0 * sin(pi * x)^3 * sin(pi * y) * cos(pi * t)^3 * cos(pi * y)^2 -\n 6.0 * sin(pi * x)^2 * sin(pi * y) * cos(pi * t)^2 * cos(pi * y) +\n 1.8 * sin(pi * x)^2 * cos(pi * t)^3 * cos(pi * x) * cos(pi * y)^3 +\n 3.6 * sin(pi * x) * cos(pi * t)^2 * cos(pi * x) * cos(pi * y)^2 +\n 0.0233333333333333 * pi * sin(pi * x) * cos(pi * t) * cos(pi * y) +\n 0.00333333333333333 * pi * sin(pi * y) * cos(pi * t) * cos(pi * x) +\n 0.4 * cos(pi * t) * cos(pi * x) * cos(pi * y)\n )\n@noinline SV_g(t, x, y, z, ::Val{2}) =\n pi * (\n -2.0 * sin(pi * t) * sin(pi * x)^2 * cos(pi * t) * cos(pi * y)^2 -\n 3.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) -\n 1.8 * sin(pi * x)^3 * sin(pi * y) * cos(pi * t)^3 * cos(pi * y)^2 -\n 3.6 * sin(pi * x)^2 * sin(pi * y) * cos(pi * t)^2 * cos(pi * y) +\n 3.0 * sin(pi * x)^2 * cos(pi * t)^3 * cos(pi * x) * cos(pi * y)^3 -\n 0.4 * sin(pi * x) * sin(pi * y) * cos(pi * t) +\n 6.0 * sin(pi * x) * cos(pi * t)^2 * cos(pi * x) * cos(pi * y)^2 +\n 0.0233333333333333 * pi * sin(pi * x) * cos(pi * t) * cos(pi * y) +\n 0.00333333333333333 * pi * sin(pi * y) * cos(pi * t) * cos(pi * x)\n )\n@noinline SW_g(t, x, y, z, ::Val{2}) = 0\n@noinline SE_g(t, x, y, z, ::Val{2}) =\n pi * (\n -1.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) +\n 1.6 * sin(pi * x)^4 * sin(pi * y) * cos(pi * t)^4 * cos(pi * y)^3 +\n 3.6 * sin(pi * x)^3 * sin(pi * y) * cos(pi * t)^3 * cos(pi * y)^2 -\n 1.6 * sin(pi * x)^3 * cos(pi * t)^4 * cos(pi * x) * cos(pi * y)^4 -\n 0.0233333333333333 *\n pi *\n sin(pi * x)^2 *\n sin(pi * y)^2 *\n cos(pi * t)^2 -\n 2.8 * sin(pi * x)^2 * sin(pi * y) * cos(pi * t)^2 * cos(pi * y) -\n 3.6 * sin(pi * x)^2 * cos(pi * t)^3 * cos(pi * x) * cos(pi * y)^3 +\n 0.0466666666666667 *\n pi *\n sin(pi * x)^2 *\n cos(pi * t)^2 *\n cos(pi * y)^2 +\n 0.0133333333333333 *\n pi *\n sin(pi * x) *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y) - 140.0 * sin(pi * x) * sin(pi * y) * cos(pi * t) +\n 2.8 * sin(pi * x) * cos(pi * t)^2 * cos(pi * x) * cos(pi * y)^2 -\n 0.0233333333333333 *\n pi *\n cos(pi * t)^2 *\n cos(pi * x)^2 *\n cos(pi * y)^2 + 140.0 * cos(pi * t) * cos(pi * x) * cos(pi * y)\n )\n@noinline ρ_g(t, x, y, z, ::Val{3}) =\n sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) + 3\n@noinline U_g(t, x, y, z, ::Val{3}) =\n (sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) + 3) *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y) *\n cos(pi * z)\n@noinline V_g(t, x, y, z, ::Val{3}) =\n (sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) + 3) *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y) *\n cos(pi * z)\n@noinline W_g(t, x, y, z, ::Val{3}) =\n (sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) + 3) *\n sin(pi * x) *\n sin(pi * z) *\n cos(pi * t) *\n cos(pi * y)\n@noinline E_g(t, x, y, z, ::Val{3}) =\n sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) + 100\n@noinline Sρ_g(t, x, y, z, ::Val{3}) =\n pi * (\n -sin(pi * t) * sin(pi * x) * cos(pi * y) * cos(pi * z) -\n 2 *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z)^2 -\n sin(pi * x)^2 * sin(pi * z)^2 * cos(pi * t)^2 * cos(pi * y)^2 +\n sin(pi * x)^2 * cos(pi * t)^2 * cos(pi * y)^2 * cos(pi * z)^2 -\n 3 * sin(pi * x) * sin(pi * y) * cos(pi * t) * cos(pi * z) +\n 2 *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 3 * sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) +\n 3 * cos(pi * t) * cos(pi * x) * cos(pi * y) * cos(pi * z)\n )\n@noinline SU_g(t, x, y, z, ::Val{3}) =\n pi * (\n -2.0 *\n sin(pi * t) *\n sin(pi * x)^2 *\n cos(pi * t) *\n cos(pi * y)^2 *\n cos(pi * z)^2 -\n 3.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) * cos(pi * z) -\n 3.0 *\n sin(pi * x)^3 *\n sin(pi * y) *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z)^3 -\n 2.0 *\n sin(pi * x)^3 *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * y)^3 *\n cos(pi * z) +\n 1.0 * sin(pi * x)^3 * cos(pi * t)^3 * cos(pi * y)^3 * cos(pi * z)^3 -\n 6.0 *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z)^2 -\n 0.6 *\n sin(pi * x)^2 *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z) -\n 3.0 * sin(pi * x)^2 * sin(pi * z)^2 * cos(pi * t)^2 * cos(pi * y)^2 +\n 1.8 *\n sin(pi * x)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z)^3 +\n 3.0 * sin(pi * x)^2 * cos(pi * t)^2 * cos(pi * y)^2 * cos(pi * z)^2 -\n 1.2 *\n sin(pi * x) *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 +\n 3.6 *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 0.0333333333333333 *\n pi *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y) *\n cos(pi * z) +\n 0.00333333333333333 *\n pi *\n sin(pi * y) *\n cos(pi * t) *\n cos(pi * x) *\n cos(pi * z) -\n 0.00333333333333333 *\n pi *\n cos(pi * t) *\n cos(pi * x) *\n cos(pi * y) *\n cos(pi * z) +\n 0.4 * cos(pi * t) * cos(pi * x) * cos(pi * y) * cos(pi * z)\n )\n@noinline SV_g(t, x, y, z, ::Val{3}) =\n pi * (\n -2.0 *\n sin(pi * t) *\n sin(pi * x)^2 *\n cos(pi * t) *\n cos(pi * y)^2 *\n cos(pi * z)^2 -\n 3.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) * cos(pi * z) +\n 0.6 *\n sin(pi * x)^3 *\n sin(pi * y) *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z) -\n 1.8 *\n sin(pi * x)^3 *\n sin(pi * y) *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z)^3 -\n 2.0 *\n sin(pi * x)^3 *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * y)^3 *\n cos(pi * z) +\n 1.0 * sin(pi * x)^3 * cos(pi * t)^3 * cos(pi * y)^3 * cos(pi * z)^3 +\n 1.2 *\n sin(pi * x)^2 *\n sin(pi * y) *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * y) -\n 3.6 *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z)^2 -\n 3.0 * sin(pi * x)^2 * sin(pi * z)^2 * cos(pi * t)^2 * cos(pi * y)^2 +\n 3.0 *\n sin(pi * x)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z)^3 +\n 3.0 * sin(pi * x)^2 * cos(pi * t)^2 * cos(pi * y)^2 * cos(pi * z)^2 -\n 0.4 * sin(pi * x) * sin(pi * y) * cos(pi * t) * cos(pi * z) +\n 0.00333333333333333 *\n pi *\n sin(pi * x) *\n sin(pi * y) *\n cos(pi * t) *\n cos(pi * z) +\n 6.0 *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 0.0333333333333333 *\n pi *\n sin(pi * x) *\n cos(pi * t) *\n cos(pi * y) *\n cos(pi * z) +\n 0.00333333333333333 *\n pi *\n sin(pi * y) *\n cos(pi * t) *\n cos(pi * x) *\n cos(pi * z)\n )\n@noinline SW_g(t, x, y, z, ::Val{3}) =\n pi *\n (\n -2.0 *\n sin(pi * t) *\n sin(pi * x)^2 *\n cos(pi * t) *\n cos(pi * y)^2 *\n cos(pi * z) - 3.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) -\n 3.0 *\n sin(pi * x)^3 *\n sin(pi * y) *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z)^2 -\n 0.8 * sin(pi * x)^3 * sin(pi * z)^2 * cos(pi * t)^3 * cos(pi * y)^3 +\n 2.8 * sin(pi * x)^3 * cos(pi * t)^3 * cos(pi * y)^3 * cos(pi * z)^2 -\n 6.0 *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z) +\n 3.0 *\n sin(pi * x)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z)^2 +\n 7.2 * sin(pi * x)^2 * cos(pi * t)^2 * cos(pi * y)^2 * cos(pi * z) -\n 0.00333333333333333 * pi * sin(pi * x) * sin(pi * y) * cos(pi * t) +\n 6.0 *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z) - 0.4 * sin(pi * x) * cos(pi * t) * cos(pi * y) +\n 0.0333333333333333 * pi * sin(pi * x) * cos(pi * t) * cos(pi * y) +\n 0.00333333333333333 * pi * cos(pi * t) * cos(pi * x) * cos(pi * y)\n ) *\n sin(pi * z)\n@noinline SE_g(t, x, y, z, ::Val{3}) =\n pi * (\n 0.001171875 *\n (1 - cos(2 * pi * x))^2 *\n (1 - cos(4 * pi * z)) *\n (cos(2 * pi * t) + 1)^2 *\n (cos(2 * pi * y) + 1)^2 -\n 1.0 * sin(pi * t) * sin(pi * x) * cos(pi * y) * cos(pi * z) +\n 0.8 *\n sin(pi * x)^4 *\n sin(pi * y) *\n sin(pi * z)^2 *\n cos(pi * t)^4 *\n cos(pi * y)^3 *\n cos(pi * z)^2 +\n 1.6 *\n sin(pi * x)^4 *\n sin(pi * y) *\n cos(pi * t)^4 *\n cos(pi * y)^3 *\n cos(pi * z)^4 +\n 0.2 * sin(pi * x)^4 * sin(pi * z)^4 * cos(pi * t)^4 * cos(pi * y)^4 -\n 0.4 * sin(pi * x)^4 * cos(pi * t)^4 * cos(pi * y)^4 * cos(pi * z)^4 +\n 1.8 *\n sin(pi * x)^3 *\n sin(pi * y) *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z) +\n 3.6 *\n sin(pi * x)^3 *\n sin(pi * y) *\n cos(pi * t)^3 *\n cos(pi * y)^2 *\n cos(pi * z)^3 -\n 0.8 *\n sin(pi * x)^3 *\n sin(pi * z)^2 *\n cos(pi * t)^4 *\n cos(pi * x) *\n cos(pi * y)^4 *\n cos(pi * z)^2 +\n 0.6 *\n sin(pi * x)^3 *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * y)^3 *\n cos(pi * z) -\n 1.6 *\n sin(pi * x)^3 *\n cos(pi * t)^4 *\n cos(pi * x) *\n cos(pi * y)^4 *\n cos(pi * z)^4 -\n 1.2 * sin(pi * x)^3 * cos(pi * t)^3 * cos(pi * y)^3 * cos(pi * z)^3 -\n 0.01 *\n pi *\n sin(pi * x)^2 *\n sin(pi * y)^2 *\n sin(pi * z)^2 *\n cos(pi * t)^2 -\n 0.0233333333333333 *\n pi *\n sin(pi * x)^2 *\n sin(pi * y)^2 *\n cos(pi * t)^2 *\n cos(pi * z)^2 -\n 0.0233333333333333 *\n pi *\n sin(pi * x)^2 *\n sin(pi * y) *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * y) -\n 2.8 *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z)^2 -\n 0.01 *\n pi *\n sin(pi * x)^2 *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * y) *\n cos(pi * z)^2 -\n 1.8 *\n sin(pi * x)^2 *\n sin(pi * z)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z) -\n 1.4 * sin(pi * x)^2 * sin(pi * z)^2 * cos(pi * t)^2 * cos(pi * y)^2 +\n 0.0133333333333333 *\n pi *\n sin(pi * x)^2 *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * y)^2 -\n 3.6 *\n sin(pi * x)^2 *\n cos(pi * t)^3 *\n cos(pi * x) *\n cos(pi * y)^3 *\n cos(pi * z)^3 +\n 0.0533333333333333 *\n pi *\n sin(pi * x)^2 *\n cos(pi * t)^2 *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 1.4 * sin(pi * x)^2 * cos(pi * t)^2 * cos(pi * y)^2 * cos(pi * z)^2 +\n 0.0133333333333333 *\n pi *\n sin(pi * x) *\n sin(pi * y) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y) *\n cos(pi * z)^2 -\n 140.0 * sin(pi * x) * sin(pi * y) * cos(pi * t) * cos(pi * z) +\n 0.0233333333333333 *\n pi *\n sin(pi * x) *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 +\n 0.01 *\n pi *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 2.8 *\n sin(pi * x) *\n cos(pi * t)^2 *\n cos(pi * x) *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 140.0 * sin(pi * x) * cos(pi * t) * cos(pi * y) * cos(pi * z) -\n 0.01 *\n pi *\n sin(pi * z)^2 *\n cos(pi * t)^2 *\n cos(pi * x)^2 *\n cos(pi * y)^2 -\n 0.0233333333333333 *\n pi *\n cos(pi * t)^2 *\n cos(pi * x)^2 *\n cos(pi * y)^2 *\n cos(pi * z)^2 +\n 140.0 * cos(pi * t) * cos(pi * x) * cos(pi * y) * cos(pi * z)\n )\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/plotting/bigfileofstuff.jl", "content": "using ClimateMachine\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Elements\nimport ClimateMachine.Mesh.Elements: baryweights\nusing ClimateMachine.Mesh.Grids: polynomialorders\nusing GaussQuadrature\nusing Base.Threads\n\n\n# Depending on CliMa version \n# old, should return a tuple of polynomial orders\n# polynomialorders(::DiscontinuousSpectralElementGrid{T, dim, N}) where {T, dim, N} = Tuple([N for i in 1:dim])\n# new, should return a tuple of polynomial orders\n# polynomialorders(::DiscontinuousSpectralElementGrid{T, dim, N}) where {T, dim, N} = N\n\n# utils.jl\n\"\"\"\nfunction cellaverage(Q; M = nothing)\n\n# Description\nCompute the cell-average of Q given the mass matrix M.\nAssumes that Q and M are the same size\n\n# Arguments\n`Q`: MPIStateArrays (array)\n\n# Keyword Arguments\n`M`: Mass Matrix (array)\n\n# Return\nThe cell-average of Q\n\"\"\"\nfunction cellaverage(Q; M = nothing)\n if M == nothing\n return nothing\n end\n return (sum(M .* Q, dims = 1) ./ sum(M, dims = 1))[:]\nend\n\n\"\"\"\nfunction coordinates(grid::DiscontinuousSpectralElementGrid)\n\n# Description\nGets the (x,y,z) coordinates corresponding to the grid\n\n# Arguments\n- `grid`: DiscontinuousSpectralElementGrid\n\n# Return\n- `x, y, z`: views of x, y, z coordinates\n\"\"\"\nfunction coordinates(grid::DiscontinuousSpectralElementGrid)\n x = view(grid.vgeo, :, grid.x1id, :) # x-direction\t\n y = view(grid.vgeo, :, grid.x2id, :) # y-direction\t\n z = view(grid.vgeo, :, grid.x3id, :) # z-direction\n return x, y, z\nend\n\n\"\"\"\nfunction cellcenters(Q; M = nothing)\n\n# Description\nGet the cell-centers of every element in the grid\n\n# Arguments\n- `grid`: DiscontinuousSpectralElementGrid\n\n# Return\n- Tuple of cell-centers\n\"\"\"\nfunction cellcenters(grid::DiscontinuousSpectralElementGrid)\n x, y, z = coordinates(grid)\n M = view(grid.vgeo, :, grid.Mid, :) # mass matrix\n xC = cellaverage(x, M = M)\n yC = cellaverage(y, M = M)\n zC = cellaverage(z, M = M)\n return xC[:], yC[:], zC[:]\nend\n\n\"\"\"\nfunction massmatrix(grid; M = nothing)\n\n# Description\nGet the mass matrix of the grid\n\n# Arguments\n- `grid`: DiscontinuousSpectralElementGrid\n\n# Return\n- Tuple of cell-centers\n\"\"\"\nfunction massmatrix(grid)\n return view(grid.vgeo, :, grid.Mid, :)\nend\n\n# find_element.jl\n# 3D version\nfunction findelement(xC, yC, zC, location, p, lin)\n ex, ey, ez = size(lin)\n # i \n currentmin = ones(1)\n minind = ones(Int64, 1)\n currentmin[1] = abs.(xC[p[lin[1, 1, 1]]] .- location[1])\n for i in 2:ex\n current = abs.(xC[p[lin[i, 1, 1]]] .- location[1])\n if current < currentmin[1]\n currentmin[1] = current\n minind[1] = i\n end\n end\n i = minind[1]\n # j \n currentmin[1] = abs.(yC[p[lin[1, 1, 1]]] .- location[2])\n minind[1] = 1\n for i in 2:ey\n current = abs.(yC[p[lin[1, i, 1]]] .- location[2])\n if current < currentmin[1]\n currentmin[1] = current\n minind[1] = i\n end\n end\n j = minind[1]\n # k \n currentmin[1] = abs.(zC[p[lin[1, 1, 1]]] .- location[3])\n minind[1] = 1\n for i in 2:ez\n current = abs.(zC[p[lin[1, 1, i]]] .- location[3])\n if current < currentmin[1]\n currentmin[1] = current\n minind[1] = i\n end\n end\n k = minind[1]\n return p[lin[i, j, k]]\nend\n\n# 2D version\nfunction findelement(xC, yC, location, p, lin)\n ex, ey = size(lin)\n # i \n currentmin = ones(1)\n minind = ones(Int64, 1)\n currentmin[1] = abs.(xC[p[lin[1, 1]]] .- location[1])\n for i in 2:ex\n current = abs.(xC[p[lin[i, 1]]] .- location[1])\n if current < currentmin[1]\n currentmin[1] = current\n minind[1] = i\n end\n end\n i = minind[1]\n # j \n currentmin[1] = abs.(yC[p[lin[1, 1]]] .- location[2])\n minind[1] = 1\n for i in 2:ey\n current = abs.(yC[p[lin[1, i]]] .- location[2])\n if current < currentmin[1]\n currentmin[1] = current\n minind[1] = i\n end\n end\n j = minind[1]\n return p[lin[i, j]]\nend\n\n# gridhelper.jl\n\nstruct InterpolationHelper{S, T}\n points::S\n quadrature::S\n interpolation::S\n cartesianindex::T\nend\n\nfunction InterpolationHelper(g::DiscontinuousSpectralElementGrid)\n porders = polynomialorders(g)\n if length(porders) == 3\n npx, npy, npz = porders\n rx, wx = GaussQuadrature.legendre(npx + 1, both)\n ωx = baryweights(rx)\n ry, wy = GaussQuadrature.legendre(npy + 1, both)\n ωy = baryweights(ry)\n rz, wz = GaussQuadrature.legendre(npz + 1, both)\n ωz = baryweights(rz)\n linlocal = reshape(\n collect(1:((npx + 1) * (npy + 1) * (npz + 1))),\n (npx + 1, npy + 1, npz + 1),\n )\n return InterpolationHelper(\n (rx, ry, rz),\n (wx, wy, wz),\n (ωx, ωy, ωz),\n linlocal,\n )\n elseif length(porders) == 2\n npx, npy = porders\n rx, wx = GaussQuadrature.legendre(npx + 1, both)\n ωx = baryweights(rx)\n ry, wy = GaussQuadrature.legendre(npy + 1, both)\n ωy = baryweights(ry)\n linlocal =\n reshape(collect(1:((npx + 1) * (npy + 1))), (npx + 1, npy + 1))\n return InterpolationHelper((rx, ry), (wx, wy), (ωx, ωy), linlocal)\n else\n println(\"Not supported\")\n return nothing\n end\n return nothing\nend\n\nstruct ElementHelper{S, T, U, Q, V, W}\n cellcenters::S\n coordinates::T\n cartesiansizes::U\n polynomialorders::Q\n permutation::V\n cartesianindex::W\nend\n\naddup(xC, tol) = sum(abs.(xC[1] .- xC) .≤ tol)\n\n# only valid for cartesian domains\nfunction ElementHelper(g::DiscontinuousSpectralElementGrid)\n porders = polynomialorders(g)\n x, y, z = coordinates(g)\n xC, yC, zC = cellcenters(g)\n ne = size(x)[2]\n ex = round(Int64, ne / addup(xC, 10^4 * eps(maximum(abs.(x)))))\n ey = round(Int64, ne / addup(yC, 10^4 * eps(maximum(abs.(y)))))\n ez = round(Int64, ne / addup(zC, 10^4 * eps(maximum(abs.(z)))))\n\n check = ne == ex * ey * ez\n check ? true : error(\"improper counting\")\n p = getperm(xC, yC, zC, ex, ey, ez)\n # should use dispatch ...\n if length(porders) == 3\n npx, npy, npz = porders\n lin = reshape(collect(1:length(xC)), (ex, ey, ez))\n return ElementHelper(\n (xC, yC, zC),\n (x, y, z),\n (ex, ey, ez),\n porders,\n p,\n lin,\n )\n elseif length(porders) == 2\n npx, npy = porders\n lin = reshape(collect(1:length(xC)), (ex, ey))\n check = ne == ex * ey\n check ? true : error(\"improper counting\")\n return ElementHelper((xC, yC), (x, y), (ex, ey), porders, p, lin)\n else\n println(\"no constructor for polynomial order = \", porders)\n return nothing\n end\n return nothing\nend\n\nstruct GridHelper{S, T, V}\n interpolation::S\n element::T\n grid::V\nend\n\nfunction GridHelper(g::DiscontinuousSpectralElementGrid)\n return GridHelper(InterpolationHelper(g), ElementHelper(g), g)\nend\n\nfunction getvalue(f, location, gridhelper::GridHelper)\n ih = gridhelper.interpolation\n eh = gridhelper.element\n porders = gridhelper.element.polynomialorders\n if length(porders) == 3\n npx, npy, npz = gridhelper.element.polynomialorders\n fl = reshape(f, (npx + 1, npy + 1, npz + 1, prod(eh.cartesiansizes)))\n ip = getvalue(\n fl,\n eh.cellcenters...,\n location,\n eh.permutation,\n eh.cartesianindex,\n ih.cartesianindex,\n eh.coordinates...,\n ih.points...,\n ih.interpolation...,\n )\n return ip\n elseif length(porders) == 2\n npx, npy = gridhelper.element.polynomialorders\n fl = reshape(f, (npx + 1, npy + 1, prod(eh.cartesiansizes)))\n ip = getvalue(\n fl,\n eh.cellcenters...,\n location,\n eh.permutation,\n eh.cartesianindex,\n ih.cartesianindex,\n eh.coordinates...,\n ih.points...,\n ih.interpolation...,\n )\n return ip\n end\n return nothing\nend\n\n# lagrange_interpolation.jl\nfunction checkgl(x, rx)\n for i in eachindex(rx)\n if abs(x - rx[i]) ≤ eps(rx[i])\n return i\n end\n end\n return 0\nend\n\nfunction lagrange_eval(f, newx, newy, newz, rx, ry, rz, ωx, ωy, ωz)\n icheck = checkgl(newx, rx)\n jcheck = checkgl(newy, ry)\n kcheck = checkgl(newz, rz)\n numerator = zeros(1)\n denominator = zeros(1)\n for k in eachindex(rz)\n if kcheck == 0\n Δz = (newz .- rz[k])\n polez = ωz[k] ./ Δz\n kk = k\n else\n polez = 1.0\n k = eachindex(rz)[end]\n kk = kcheck\n end\n for j in eachindex(ry)\n if jcheck == 0\n Δy = (newy .- ry[j])\n poley = ωy[j] ./ Δy\n jj = j\n else\n poley = 1.0\n j = eachindex(ry)[end]\n jj = jcheck\n end\n for i in eachindex(rx)\n if icheck == 0\n Δx = (newx .- rx[i])\n polex = ωx[i] ./ Δx\n ii = i\n else\n polex = 1.0\n i = eachindex(rx)[end]\n ii = icheck\n end\n numerator[1] += f[ii, jj, kk] * polex * poley * polez\n denominator[1] += polex * poley * polez\n end\n end\n end\n return numerator[1] / denominator[1]\nend\n\nfunction lagrange_eval(f, newx, newy, rx, ry, ωx, ωy)\n icheck = checkgl(newx, rx)\n jcheck = checkgl(newy, ry)\n numerator = zeros(1)\n denominator = zeros(1)\n for j in eachindex(ry)\n if jcheck == 0\n Δy = (newy .- ry[j])\n poley = ωy[j] ./ Δy\n jj = j\n else\n poley = 1.0\n j = eachindex(ry)[end]\n jj = jcheck\n end\n for i in eachindex(rx)\n if icheck == 0\n Δx = (newx .- rx[i])\n polex = ωx[i] ./ Δx\n ii = i\n else\n polex = 1.0\n i = eachindex(rx)[end]\n ii = icheck\n end\n numerator[1] += f[ii, jj] * polex * poley\n denominator[1] += polex * poley\n end\n end\n return numerator[1] / denominator[1]\nend\n\n\nfunction lagrange_eval_nocheck(f, newx, newy, newz, rx, ry, rz, ωx, ωy, ωz)\n numerator = zeros(1)\n denominator = zeros(1)\n for k in eachindex(rz)\n Δz = (newz .- rz[k])\n polez = ωz[k] ./ Δz\n for j in eachindex(ry)\n Δy = (newy .- ry[j])\n poley = ωy[j] ./ Δy\n for i in eachindex(rx)\n Δx = (newx .- rx[i])\n polex = ωx[i] ./ Δx\n numerator[1] += f[i, j, k] * polex * poley * polez\n denominator[1] += polex * poley * polez\n end\n end\n end\n return numerator[1] / denominator[1]\nend\n\nfunction lagrange_eval_nocheck(f, newx, newy, rx, ry, ωx, ωy)\n numerator = zeros(1)\n denominator = zeros(1)\n for j in eachindex(ry)\n Δy = (newy .- ry[j])\n poley = ωy[j] ./ Δy\n for i in eachindex(rx)\n Δx = (newx .- rx[i])\n polex = ωx[i] ./ Δx\n numerator[1] += f[i, j] * polex * poley\n denominator[1] += polex * poley\n end\n end\n return numerator[1] / denominator[1]\nend\n\nfunction lagrange_eval_nocheck(f, newx, rx, ωx)\n numerator = zeros(1)\n denominator = zeros(1)\n for i in eachindex(rx)\n Δx = (newx .- rx[i])\n polex = ωx[i] ./ Δx\n numerator[1] += f[i] * polex * poley\n denominator[1] += polex * poley\n end\n return numerator[1] / denominator[1]\nend\n\n# 3D, only valid for rectangles\nfunction getvalue(\n fl,\n xC,\n yC,\n zC,\n location,\n p,\n lin,\n linlocal,\n x,\n y,\n z,\n rx,\n ry,\n rz,\n ωx,\n ωy,\n ωz,\n)\n e = findelement(xC, yC, zC, location, p, lin)\n # need bounds to rescale, only value for cartesian\n\n xmax = x[linlocal[length(rx), 1, 1], e]\n xmin = x[linlocal[1, 1, 1], e]\n ymax = y[linlocal[1, length(ry), 1], e]\n ymin = y[linlocal[1, 1, 1], e]\n zmax = z[linlocal[1, 1, length(rz)], e]\n zmin = z[linlocal[1, 1, 1], e]\n\n # rescale new point to [-1,1]³\n newx = 2 * (location[1] - xmin) / (xmax - xmin) - 1\n newy = 2 * (location[2] - ymin) / (ymax - ymin) - 1\n newz = 2 * (location[3] - zmin) / (zmax - zmin) - 1\n\n return lagrange_eval(\n view(fl, :, :, :, e),\n newx,\n newy,\n newz,\n rx,\n ry,\n rz,\n ωx,\n ωy,\n ωz,\n )\nend\n\n# 2D\nfunction getvalue(fl, xC, yC, location, p, lin, linlocal, x, y, rx, ry, ωx, ωy)\n e = findelement(xC, yC, location, p, lin)\n # need bounds to rescale\n xmax = x[linlocal[length(rx), 1, 1], e]\n xmin = x[linlocal[1, 1, 1], e]\n ymax = y[linlocal[1, length(ry), 1], e]\n ymin = y[linlocal[1, 1, 1], e]\n\n # rescale new point to [-1,1]²\n newx = 2 * (location[1] - xmin) / (xmax - xmin) - 1\n newy = 2 * (location[2] - ymin) / (ymax - ymin) - 1\n\n return lagrange_eval(view(fl, :, :, e), newx, newy, rx, ry, ωx, ωy)\nend\n\n# permutations.jl\nfunction getperm(xC, yC, zC, ex, ey, ez)\n pz = sortperm(zC)\n tmpY = reshape(yC[pz], (ex * ey, ez))\n tmp_py = [sortperm(tmpY[:, i]) for i in 1:ez]\n py = zeros(Int64, length(pz))\n for i in eachindex(tmp_py)\n n = length(tmp_py[i])\n ii = (i - 1) * n + 1\n py[ii:(ii + n - 1)] .= tmp_py[i] .+ ii .- 1\n end\n tmpX = reshape(xC[pz][py], (ex, ey * ez))\n tmp_px = [sortperm(tmpX[:, i]) for i in 1:(ey * ez)]\n px = zeros(Int64, length(pz))\n for i in eachindex(tmp_px)\n n = length(tmp_px[i])\n ii = (i - 1) * n + 1\n px[ii:(ii + n - 1)] .= tmp_px[i] .+ ii .- 1\n end\n p = [pz[py[px[i]]] for i in eachindex(px)]\n return p\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/plotting/plot_output.jl", "content": "using JLD2\nusing GLMakie\n\nfilename = \"/Users/ballen/Projects/Clima/CLIMA/output/vtk_bickley_2D/bickley_jet.jld2\"\nDOF = 32\n\nf = jldopen(filename, \"r+\")\ninclude(\"bigfileofstuff.jl\")\ninclude(\"vizinanigans.jl\")\ninclude(\"ScalarFields.jl\")\n\ndg_grid = f[\"grid\"]\ngridhelper = GridHelper(dg_grid)\nx, y, z = coordinates(dg_grid)\nxC, yC, zC = cellcenters(dg_grid)\n\nϕ = ScalarField(copy(x), gridhelper)\n\nnewx = range(-2π, 2π, length = DOF)\nnewy = range(-2π, 2π, length = DOF)\nnorm(f[\"100\"])\n\n##\n\nQ = f[\"0\"]\ndof, nstates, nelems = size(Q)\n\nstates = Array{Float64, 3}[]\nstatenames = [\"ρ\", \"ρu\", \"ρv\", \"ρθ\"]\n\nfor i in 1:nstates\n state = zeros(length(newx), length(newy), 101)\n push!(states, state)\nend\n\nfor i in 0:100\n println(\"interpolating step \" * string(i))\n Qⁱ = f[string(i)]\n\n for j in 1:nstates\n ϕ .= Qⁱ[:, j, :]\n states[j][:, :, i + 1] .= ϕ(newx, newy, threads = true)\n end\nend\n\n# f[\"states\"] = states\nclose(f)\n\n##\n\nscene = volumeslice(states, statenames = statenames)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/plotting/vizinanigans.jl", "content": "using GLMakie, Statistics, Printf\n\n\"\"\"\nvisualize(states::AbstractArray; statenames = string.(1:length(states)), quantiles = (0.1, 0.99), aspect = (1,1,1), resolution = (1920, 1080), statistics = false, title = \"Field = \")\n# Description \nVisualize 3D states \n# Arguments\n- `states`: Array{Array{Float64,3},1}. An array of arrays containing different fields\n# Keyword Arguments\n- `statenames`: Array{String,1}. An array of stringnames\n- `aspect`: Tuple{Int64,Int64,Float64}. Determines aspect ratio of box for volumes\n- `resolution`: Resolution of preliminary makie window\n- `statistics`: boolean. toggle for displaying statistics \n# Return\n- `scene`: Scene. A preliminary scene object for manipulation\n\"\"\"\nfunction visualize(\n states::AbstractArray;\n statenames = string.(1:length(states)),\n units = [\"\" for i in eachindex(states)],\n aspect = (1, 1, 1),\n resolution = (1920, 1080),\n statistics = false,\n title = \"Field = \",\n bins = 300,\n)\n # Create scene\n scene, layout = layoutscene(resolution = resolution)\n lscene = layout[2:4, 2:4] = LScene(scene)\n width = round(Int, resolution[1] / 4) # make menu 1/4 of preliminary resolution\n\n # Create choices and nodes\n stateindex = collect(1:length(states))\n statenode = Node(stateindex[1])\n\n colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]\n colornode = Node(colorchoices[1])\n\n if statistics\n llscene =\n layout[4, 1] = Axis(\n scene,\n xlabel = @lift(statenames[$statenode] * units[$statenode]),\n xlabelcolor = :black,\n ylabel = \"pdf\",\n ylabelcolor = :black,\n xlabelsize = 40,\n ylabelsize = 40,\n xticklabelsize = 25,\n yticklabelsize = 25,\n xtickcolor = :black,\n ytickcolor = :black,\n xticklabelcolor = :black,\n yticklabelcolor = :black,\n )\n layout[3, 1] = Label(scene, \"Statistics\", width = width, textsize = 50)\n end\n\n # x,y,z are for determining the aspect ratio of the box\n if (typeof(aspect) <: Tuple) & (length(aspect) == 3)\n x, y, z = aspect\n else\n x, y, z = size(states[1])\n end\n\n # Clim sliders\n upperclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)\n upperclim_node = upperclim_slider.value\n lowerclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)\n lowerclim_node = lowerclim_slider.value\n\n # Lift Nodes\n state = @lift(states[$statenode])\n statename = @lift(statenames[$statenode])\n clims = @lift((\n quantile($state[:], $lowerclim_node),\n quantile($state[:], $upperclim_node),\n ))\n cmap_rgb = @lift(to_colormap($colornode))\n titlename = @lift(title * $statename) # use padding and appropriate centering\n\n # Statistics\n if statistics\n histogram_node = @lift(histogram($state, bins = bins))\n xs = @lift($histogram_node[1])\n ys = @lift($histogram_node[2])\n pdf = GLMakie.AbstractPlotting.barplot!(\n llscene,\n xs,\n ys,\n color = :red,\n strokecolor = :red,\n strokewidth = 1,\n )\n @lift(GLMakie.AbstractPlotting.xlims!(llscene, extrema($state)))\n @lift(GLMakie.AbstractPlotting.ylims!(\n llscene,\n extrema($histogram_node[2]),\n ))\n vlines!(\n llscene,\n @lift($clims[1]),\n color = :black,\n linewidth = width / 100,\n )\n vlines!(\n llscene,\n @lift($clims[2]),\n color = :black,\n linewidth = width / 100,\n )\n end\n\n # Volume Plot \n volume!(\n lscene,\n 0..x,\n 0..y,\n 0..z,\n state,\n camera = cam3d!,\n colormap = cmap_rgb,\n colorrange = clims,\n )\n # Camera\n cam = cameracontrols(scene.children[1])\n eyeposition = Float32[2, 2, 1.3]\n lookat = Float32[0.82, 0.82, 0.1]\n # Title\n supertitle =\n layout[1, 2:4] = Label(scene, titlename, textsize = 50, color = :black)\n\n\n # Menus\n statemenu = Menu(scene, options = zip(statenames, stateindex))\n on(statemenu.selection) do s\n statenode[] = s\n end\n\n colormenu = Menu(scene, options = zip(colorchoices, colorchoices))\n on(colormenu.selection) do s\n colornode[] = s\n end\n lowerclim_string = @lift(\n \"lower clim quantile = \" *\n @sprintf(\"%0.2f\", $lowerclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[1])\n )\n upperclim_string = @lift(\n \"upper clim quantile = \" *\n @sprintf(\"%0.2f\", $upperclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[2])\n )\n # depends on makie version, vbox for old, vgrid for new\n layout[2, 1] = vgrid!(\n Label(scene, \"State\", width = nothing),\n statemenu,\n Label(scene, \"Color\", width = nothing),\n colormenu,\n Label(scene, lowerclim_string, width = nothing),\n lowerclim_slider,\n Label(scene, upperclim_string, width = nothing),\n upperclim_slider,\n )\n layout[1, 1] = Label(scene, \"Menu\", width = width, textsize = 50)\n\n # Modify Axis\n axis = scene.children[1][OldAxis]\n # axis[:names][:axisnames] = (\"↓ Zonal [m] \", \"Meriodonal [m]↓ \", \"Depth [m]↓ \")\n axis[:names][:axisnames] = (\"↓\", \"↓ \", \"↓ \")\n axis[:names][:align] =\n ((:left, :center), (:right, :center), (:right, :center))\n # need to adjust size of ticks first and then size of axis names\n axis[:names][:textsize] = (50.0, 50.0, 50.0)\n axis[:ticks][:textsize] = (00.0, 00.0, 00.0)\n # axis[:ticks][:ranges_labels].val # current axis labels\n xticks = collect(range(-0, aspect[1], length = 2))\n yticks = collect(range(-0, aspect[2], length = 6))\n zticks = collect(range(-0, aspect[3], length = 2))\n ticks = (xticks, yticks, zticks)\n axis[:ticks][:ranges] = ticks\n xtickslabels = [@sprintf(\"%0.1f\", (xtick)) for xtick in xticks]\n xtickslabels[end] = \"1e6\"\n ytickslabels = [\"\", \"south\", \"\", \"\", \"north\", \"\"]\n ztickslabels = [@sprintf(\"%0.1f\", (xtick)) for xtick in xticks]\n labels = (xtickslabels, ytickslabels, ztickslabels)\n axis[:ticks][:labels] = labels\n\n display(scene)\n # Change the default camera position after the fact\n # note that these change dynamically as the plot is manipulated\n return scene\nend\n\n\n\"\"\"\nhistogram(array; bins = 100)\n# Description\nreturn arrays for plotting histogram\n\"\"\"\nfunction histogram(\n array;\n bins = minimum([100, length(array)]),\n normalize = true,\n)\n tmp = zeros(bins)\n down, up = extrema(array)\n down, up = down == up ? (down - 1, up + 1) : (down, up) # edge case\n bucket = collect(range(down, up, length = bins + 1))\n normalization = normalize ? length(array) : 1\n for i in eachindex(array)\n # normalize then multiply by bins\n val = (array[i] - down) / (up - down) * bins\n ind = ceil(Int, val)\n # handle edge cases\n ind = maximum([ind, 1])\n ind = minimum([ind, bins])\n tmp[ind] += 1 / normalization\n end\n return (bucket[2:end] + bucket[1:(end - 1)]) .* 0.5, tmp\nend\n\n# 2D visualization\nfunction visualize(\n states::Array{Array{S, 2}, 1};\n statenames = string.(1:length(states)),\n units = [\"\" for i in eachindex(states)],\n aspect = (1, 1, 1),\n resolution = (2412, 1158),\n title = \"Zonal and Temporal Average of \",\n xlims = (0, 1),\n ylims = (0, 1),\n bins = 300,\n) where {S}\n # Create scene\n scene, layout = layoutscene(resolution = resolution)\n lscene =\n layout[2:4, 2:4] = Axis(\n scene,\n xlabel = \"South to North [m]\",\n xlabelcolor = :black,\n ylabel = \"Depth [m]\",\n ylabelcolor = :black,\n xlabelsize = 40,\n ylabelsize = 40,\n xticklabelsize = 25,\n yticklabelsize = 25,\n xtickcolor = :black,\n ytickcolor = :black,\n xticklabelcolor = :black,\n yticklabelcolor = :black,\n titlesize = 50,\n )\n width = round(Int, resolution[1] / 4) # make menu 1/4 of preliminary resolution\n\n # Create choices and nodes\n stateindex = collect(1:length(states))\n statenode = Node(stateindex[1])\n\n colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]\n colornode = Node(colorchoices[1])\n\n interpolationlabels = [\"contour\", \"heatmap\"]\n interpolationchoices = [true, false]\n interpolationnode = Node(interpolationchoices[1])\n\n # Statistics\n llscene =\n layout[4, 1] = Axis(\n scene,\n xlabel = @lift(statenames[$statenode] * \" \" * units[$statenode]),\n xlabelcolor = :black,\n ylabel = \"pdf\",\n ylabelcolor = :black,\n xlabelsize = 40,\n ylabelsize = 40,\n xticklabelsize = 25,\n yticklabelsize = 25,\n xtickcolor = :black,\n ytickcolor = :black,\n xticklabelcolor = :black,\n yticklabelcolor = :black,\n )\n layout[3, 1] = Label(scene, \"Statistics\", width = width, textsize = 50)\n\n # Clim sliders\n upperclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)\n upperclim_node = upperclim_slider.value\n lowerclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)\n lowerclim_node = lowerclim_slider.value\n\n #ylims = @lift(range($lowerval, $upperval, length = $))\n # Lift Nodes\n state = @lift(states[$statenode])\n statename = @lift(statenames[$statenode])\n unit = @lift(units[$statenode])\n oclims = @lift((\n quantile($state[:], $lowerclim_node),\n quantile($state[:], $upperclim_node),\n ))\n cmap_rgb = @lift(\n $oclims[1] < $oclims[2] ? to_colormap($colornode) :\n reverse(to_colormap($colornode))\n )\n clims = @lift(\n $oclims[1] != $oclims[2] ? (minimum($oclims), maximum($oclims)) :\n (minimum($oclims) - 1, maximum($oclims) + 1)\n )\n xlims = Array(range(xlims[1], xlims[2], length = 4)) #collect(range(xlims[1], xlims[2], length = size(states[1])[1]))\n ylims = Array(range(ylims[1], ylims[2], length = 4)) #@lift(collect(range($lowerval], $upperval, length = size($state)[2])))\n # newrange = @lift(range($lowerval, $upperval, length = 4))\n # lscene.yticks = @lift(Array($newrange))\n titlename = @lift(title * $statename) # use padding and appropriate centering\n layout[1, 2:4] = Label(scene, titlename, textsize = 50)\n # heatmap \n heatmap1 = heatmap!(\n lscene,\n xlims,\n ylims,\n state,\n interpolate = interpolationnode,\n colormap = cmap_rgb,\n colorrange = clims,\n )\n\n\n # statistics\n histogram_node = @lift(histogram($state, bins = bins))\n xs = @lift($histogram_node[1])\n ys = @lift($histogram_node[2])\n pdf = GLMakie.AbstractPlotting.barplot!(\n llscene,\n xs,\n ys,\n color = :red,\n strokecolor = :red,\n strokewidth = 1,\n )\n @lift(GLMakie.AbstractPlotting.xlims!(llscene, extrema($state)))\n @lift(GLMakie.AbstractPlotting.ylims!(llscene, extrema($histogram_node[2])))\n vlines!(llscene, @lift($clims[1]), color = :black, linewidth = width / 100)\n vlines!(llscene, @lift($clims[2]), color = :black, linewidth = width / 100)\n\n # Menus\n statemenu = Menu(scene, options = zip(statenames, stateindex))\n on(statemenu.selection) do s\n statenode[] = s\n end\n\n colormenu = Menu(scene, options = zip(colorchoices, colorchoices))\n on(colormenu.selection) do s\n colornode[] = s\n end\n\n interpolationmenu =\n Menu(scene, options = zip(interpolationlabels, interpolationchoices))\n on(interpolationmenu.selection) do s\n interpolationnode[] = s\n heatmap1 = heatmap!(\n lscene,\n xlims,\n ylims,\n state,\n interpolate = s,\n colormap = cmap_rgb,\n colorrange = clims,\n )\n end\n\n newlabel = @lift($statename * \" \" * $unit)\n cbar = Colorbar(scene, heatmap1, label = newlabel)\n cbar.width = Relative(1 / 3)\n cbar.height = Relative(5 / 6)\n cbar.halign = :center\n cbar.flipaxisposition = true\n # cbar.labelpadding = -350\n cbar.labelsize = 50\n\n lowerclim_string = @lift(\n \"clim quantile = \" *\n @sprintf(\"%0.2f\", $lowerclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[1])\n )\n upperclim_string = @lift(\n \"clim quantile = \" *\n @sprintf(\"%0.2f\", $upperclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[2])\n )\n\n # depends on makie version, vbox for old, vgrid for new\n layout[2, 1] = vgrid!(\n Label(scene, \"State\", width = nothing),\n statemenu,\n Label(\n scene,\n \"plotting options\",\n width = width,\n textsize = 30,\n padding = (0, 0, 10, 0),\n ),\n interpolationmenu,\n Label(scene, \"Color\", width = nothing),\n colormenu,\n Label(scene, lowerclim_string, width = nothing),\n lowerclim_slider,\n Label(scene, upperclim_string, width = nothing),\n upperclim_slider,\n )\n\n layout[2:4, 5] = vgrid!(\n Label(\n scene,\n \"Color Bar\",\n width = width / 2,\n textsize = 50,\n padding = (25, 0, 0, 00),\n ),\n cbar,\n )\n layout[1, 1] = Label(scene, \"Menu\", width = width, textsize = 50)\n display(scene)\n return scene\nend\n\nfunction volumeslice(\n states::AbstractArray;\n statenames = string.(1:length(states)),\n units = [\"\" for i in eachindex(states)],\n aspect = (1, 1, 32 / 192),\n resolution = (2678, 1030),\n statistics = false,\n title = \"Volume plot of \",\n bins = 300,\n statlabelsize = (20, 20),\n)\n scene, layout = layoutscene(resolution = resolution)\n volumescene = layout[2:4, 2:4] = LScene(scene)\n menuwidth = round(Int, 350)\n layout[1, 1] = Label(scene, \"Menu\", width = menuwidth, textsize = 50)\n\n slice_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.0)\n slice_node = slice_slider.value\n\n directionindex = [1, 2, 3]\n directionnames = [\"x-slice\", \"y-slice\", \"z-slice\"]\n directionnode = Node(directionindex[1])\n\n stateindex = collect(1:length(states))\n statenode = Node(stateindex[1])\n\n layout[1, 2:4] =\n Label(scene, @lift(title * statenames[$statenode]), textsize = 50)\n\n colorchoices = [:balance, :thermal, :dense, :deep, :curl, :thermometer]\n colornode = Node(colorchoices[1])\n\n state = @lift(states[$statenode])\n statename = @lift(statenames[$statenode])\n unit = @lift(units[$statenode])\n nx = @lift(size($state)[1])\n ny = @lift(size($state)[2])\n nz = @lift(size($state)[3])\n nr = @lift([$nx, $ny, $nz])\n\n nslider = 100\n xrange = range(0.00, aspect[1], length = nslider)\n yrange = range(0.00, aspect[2], length = nslider)\n zrange = range(0.00, aspect[3], length = nslider)\n constx = collect(reshape(xrange, (nslider, 1, 1)))\n consty = collect(reshape(yrange, (1, nslider, 1)))\n constz = collect(reshape(zrange, (1, 1, nslider)))\n matx = zeros(nslider, nslider, nslider)\n maty = zeros(nslider, nslider, nslider)\n matz = zeros(nslider, nslider, nslider)\n matx .= constx\n maty .= consty\n matz .= constz\n sliceconst = [matx, maty, matz]\n planeslice = @lift(sliceconst[$directionnode])\n\n upperclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)\n upperclim_node = upperclim_slider.value\n lowerclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)\n lowerclim_node = lowerclim_slider.value\n\n clims = @lift((\n quantile($state[:], $lowerclim_node),\n quantile($state[:], $upperclim_node),\n ))\n\n volume!(\n volumescene,\n 0..aspect[1],\n 0..aspect[2],\n 0..aspect[3],\n state,\n overdraw = false,\n colorrange = clims,\n colormap = @lift(to_colormap($colornode)),\n )\n\n\n alpha_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.5)\n alphanode = alpha_slider.value\n\n slicecolormap = @lift(cgrad(:viridis, alpha = $alphanode))\n v = volume!(\n volumescene,\n 0..aspect[1],\n 0..aspect[2],\n 0..aspect[3],\n planeslice,\n algorithm = :iso,\n isorange = 0.005,\n isovalue = @lift($slice_node * aspect[$directionnode]),\n transparency = true,\n overdraw = false,\n visible = true,\n colormap = slicecolormap,\n colorrange = [-1, 0],\n )\n\n # Volume histogram\n\n layout[3, 1] = Label(scene, \"Statistics\", textsize = 50)\n hscene =\n layout[4, 1] = Axis(\n scene,\n xlabel = @lift(statenames[$statenode] * \" \" * units[$statenode]),\n xlabelcolor = :black,\n ylabel = \"pdf\",\n ylabelcolor = :black,\n xlabelsize = 40,\n ylabelsize = 40,\n xticklabelsize = statlabelsize[1],\n yticklabelsize = statlabelsize[2],\n xtickcolor = :black,\n ytickcolor = :black,\n xticklabelcolor = :black,\n yticklabelcolor = :black,\n )\n\n histogram_node = @lift(histogram($state, bins = bins))\n vxs = @lift($histogram_node[1])\n vys = @lift($histogram_node[2])\n pdf = GLMakie.AbstractPlotting.barplot!(\n hscene,\n vxs,\n vys,\n color = :red,\n strokecolor = :red,\n strokewidth = 1,\n )\n\n @lift(GLMakie.AbstractPlotting.xlims!(hscene, extrema($vxs)))\n @lift(GLMakie.AbstractPlotting.ylims!(hscene, extrema($vys)))\n vlines!(\n hscene,\n @lift($clims[1]),\n color = :black,\n linewidth = menuwidth / 100,\n )\n vlines!(\n hscene,\n @lift($clims[2]),\n color = :black,\n linewidth = menuwidth / 100,\n )\n\n\n # Slice\n sliceupperclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.99)\n sliceupperclim_node = sliceupperclim_slider.value\n slicelowerclim_slider =\n Slider(scene, range = range(0, 1, length = 101), startvalue = 0.01)\n slicelowerclim_node = slicelowerclim_slider.value\n\n\n slicexaxislabel = @lift([\"y\", \"x\", \"x\"][$directionnode])\n sliceyaxislabel = @lift([\"z\", \"z\", \"y\"][$directionnode])\n\n slicexaxis = @lift([[1, $ny], [1, $nx], [1, $nx]][$directionnode])\n sliceyaxis = @lift([[1, $nz], [1, $nz], [1, $ny]][$directionnode])\n\n slicescene =\n layout[2:4, 5:6] =\n Axis(scene, xlabel = slicexaxislabel, ylabel = sliceyaxislabel)\n\n sliced_state1 = @lift( $state[\n round(Int, 1 + $slice_node * (size($state)[1] - 1)),\n 1:size($state)[2],\n 1:size($state)[3],\n ])\n sliced_state2 = @lift( $state[\n 1:size($state)[1],\n round(Int, 1 + $slice_node * (size($state)[2] - 1)),\n 1:size($state)[3],\n ])\n sliced_state3 = @lift( $state[\n 1:size($state)[1],\n 1:size($state)[2],\n round(Int, 1 + $slice_node * (size($state)[3] - 1)),\n ])\n sliced_states = @lift([$sliced_state1, $sliced_state2, $sliced_state3])\n sliced_state = @lift($sliced_states[$directionnode])\n\n oclims = @lift((\n quantile($sliced_state[:], $slicelowerclim_node),\n quantile($sliced_state[:], $sliceupperclim_node),\n ))\n slicecolormapnode = @lift($oclims[1] < $oclims[2] ? $colornode : $colornode)\n sliceclims = @lift(\n $oclims[1] != $oclims[2] ? (minimum($oclims), maximum($oclims)) :\n (minimum($oclims) - 1, maximum($oclims) + 1)\n )\n\n heatmap1 = heatmap!(\n slicescene,\n slicexaxis,\n sliceyaxis,\n sliced_state,\n interpolate = true,\n colormap = slicecolormapnode,\n colorrange = sliceclims,\n )\n\n # Colorbar\n newlabel = @lift($statename * \" \" * $unit)\n cbar = Colorbar(scene, heatmap1, label = newlabel)\n cbar.width = Relative(1 / 3)\n # cbar.height = Relative(5/6)\n cbar.halign = :left\n # cbar.flipaxisposition = true\n # cbar.labelpadding = -250\n cbar.labelsize = 50\n\n @lift(GLMakie.AbstractPlotting.xlims!(slicescene, extrema($slicexaxis)))\n @lift(GLMakie.AbstractPlotting.ylims!(slicescene, extrema($sliceyaxis)))\n\n sliceindex = @lift([\n round(Int, 1 + $slice_node * ($nx - 1)),\n round(Int, 1 + $slice_node * ($ny - 1)),\n round(Int, 1 + $slice_node * ($nz - 1)),\n ][$directionnode])\n slicestring =\n @lift(directionnames[$directionnode] * \" of \" * statenames[$statenode])\n layout[1, 5:6] = Label(scene, slicestring, textsize = 50)\n\n\n axis = scene.children[1][OldAxis]\n axis[:names][:axisnames] = (\"↓\", \"↓ \", \"↓ \")\n axis[:names][:align] =\n ((:left, :center), (:right, :center), (:right, :center))\n axis[:names][:textsize] = (50.0, 50.0, 50.0)\n axis[:ticks][:textsize] = (00.0, 00.0, 00.0)\n\n\n # Menus\n statemenu = Menu(scene, options = zip(statenames, stateindex))\n on(statemenu.selection) do s\n statenode[] = s\n end\n\n colormenu = Menu(scene, options = zip(colorchoices, colorchoices))\n on(colormenu.selection) do s\n colornode[] = s\n end\n\n\n # Slice Statistics\n layout[1, 7] = Label(scene, \"Slice Menu\", width = menuwidth, textsize = 50)\n layout[3, 7] = Label(scene, \"Slice Statistics\", textsize = 50)\n hslicescene =\n layout[4, 7] = Axis(\n scene,\n xlabel = @lift(statenames[$statenode] * \" \" * units[$statenode]),\n xlabelcolor = :black,\n ylabel = \"pdf\",\n ylabelcolor = :black,\n xlabelsize = 40,\n ylabelsize = 40,\n xticklabelsize = statlabelsize[1],\n yticklabelsize = statlabelsize[2],\n xtickcolor = :black,\n ytickcolor = :black,\n xticklabelcolor = :black,\n yticklabelcolor = :black,\n )\n\n slicehistogram_node = @lift(histogram($sliced_state, bins = bins))\n xs = @lift($slicehistogram_node[1])\n ys = @lift($slicehistogram_node[2])\n pdf = GLMakie.AbstractPlotting.barplot!(\n hslicescene,\n xs,\n ys,\n color = :blue,\n strokecolor = :blue,\n strokewidth = 1,\n )\n\n @lift(GLMakie.AbstractPlotting.xlims!(hslicescene, extrema($xs)))\n @lift(GLMakie.AbstractPlotting.ylims!(hslicescene, extrema($ys)))\n vlines!(\n hslicescene,\n @lift($sliceclims[1]),\n color = :black,\n linewidth = menuwidth / 100,\n )\n vlines!(\n hslicescene,\n @lift($sliceclims[2]),\n color = :black,\n linewidth = menuwidth / 100,\n )\n\n interpolationnames = [\"contour\", \"heatmap\"]\n interpolationchoices = [true, false]\n interpolationnode = Node(interpolationchoices[1])\n interpolationmenu =\n Menu(scene, options = zip(interpolationnames, interpolationchoices))\n\n on(interpolationmenu.selection) do s\n interpolationnode[] = s\n # hack\n heatmap!(\n slicescene,\n slicexaxis,\n sliceyaxis,\n sliced_state,\n interpolate = s,\n colormap = slicecolormapnode,\n colorrange = sliceclims,\n )\n end\n\n directionmenu = Menu(scene, options = zip(directionnames, directionindex))\n\n on(directionmenu.selection) do s\n directionnode[] = s\n end\n\n slicemenustring = @lift(\n directionnames[$directionnode] *\n \" at index \" *\n string(round(Int, 1 + $slice_node * ($nr[$directionnode] - 1)))\n )\n lowerclim_string = @lift(\n \"quantile = \" *\n @sprintf(\"%0.2f\", $lowerclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[1])\n )\n upperclim_string = @lift(\n \"quantile = \" *\n @sprintf(\"%0.2f\", $upperclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $clims[2])\n )\n alphastring = @lift(\"Slice alpha = \" * @sprintf(\"%0.2f\", $alphanode))\n layout[2, 1] = vgrid!(\n Label(scene, \"State\", width = nothing),\n statemenu,\n Label(scene, \"Color\", width = nothing),\n colormenu,\n Label(scene, \"Slice Direction\", width = nothing),\n directionmenu,\n Label(scene, alphastring, width = nothing),\n alpha_slider,\n Label(scene, slicemenustring, width = nothing),\n slice_slider,\n Label(scene, lowerclim_string, width = nothing),\n lowerclim_slider,\n Label(scene, upperclim_string, width = nothing),\n upperclim_slider,\n )\n\n slicelowerclim_string = @lift(\n \"quantile = \" *\n @sprintf(\"%0.2f\", $slicelowerclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $sliceclims[1])\n )\n sliceupperclim_string = @lift(\n \"quantile = \" *\n @sprintf(\"%0.2f\", $sliceupperclim_node) *\n \", value = \" *\n @sprintf(\"%0.1e\", $sliceclims[2])\n )\n\n layout[2, 7] = vgrid!(\n Label(scene, \"Contour Plot Type\", width = nothing),\n interpolationmenu,\n Label(scene, slicelowerclim_string, width = nothing),\n slicelowerclim_slider,\n Label(scene, sliceupperclim_string, width = nothing),\n sliceupperclim_slider,\n cbar,\n )\n\n display(scene)\n return scene\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/FluidBC.jl", "content": "abstract type BoundaryCondition end\n\n\"\"\"\n FluidBC(momentum = Impenetrable(NoSlip())\n temperature = Insulating())\n\nThe standard boundary condition for CNSEModel. The default options imply a \"no flux\" boundary condition.\n\"\"\"\nBase.@kwdef struct FluidBC{M, T} <: BoundaryCondition\n momentum::M = Impenetrable(NoSlip())\n temperature::T = Insulating()\nend\n\nabstract type StateBC end\nabstract type MomentumBC <: StateBC end\nabstract type MomentumDragBC <: StateBC end\nabstract type TemperatureBC <: StateBC end\n\n(bc::StateBC)(state, aux, t) = bc.flux(bc.params, state, aux, t)\n\n\"\"\"\n Impenetrable(drag::MomentumDragBC) :: MomentumBC\n\nDefines an impenetrable wall model for momentum. This implies:\n - no flow in the direction normal to the boundary, and\n - flow parallel to the boundary is subject to the `drag` condition.\n\"\"\"\nstruct Impenetrable{D <: MomentumDragBC} <: MomentumBC\n drag::D\nend\n\n\"\"\"\n Penetrable(drag::MomentumDragBC) :: MomentumBC\n\nDefines an penetrable wall model for momentum. This implies:\n - no constraint on flow in the direction normal to the boundary, and\n - flow parallel to the boundary is subject to the `drag` condition.\n\"\"\"\nstruct Penetrable{D <: MomentumDragBC} <: MomentumBC\n drag::D\nend\n\n\"\"\"\n NoSlip() :: MomentumDragBC\n\nZero momentum at the boundary.\n\"\"\"\nstruct NoSlip <: MomentumDragBC end\n\n\"\"\"\n FreeSlip() :: MomentumDragBC\n\nNo surface drag on momentum parallel to the boundary.\n\"\"\"\nstruct FreeSlip <: MomentumDragBC end\n\n\"\"\"\n MomentumFlux(stress) :: MomentumDragBC\n\nApplies the specified kinematic stress on momentum normal to the boundary.\nPrescribe the net inward kinematic stress across the boundary by `stress`,\na function with signature `stress(problem, state, aux, t)`, returning the flux (in m²/s²).\n\"\"\"\nBase.@kwdef struct MomentumFlux{𝒯, 𝒫} <: MomentumDragBC\n flux::𝒯 = nothing\n params::𝒫 = nothing\nend\n\n\"\"\"\n Insulating() :: TemperatureBC\n\nNo temperature flux across the boundary\n\"\"\"\nstruct Insulating <: TemperatureBC end\n\n\"\"\"\n TemperatureFlux(flux) :: TemperatureBC\n\nPrescribe the net inward temperature flux across the boundary by `flux`,\na function with signature `flux(problem, state, aux, t)`, returning the flux (in m⋅K/s).\n\"\"\"\nBase.@kwdef struct TemperatureFlux{𝒯, 𝒫} <: TemperatureBC\n flux::𝒯 = nothing\n params::𝒫 = nothing\nend\n\nfunction check_bc(bcs, label)\n bctype = FluidBC\n\n bc_ρu = check_bc(bcs, Val(:ρu), label)\n bc_ρθ = check_bc(bcs, Val(:ρθ), label)\n\n return bctype(bc_ρu, bc_ρθ)\nend\n\nfunction check_bc(bcs, ::Val{:ρθ}, label)\n if haskey(bcs, :ρθ)\n if haskey(bcs[:ρθ], label)\n return bcs[:ρθ][label]\n end\n end\n\n return Insulating()\nend\n\nfunction check_bc(bcs, ::Val{:ρu}, label)\n if haskey(bcs, :ρu)\n if haskey(bcs[:ρu], label)\n return bcs[:ρu][label]\n end\n end\n\n return Impenetrable(FreeSlip())\nend\n\n# these functions just trim off the extra arguments\nfunction _cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n bc,\n model,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n t,\n _...,\n)\n return cnse_boundary_state!(nf, bc, model, state⁺, aux⁺, n, state⁻, aux⁻, t)\nend\n\nfunction _cnse_boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc,\n model,\n state⁺,\n gradflux⁺,\n hyperflux⁺,\n aux⁺,\n n,\n state⁻,\n gradflux⁻,\n hyperflux⁻,\n aux⁻,\n t,\n _...,\n)\n return cnse_boundary_state!(\n nf,\n bc,\n model,\n state⁺,\n gradflux⁺,\n aux⁺,\n n,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n )\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/ScalarFields.jl", "content": "using Base.Threads, LinearAlgebra\nimport Base: getindex, materialize!, broadcasted\n\nabstract type AbstractField end\nstruct ScalarField{S, T} <: AbstractField\n data::S\n grid::T\nend\n\nfunction (ϕ::ScalarField)(x::Tuple)\n return getvalue(ϕ.data, x, ϕ.grid)\nend\n\nfunction (ϕ::ScalarField)(x::Number, y::Number, z::Number)\n return getvalue(ϕ.data, (x, y, z), ϕ.grid)\nend\n\nfunction (ϕ::ScalarField)(x::Number, y::Number)\n return getvalue(ϕ.data, (x, y), ϕ.grid)\nend\n\ngetindex(ϕ::ScalarField, i::Int) = ϕ.data[i]\n\nmaterialize!(ϕ::ScalarField, f::Base.Broadcast.Broadcasted) =\n materialize!(ϕ.data, f)\nbroadcasted(identity, ϕ::ScalarField) = broadcasted(Base.identity, ϕ.data)\n\nfunction (ϕ::ScalarField)(\n xlist::StepRangeLen,\n ylist::StepRangeLen,\n zlist::StepRangeLen;\n threads = false,\n)\n newfield = zeros(length(xlist), length(ylist), length(zlist))\n if threads\n @threads for k in eachindex(zlist)\n for j in eachindex(ylist)\n for i in eachindex(xlist)\n newfield[i, j, k] =\n getvalue(ϕ.data, (xlist[i], ylist[j], zlist[k]), ϕ.grid)\n end\n end\n end\n else\n for k in eachindex(zlist)\n for j in eachindex(ylist)\n for i in eachindex(xlist)\n newfield[i, j, k] =\n getvalue(ϕ.data, (xlist[i], ylist[j], zlist[k]), ϕ.grid)\n end\n end\n end\n end\n return newfield\nend\n\nfunction (ϕ::ScalarField)(\n xlist::StepRangeLen,\n ylist::StepRangeLen;\n threads = false,\n)\n newfield = zeros(length(xlist), length(ylist))\n if threads\n @threads for j in eachindex(ylist)\n for i in eachindex(xlist)\n newfield[i, j] = getvalue(ϕ.data, (xlist[i], ylist[j]), ϕ.grid)\n end\n end\n else\n for j in eachindex(ylist)\n for i in eachindex(xlist)\n newfield[i, j] = getvalue(ϕ.data, (xlist[i], ylist[j]), ϕ.grid)\n end\n end\n end\n return newfield\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/VectorFields.jl", "content": "import Base: getindex\n\nabstract type AbstractRepresentation end\n\nstruct Cartesian <: AbstractRepresentation end\nstruct Spherical <: AbstractRepresentation end\nstruct Covariant <: AbstractRepresentation end\nstruct Contravariant <: AbstractRepresentation end\n\nBase.@kwdef struct VectorField{S, T, C} <: AbstractField\n data::S\n grid::T\n representation::C = Cartesian()\nend\n\nfunction VectorField(ϕ::VectorField; representation = Cartesian())\n return VectorField(\n data = ϕ.data,\n grid = ϕ.grid,\n representation = representation,\n )\nend\n\nfunction (ϕ::VectorField)(representation::AbstractRepresentation)\n return VectorField(ϕ, representation = representation)\nend\n\nfunction components(data, ::Cartesian)\n one = @sprintf(\"%0.2e\", data[1])\n two = @sprintf(\"%0.2e\", data[2])\n three = @sprintf(\"%0.2e\", data[3])\n println(one, \"x̂ +\", two, \"ŷ +\", three, \"ẑ \")\n return data\nend\n\ngetindex(ϕ::VectorField, ijk, e; verbose = true) = components(\n getindex.(ϕ.data, ijk, e),\n ϕ.grid,\n ijk,\n e,\n ϕ.representation,\n verbose = verbose,\n)\n\n## Component Grabber\n\nfunction convenience_print(grid, ijk, e)\n print(\"location: \")\n x, y, z = get_position(grid, ijk, e)\n println(\n \"x=\",\n @sprintf(\"%0.2e\", x),\n \",y=\",\n @sprintf(\"%0.2e\", y),\n \",z=\",\n @sprintf(\"%0.2e\", z),\n \" \",\n )\n print(\"field: \")\nend\n\nfunction components(v⃗, grid, ijk, e, ::Cartesian; verbose = true)\n if verbose\n one = @sprintf(\"%0.2e\", v⃗[1])\n two = @sprintf(\"%0.2e\", v⃗[2])\n three = @sprintf(\"%0.2e\", v⃗[3])\n convenience_print(grid, ijk, e)\n println(one, \"x̂ +\", two, \"ŷ +\", three, \"ẑ \")\n end\n return v⃗\nend\n\nfunction components(v⃗, grid, ijk, e, ::Covariant; verbose = true)\n v⃗¹, v⃗², v⃗³ = get_contravariant(grid, ijk, e)\n v₁ = dot(v⃗¹, v⃗)\n v₂ = dot(v⃗², v⃗)\n v₃ = dot(v⃗³, v⃗)\n if verbose\n one = @sprintf(\"%0.2e\", v₁)\n two = @sprintf(\"%0.2e\", v₂)\n three = @sprintf(\"%0.2e\", v₃)\n convenience_print(grid, ijk, e)\n println(one, \"v⃗¹ +\", two, \"v⃗² +\", three, \"v⃗³ \")\n end\n return (; v₁, v₂, v₃)\nend\n\nfunction components(v⃗, grid, ijk, e, ::Contravariant; verbose = true)\n v⃗₁, v⃗₂, v⃗₃ = get_covariant(grid, ijk, e)\n v¹ = dot(v⃗₁, v⃗)\n v² = dot(v⃗₂, v⃗)\n v³ = dot(v⃗₃, v⃗)\n if verbose\n one = @sprintf(\"%0.2e\", v¹)\n two = @sprintf(\"%0.2e\", v²)\n three = @sprintf(\"%0.2e\", v³)\n convenience_print(grid, ijk, e)\n println(one, \"v⃗₁ +\", two, \"v⃗₂ +\", three, \"v⃗₃ \")\n end\n return (; v¹, v², v³)\nend\n\nfunction components(v⃗, grid, ijk, e, ::Spherical; verbose = true)\n r̂, θ̂, φ̂ = get_spherical(grid, ijk, e)\n vʳ = dot(r̂, v⃗)\n vᶿ = dot(θ̂, v⃗)\n vᵠ = dot(φ̂, v⃗)\n if verbose\n one = @sprintf(\"%0.2e\", vʳ)\n two = @sprintf(\"%0.2e\", vᶿ)\n three = @sprintf(\"%0.2e\", vᵠ)\n convenience_print(grid, ijk, e)\n println(one, \"r̂ +\", two, \"θ̂ +\", three, \"φ̂ \")\n end\n return (; vʳ, vᶿ, vᵠ)\nend\n\n## Helper functions\nfunction get_jacobian(grid, ijk, e)\n ξ1x1 = grid.vgeo[ijk, grid.ξ1x1id, e]\n ξ1x2 = grid.vgeo[ijk, grid.ξ1x2id, e]\n ξ1x3 = grid.vgeo[ijk, grid.ξ1x3id, e]\n\n ξ2x1 = grid.vgeo[ijk, grid.ξ2x1id, e]\n ξ2x2 = grid.vgeo[ijk, grid.ξ2x2id, e]\n ξ2x3 = grid.vgeo[ijk, grid.ξ2x3id, e]\n\n ξ3x1 = grid.vgeo[ijk, grid.ξ3x1id, e]\n ξ3x2 = grid.vgeo[ijk, grid.ξ3x2id, e]\n ξ3x3 = grid.vgeo[ijk, grid.ξ3x3id, e]\n\n J = [\n ξ1x1 ξ1x2 ξ1x3\n ξ2x1 ξ2x2 ξ2x3\n ξ3x1 ξ3x2 ξ3x3\n ]\n # rows are the contravariant vectors\n # columns of the inverse are the covariant vectors\n return J\nend\n\nfunction get_contravariant(grid, ijk, e)\n J = get_jacobian(grid, ijk, e)\n a⃗¹ = J[1, :]\n a⃗² = J[2, :]\n a⃗³ = J[3, :]\n return (; a⃗¹, a⃗², a⃗³)\nend\n\nfunction get_covariant(grid, ijk, e)\n J = inv(get_jacobian(grid, ijk, e))\n a⃗₁ = J[:, 1]\n a⃗₂ = J[:, 2]\n a⃗₃ = J[:, 3]\n return (; a⃗₁, a⃗₂, a⃗₃)\nend\n\nfunction get_spherical(grid, ijk, e)\n x, y, z = get_position(grid, ijk, e)\n r̂ = [x, y, z] ./ norm([x, y, z])\n θ̂ = [x * z, y * z, -(x^2 + y^2)] ./ (norm([x, y, z]) * norm([x, y, 0]))\n φ̂ = [-y, x, 0] ./ norm([x, y, 0])\n return (; r̂, θ̂, φ̂)\nend\n\nfunction get_position(grid, ijk, e)\n x1 = grid.vgeo[ijk, grid.x1id, e]\n x2 = grid.vgeo[ijk, grid.x2id, e]\n x3 = grid.vgeo[ijk, grid.x3id, e]\n r = [x1 x2 x3]\n return r\nend\n\nfunction construct_determinant(grid)\n M = grid.vgeo[:, grid.Mid, :]\n\n ωx = reshape(grid.ω[1], (length(grid.ω[1]), 1, 1, 1))\n ωy = reshape(grid.ω[2], (1, length(grid.ω[2]), 1, 1))\n ωz = reshape(grid.ω[3], (1, 1, length(grid.ω[3]), 1))\n ω = reshape(ωx .* ωy .* ωz, (size(M)[1], 1))\n J = M ./ ω\n return J\nend\n\nfunction get_jacobian_determinant(grid, ijk, e)\n M = grid.vgeo[ijk, grid.Mid, e]\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/abstractions.jl", "content": "#######\n# useful concepts for dispatch\n#######\n\n\"\"\"\nAdvection terms\n\nright now really only non-linear or ::Nothing\n\"\"\"\nabstract type AdvectionTerm end\nstruct NonLinearAdvectionTerm <: AdvectionTerm end\n\n\"\"\"\nTurbulence Closures\n\nways to handle drag and diffusion and such\n\"\"\"\nabstract type TurbulenceClosure end\n\nstruct LinearDrag{T} <: TurbulenceClosure\n λ::T\nend\n\nstruct ConstantViscosity{T} <: TurbulenceClosure\n μ::T\n ν::T\n κ::T\n function ConstantViscosity{T}(;\n μ = T(1e-6), # m²/s\n ν = T(1e-6), # m²/s\n κ = T(1e-6), # m²/s\n ) where {T <: AbstractFloat}\n return new{T}(μ, ν, κ)\n end\nend\n\n\"\"\"\nForcings\n\nways to add body terms and sources\n\"\"\"\nabstract type Forcing end\nabstract type CoriolisForce <: Forcing end\n\nstruct fPlaneCoriolis{T} <: CoriolisForce\n fₒ::T\n β::T\n function fPlaneCoriolis{T}(;\n fₒ = T(1e-4), # Hz\n β = T(1e-11), # Hz/m\n ) where {T <: AbstractFloat}\n return new{T}(fₒ, β)\n end\nend\n\nstruct SphereCoriolis{T} <: CoriolisForce\n Ω::T\n function SphereCoriolis{T}(;\n Ω = T(2π / 86400), # Hz\n ) where {T <: AbstractFloat}\n return new{T}(Ω)\n end\nend\nstruct KinematicStress{T} <: Forcing\n τₒ::T\n function KinematicStress{T}(; τₒ = T(1e-4)) where {T <: AbstractFloat}\n return new{T}(τₒ)\n end\nend\n\nstruct Buoyancy{T} <: Forcing\n α::T # 1/K\n g::T # m/s²\n function Buoyancy{T}(; α = T(2e-4), g = T(10)) where {T <: AbstractFloat}\n return new{T}(α, g)\n end\nend\n\n\"\"\"\nGrouping structs\n\"\"\"\nabstract type AbstractModel end\n\nBase.@kwdef struct SpatialModel{𝒜, ℬ, 𝒞, 𝒟, ℰ, ℱ} <: AbstractModel\n balance_law::𝒜\n physics::ℬ\n numerics::𝒞\n grid::𝒟\n boundary_conditions::ℰ\n parameters::ℱ\nend\n\npolynomialorders(model::SpatialModel) = convention(\n model.grid.resolution.polynomial_order,\n Val(ndims(model.grid.domain)),\n)\n\nabstract type ModelPhysics end\n\nBase.@kwdef struct FluidPhysics{𝒪, 𝒜, 𝒟, 𝒞, ℬ} <: ModelPhysics\n orientation::𝒪 = ClimateMachine.Orientations.FlatOrientation()\n advection::𝒜 = NonLinearAdvectionTerm()\n dissipation::𝒟 = nothing\n coriolis::𝒞 = nothing\n buoyancy::ℬ = nothing\nend\n\nabstract type AbstractInitialValueProblem end\n\nBase.@kwdef struct InitialValueProblem{𝒫, ℐ𝒱} <: AbstractInitialValueProblem\n params::𝒫 = nothing\n initial_conditions::ℐ𝒱 = nothing\nend\n\nabstract type AbstractSimulation end\n\nstruct Simulation{𝒜, ℬ, 𝒞, 𝒟, ℰ, ℱ} <: AbstractSimulation\n model::𝒜\n state::ℬ\n timestepper::𝒞\n initial_conditions::𝒟\n callbacks::ℰ\n time::ℱ\nend\n\nfunction Simulation(;\n model = nothing,\n state = nothing,\n timestepper = nothing,\n initial_conditions = nothing,\n callbacks = nothing,\n time = nothing,\n)\n model = DGModel(model, initial_conditions = initial_conditions)\n\n FT = eltype(model.grid.vgeo)\n\n if state == nothing\n state = init_ode_state(model, FT(0); init_on_cpu = true)\n end\n # model = (discrete = dgmodel, spatial = model)\n return Simulation(\n model,\n state,\n timestepper,\n initial_conditions,\n callbacks,\n time,\n )\nend\n\ncoordinates(simulation::Simulation) = coordinates(simulation.model.grid)\npolynomialorders(simulation::Simulation) =\n polynomialorders(simulation.model.grid)\n\nabstract type AbstractTimestepper end\n\nBase.@kwdef struct TimeStepper{S, T} <: AbstractTimestepper\n method::S\n timestep::T\nend\n\n\"\"\"\ncalculate_dt(grid, wavespeed = nothing, diffusivity = nothing, viscocity = nothing, cfl = 0.1)\n\"\"\"\nfunction calculate_dt(\n grid;\n wavespeed = nothing,\n diffusivity = nothing,\n viscocity = nothing,\n cfl = 1.0,\n)\n Δx = min_node_distance(grid)\n Δts = []\n if wavespeed != nothing\n push!(Δts, Δx / wavespeed)\n end\n if diffusivity != nothing\n push!(Δts, Δx^2 / diffusivity)\n end\n if viscocity != nothing\n push!(Δts, Δx^2 / viscocity)\n end\n if Δts == []\n @error(\"Please provide characteristic speed or diffusivities\")\n return nothing\n end\n return cfl * minimum(Δts)\nend\n\nfunction calculate_dt(\n grid::DiscretizedDomain;\n wavespeed = nothing,\n diffusivity = nothing,\n viscocity = nothing,\n cfl = 1.0,\n)\n return calculate_dt(\n grid.numerical;\n wavespeed = wavespeed,\n diffusivity = diffusivity,\n viscocity = viscocity,\n cfl = cfl,\n )\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/boilerplate.jl", "content": "using MPI\nusing JLD2\nusing Test\nusing Dates\nusing Printf\nusing Logging\nusing StaticArrays\nusing LinearAlgebra\n\nusing ClimateMachine\nusing ClimateMachine.VTK\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\n\n# ×(a::SVector, b::SVector) = StaticArrays.cross(a, b)\n⋅(a::SVector, b::SVector) = StaticArrays.dot(a, b)\n⊗(a::SVector, b::SVector) = a * b'\n\nabstract type AbstractFluid <: BalanceLaw end\nstruct Fluid <: AbstractFluid end\n\ninclude(\"domains.jl\")\ninclude(\"grids.jl\")\ninclude(\"abstractions.jl\")\ninclude(\"callbacks.jl\")\ninclude(\"FluidBC.jl\")\ninclude(\"ScalarFields.jl\")\ninclude(\"VectorFields.jl\")\n# include(\"../plotting/bigfileofstuff.jl\")\n# include(\"../plotting/vizinanigans.jl\")\n\n\"\"\"\nfunction coordinates(grid::DiscontinuousSpectralElementGrid)\n# Description\nGets the (x,y,z) coordinates corresponding to the grid\n# Arguments\n- `grid`: DiscontinuousSpectralElementGrid\n# Return\n- `x, y, z`: views of x, y, z coordinates\n\"\"\"\n\nfunction evolve!(simulation, spatial_model; refDat = ())\n Q = simulation.state\n\n dg = simulation.model\n Ns = polynomialorders(spatial_model)\n\n if haskey(spatial_model.grid.resolution, :overintegration_order)\n Nover = spatial_model.grid.resolution.overintegration_order\n else\n Nover = (0, 0, 0)\n end\n\n # only works if Nover > 0\n overintegration_filter!(Q, dg, Ns, Nover)\n\n function custom_tendency(tendency, x...; kw...)\n dg(tendency, x...; kw...)\n overintegration_filter!(tendency, dg, Ns, Nover)\n end\n\n t0 = simulation.time.start\n Δt = simulation.timestepper.timestep\n timestepper = simulation.timestepper.method\n\n odesolver = timestepper(custom_tendency, Q, dt = Δt, t0 = t0)\n\n cbvector = create_callbacks(simulation, odesolver)\n\n if isempty(cbvector)\n solve!(Q, odesolver; timeend = simulation.time.finish)\n else\n solve!(\n Q,\n odesolver;\n timeend = simulation.time.finish,\n callbacks = cbvector,\n )\n end\n\n ## Check results against reference\n\n ClimateMachine.StateCheck.scprintref(cbvector[end])\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(cbvector[end], refDat)\n end\n\n return Q\nend\n\nfunction visualize(\n simulation::Simulation;\n statenames = [string(i) for i in 1:size(simulation.state)[2]],\n resolution = (32, 32, 32),\n)\n a_, statesize, b_ = size(simulation.state)\n mpistate = simulation.state\n grid = simulation.model.grid\n grid_helper = GridHelper(grid)\n r = coordinates(grid)\n states = []\n ϕ = ScalarField(copy(r[1]), grid_helper)\n r = uniform_grid(Ω, resolution = resolution)\n # statesymbol = vars(Q).names[i] # doesn't work for vectors\n for i in 1:statesize\n ϕ .= mpistate[:, i, :]\n ϕnew = ϕ(r...)\n push!(states, ϕnew)\n end\n visualize([states...], statenames = statenames)\nend\n\nfunction overintegration_filter!(state_array, dgmodel, Ns, Nover)\n if sum(Nover) > 0\n cutoff_order = Ns .+ 1\n\n cutoff = MassPreservingCutoffFilter(dgmodel.grid, cutoff_order)\n num_state_prognostic = number_states(dgmodel.balance_law, Prognostic())\n\n ClimateMachine.Mesh.Filters.apply!(\n state_array,\n 1:num_state_prognostic,\n dgmodel.grid,\n cutoff,\n )\n end\n\n return nothing\nend\n\n# initialized on CPU so not problem, but could do kernel launch?\nfunction set_ic!(ϕ, s::Number, _...)\n ϕ .= s\n return nothing\nend\n\nfunction set_ic!(ϕ, s::Function, p, x, y, z)\n _, nstates, _ = size(ϕ)\n @inbounds for i in eachindex(x)\n @inbounds for j in 1:nstates\n ϕʲ = view(ϕ, :, j, :)\n ϕʲ[i] = s(p, x[i], y[i], z[i])[j]\n end\n end\n return nothing\nend\n\n\n# filter below here\nusing ClimateMachine.Mesh.Filters\nusing KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\nimport ClimateMachine.Mesh.Filters: apply_async!\nimport ClimateMachine.Mesh.Filters: AbstractFilterTarget\nimport ClimateMachine.Mesh.Filters:\n number_state_filtered,\n vars_state_filtered,\n compute_filter_argument!,\n compute_filter_result!\n\nfunction modified_filter_matrix(r, Nc, σ)\n N = length(r) - 1\n T = eltype(r)\n\n @assert N >= 0\n @assert 0 <= Nc\n\n a, b = GaussQuadrature.legendre_coefs(T, N)\n V = (N == 0 ? ones(T, 1, 1) : GaussQuadrature.orthonormal_poly(r, a, b))\n\n Σ = ones(T, N + 1)\n if Nc ≤ N\n Σ[(Nc:N) .+ 1] .= σ.(((Nc:N) .- Nc) ./ (N - Nc))\n end\n\n V * Diagonal(Σ) / V\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/callbacks.jl", "content": "abstract type AbstractCallback end\n\nstruct Default <: AbstractCallback end\nstruct Info <: AbstractCallback end\nstruct StateCheck{T} <: AbstractCallback\n number_of_checks::T\nend\n\nBase.@kwdef struct JLD2State{T, V, B} <: AbstractCallback\n iteration::T\n filepath::V\n overwrite::B = true\nend\n\nBase.@kwdef struct VTKState{T, V, C, B} <: AbstractCallback\n iteration::T = 1\n filepath::V = \".\"\n counter::C = [0]\n overwrite::B = true\nend\n\nfunction create_callbacks(simulation::Simulation, odesolver)\n callbacks = simulation.callbacks\n\n if isempty(callbacks)\n return ()\n else\n cbvector = [\n create_callback(callback, simulation, odesolver)\n for callback in callbacks\n ]\n return tuple(cbvector...)\n end\nend\n\nfunction create_callback(::Default, simulation::Simulation, odesolver)\n cb_info = create_callback(Info(), simulation, odesolver)\n cb_state_check = create_callback(StateCheck(10), simulation, odesolver)\n\n return (cb_info, cb_state_check)\nend\n\nfunction create_callback(::Info, simulation::Simulation, odesolver)\n Q = simulation.state\n timeend = simulation.time.finish\n mpicomm = MPI.COMM_WORLD\n\n starttime = Ref(now())\n cbinfo = ClimateMachine.GenericCallbacks.EveryXWallTimeSeconds(\n 60,\n mpicomm,\n ) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.4f / %8.4f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ClimateMachine.ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n\n if isnan(energy)\n error(\"NaNs\")\n end\n end\n end\n\n return cbinfo\nend\n\nfunction create_callback(callback::StateCheck, simulation::Simulation, _...)\n sim_length = simulation.time.finish - simulation.time.start\n timestep = simulation.timestepper.timestep\n nChecks = callback.number_of_checks\n\n\n nt_freq = floor(Int, sim_length / timestep / nChecks)\n\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [(simulation.state, \"state\")],\n nt_freq,\n )\n\n return cbcs_dg\nend\n\nfunction create_callback(output::JLD2State, simulation::Simulation, odesolver)\n # Initialize output\n output.overwrite &&\n isfile(output.filepath) &&\n rm(output.filepath; force = output.overwrite)\n\n Q = simulation.state\n mpicomm = MPI.COMM_WORLD\n iteration = output.iteration\n\n steps = ClimateMachine.ODESolvers.getsteps(odesolver)\n time = ClimateMachine.ODESolvers.gettime(odesolver)\n\n file = jldopen(output.filepath, \"a+\")\n JLD2.Group(file, \"state\")\n JLD2.Group(file, \"time\")\n file[\"state\"][string(steps)] = Array(Q)\n file[\"time\"][string(steps)] = time\n close(file)\n\n\n jldcallback = ClimateMachine.GenericCallbacks.EveryXSimulationSteps(\n iteration,\n ) do (s = false)\n steps = ClimateMachine.ODESolvers.getsteps(odesolver)\n time = ClimateMachine.ODESolvers.gettime(odesolver)\n @info steps, time\n file = jldopen(output.filepath, \"a+\")\n file[\"state\"][string(steps)] = Array(Q)\n file[\"time\"][string(steps)] = time\n close(file)\n return nothing\n end\nend\n\nfunction create_callback(output::VTKState, simulation::Simulation, odesolver)\n # Initialize output\n output.overwrite &&\n isfile(output.filepath) &&\n rm(output.filepath; force = output.overwrite)\n mkpath(output.filepath)\n\n state = simulation.state\n model = simulation.model\n\n function do_output(counter, model, state)\n mpicomm = MPI.COMM_WORLD\n balance_law = model.balance_law\n aux_state = model.state_auxiliary\n\n outprefix = @sprintf(\n \"%s/mpirank%04d_step%04d\",\n output.filepath,\n MPI.Comm_rank(mpicomm),\n counter[1],\n )\n\n @info \"doing VTK output\" outprefix\n\n state_names =\n flattenednames(vars_state(balance_law, Prognostic(), eltype(state)))\n aux_names =\n flattenednames(vars_state(balance_law, Auxiliary(), eltype(state)))\n\n writevtk(outprefix, state, model, state_names, aux_state, aux_names)\n\n counter[1] += 1\n\n return nothing\n end\n\n do_output(output.counter, model, state)\n cbvtk =\n ClimateMachine.GenericCallbacks.EveryXSimulationSteps(output.iteration) do (\n init = false\n )\n do_output(output.counter, model, state)\n return nothing\n end\n\n return cbvtk\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/domains.jl", "content": "import Base: getindex, *, ndims, length, ^\nimport LinearAlgebra: ×\n\nabstract type AbstractDomain end\nabstract type AbstractBoundary end\n\nstruct DomainBoundary <: AbstractBoundary\n closure::Any\nend\n\nstruct PointDomain{S} <: AbstractDomain\n point::S\nend\n\nstruct IntervalDomain{AT, BT, PT} <: AbstractDomain\n min::AT\n max::BT\n periodic::PT\nend\n\nfunction IntervalDomain(min, max; periodic = false)\n @assert min < max\n return IntervalDomain(min, max, periodic)\nend\n\nfunction Periodic(min, max)\n @assert min < max\n return IntervalDomain(min, max, periodic = true)\nend\n\nS¹ = Periodic\n\nfunction Interval(min, max)\n @assert min < max\n return IntervalDomain(min, max)\nend\n\nfunction Periodic(min, max)\n @assert min < max\n return IntervalDomain(min, max, periodic = true)\nend\n\nfunction Base.show(io::IO, Ω::IntervalDomain)\n min = Ω.min\n max = Ω.max\n printstyled(io, \"[\", color = 226)\n astring = @sprintf(\"%0.2f\", min)\n bstring = @sprintf(\"%0.2f\", max)\n printstyled(astring, \", \", bstring, color = 7)\n # printstyled(\"$min, $max\", color = 7)\n Ω.periodic ? printstyled(io, \")\", color = 226) :\n printstyled(io, \"]\", color = 226)\nend\n\nfunction Base.show(io::IO, o::PointDomain)\n printstyled(\"{\", o.point, \"}\", color = 201)\nend\n\n# Product Domains\nstruct ProductDomain{DT} <: AbstractDomain\n domains::DT\nend\n\nfunction Base.show(io::IO, Ω::ProductDomain)\n for (i, domain) in enumerate(Ω.domains)\n print(domain)\n if i != length(Ω.domains)\n printstyled(io, \"×\", color = 118)\n end\n end\nend\n\nndims(p::PointDomain) = 0\nndims(Ω::IntervalDomain) = 1\nndims(Ω::ProductDomain) = +(ndims.(Ω.domains)...)\n\nlength(Ω::IntervalDomain) = Ω.max - Ω.min\nlength(Ω::ProductDomain) = length.(Ω.domains)\n\n×(arg1::AbstractDomain, arg2::AbstractDomain) = ProductDomain((arg1, arg2))\n×(args::ProductDomain, arg2::AbstractDomain) =\n ProductDomain((args.domains..., arg2))\n×(arg1::AbstractDomain, args::ProductDomain) =\n ProductDomain((arg1, args.domains...))\n×(arg1::ProductDomain, args::ProductDomain) =\n ProductDomain((arg1.domains..., args.domains...))\n×(args::AbstractDomain) = ProductDomain(args...)\n*(arg1::AbstractDomain, arg2::AbstractDomain) = arg1 × arg2\n\nfunction ^(Ω::IntervalDomain, T::Int)\n Ωᵀ = Ω\n for i in 1:(T - 1)\n Ωᵀ *= Ω\n end\n return Ωᵀ\nend\n\nfunction info(Ω::ProductDomain)\n println(\"This is a \", ndims(Ω), \"-dimensional tensor product domain.\")\n print(\"The domain is \")\n println(Ω, \".\")\n for (i, domain) in enumerate(Ω.domains)\n domain_string = domain.periodic ? \"periodic\" : \"wall-bounded\"\n length = @sprintf(\"%.2f \", domain.max - domain.min)\n println(\n \"The dimension $i domain is \",\n domain_string,\n \" with length ≈ \",\n length,\n )\n end\n return nothing\nend\n\nfunction isperiodic(Ω::ProductDomain)\n max = [Ω.domains[i].periodic for i in eachindex(Ω.domains)]\n return prod(max)\nend\n\nfunction periodicityof(Ω::ProductDomain)\n periodicity = ones(Bool, ndims(Ω))\n for i in 1:ndims(Ω)\n periodicity[i] = Ω[i].periodic\n end\n return Tuple(periodicity)\nend\n\ngetindex(Ω::ProductDomain, i::Int) = Ω.domains[i]\n\n# Boundaries\nstruct Boundaries{S}\n boundaries::S\nend\n\ngetindex(∂Ω::Boundaries, i) = ∂Ω.boundaries[i]\n\nfunction Base.show(io::IO, ∂Ω::Boundaries)\n for (i, boundary) in enumerate(∂Ω.boundaries)\n printstyled(\"boundary \", i, \": \", color = 13)\n println(boundary)\n end\nend\n\nfunction ∂(Ω::IntervalDomain)\n if Ω.periodic\n return (nothing)\n else\n return Boundaries((PointDomain(Ω.min), PointDomain(Ω.max)))\n end\n return nothing\nend\n\nfunction ∂(Ω::ProductDomain)\n ∂Ω = []\n for domain in Ω.domains\n push!(∂Ω, ∂(domain))\n end\n splitb = []\n for (i, boundary) in enumerate(∂Ω)\n tmp = Any[]\n push!(tmp, Ω.domains...)\n if boundary != nothing\n tmp1 = copy(tmp)\n tmp2 = copy(tmp)\n tmp1[i] = boundary[1]\n push!(splitb, ProductDomain(Tuple(tmp1)))\n tmp2[i] = boundary[2]\n push!(splitb, ProductDomain(Tuple(tmp2)))\n end\n end\n return Boundaries(Tuple(splitb))\nend\n\nBase.@kwdef struct SphericalShellDomain{ℛ, ℋ, 𝒟} <: AbstractDomain\n radius::ℛ = 6378e3\n height::ℋ = 30e3\n depth::𝒟 = 3e3\nend\n\n\nlength(Ω::SphericalShellDomain) = (Ω.radius, Ω.radius, Ω.height - Ω.depth)\nndims(Ω::SphericalShellDomain) = 3\n\nfunction AtmosDomain(; radius = 6378e3, height = 30e3)\n depth = 0\n return SphericalShellDomain(; radius, height, depth)\nend\n\nfunction OceanDomain(; radius = 6378e3, depth = 3e3)\n height = 0\n return SphericalShellDomain(; radius, height, depth)\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/shared_source/grids.jl", "content": "import ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid\n\nfunction coordinates(grid::DiscontinuousSpectralElementGrid)\n x = view(grid.vgeo, :, grid.x1id, :) # x-direction\t\n y = view(grid.vgeo, :, grid.x2id, :) # y-direction\t\n z = view(grid.vgeo, :, grid.x3id, :) # z-direction\n return x, y, z\nend\n\n# some convenience functions\nfunction convention(\n a::NamedTuple{(:vertical, :horizontal), T},\n ::Val{3},\n) where {T}\n return (a.horizontal, a.horizontal, a.vertical)\nend\n\nfunction convention(a::Number, ::Val{3})\n return (a, a, a)\nend\n\nfunction convention(\n a::NamedTuple{(:vertical, :horizontal), T},\n ::Val{2},\n) where {T}\n return (a.horizontal, a.vertical)\nend\n\nfunction convention(a::Number, ::Val{2})\n return (a, a)\nend\n\nfunction convention(a::Tuple, b)\n return a\nend\n\n# brick range brickbuilder\nfunction uniform_brick_builder(domain, elements; FT = Float64)\n dimension = ndims(domain)\n\n tuple_ranges = []\n for i in 1:dimension\n push!(\n tuple_ranges,\n range(\n FT(domain[i].min);\n length = elements[i] + 1,\n stop = FT(domain[i].max),\n ),\n )\n end\n\n brickrange = Tuple(tuple_ranges)\n return brickrange\nend\n\n# Grid Constructor\n\"\"\"\nfunction DiscontinuousSpectralElementGrid(domain::ProductDomain; elements = nothing, polynomialorder = nothing)\n# Description \nComputes a DiscontinuousSpectralElementGrid as specified by a product domain\n# Arguments\n-`domain`: A product domain object\n# Keyword Arguments \n-`elements`: A tuple of integers ordered by (Nx, Ny, Nz) for number of elements\n-`polynomialorder`: A tupe of integers ordered by (npx, npy, npz) for polynomial order\n-`FT`: floattype, assumed Float64 unless otherwise specified\n-`topology`: default = StackedBrickTopology\n-`mpicomm`: default = MPI.COMM_WORLD\n-`array`: default = ClimateMachine.array_type()\n-`brickbuilder`: default = uniform_brick_builder, \n brickrange=uniform_brick_builder(domain, elements)\n# Return \nA DiscontinuousSpectralElementGrid object\n\"\"\"\nfunction DiscontinuousSpectralElementGrid(\n domain::ProductDomain;\n elements = nothing,\n polynomialorder = nothing,\n FT = Float64,\n mpicomm = MPI.COMM_WORLD,\n array = ClimateMachine.array_type(),\n topology = StackedBrickTopology,\n brick_builder = uniform_brick_builder,\n)\n\n if elements == nothing\n error_message = \"Please specify the number of elements as a tuple whose size is commensurate with the domain,\"\n error_message *= \" e.g., a 3 dimensional domain would need a specification like elements = (10,10,10).\"\n error_message *= \" or elements = (vertical = 8, horizontal = 5)\"\n\n @error(error_message)\n return nothing\n end\n\n if polynomialorder == nothing\n error_message = \"Please specify the polynomial order as a tuple whose size is commensurate with the domain,\"\n error_message *= \"e.g., a 3 dimensional domain would need a specification like polynomialorder = (3,3,3).\"\n error_message *= \" or polynomialorder = (vertical = 8, horizontal = 5)\"\n\n @error(error_message)\n return nothing\n end\n\n dimension = ndims(domain)\n\n if (dimension < 2) || (dimension > 3)\n error_message = \"SpectralElementGrid only works with dimensions 2 or 3. \"\n error_message *= \"The current dimension is \" * string(ndims(domain))\n\n println(\"The domain is \", domain)\n @error(error_message)\n return nothing\n end\n\n elements = convention(elements, Val(dimension))\n if ndims(domain) != length(elements)\n @error(\"Incorrectly specified elements for the dimension of the domain\")\n return nothing\n end\n\n polynomialorder = convention(polynomialorder, Val(dimension))\n if ndims(domain) != length(polynomialorder)\n @error(\"Incorrectly specified polynomialorders for the dimension of the domain\")\n return nothing\n end\n\n brickrange = brick_builder(domain, elements, FT = FT)\n\n if dimension == 2\n boundary = ((1, 2), (3, 4))\n else\n boundary = ((1, 2), (3, 4), (5, 6))\n end\n\n periodicity = periodicityof(domain)\n connectivity = dimension == 2 ? :face : :full\n\n topl = topology(\n mpicomm,\n brickrange;\n periodicity = periodicity,\n boundary = boundary,\n connectivity = connectivity,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = array,\n polynomialorder = polynomialorder,\n )\n\n return grid\nend\n\nabstract type AbstractDiscretizedDomain end\n\nstruct DiscretizedDomain{𝒜, ℬ, 𝒞} <: AbstractDiscretizedDomain\n domain::𝒜\n resolution::ℬ\n numerical::𝒞\nend\n\nfunction DiscretizedDomain(\n domain::ProductDomain;\n elements = nothing,\n polynomial_order = nothing,\n overintegration_order = nothing,\n FT = Float64,\n mpicomm = MPI.COMM_WORLD,\n array = ClimateMachine.array_type(),\n topology = StackedBrickTopology,\n brick_builder = uniform_brick_builder,\n)\n\n grid = DiscontinuousSpectralElementGrid(\n domain,\n elements = elements,\n polynomialorder = polynomial_order .+ overintegration_order,\n FT = FT,\n mpicomm = mpicomm,\n array = array,\n topology = topology,\n brick_builder = brick_builder,\n )\n return DiscretizedDomain(\n domain,\n (; elements, polynomial_order, overintegration_order),\n grid,\n )\nend\n\n## SphericalShellDomain\n\nfunction DiscontinuousSpectralElementGrid(\n Ω::SphericalShellDomain;\n elements = (vertical = 2, horizontal = 4),\n polynomialorder = (vertical = 4, horizontal = 4),\n mpicomm = MPI.COMM_WORLD,\n boundary = (5, 6),\n FT = Float64,\n array = Array,\n)\n Rrange = grid1d(\n Ω.radius - Ω.depth,\n Ω.radius + Ω.height,\n nelem = elements.vertical,\n )\n\n topl = StackedCubedSphereTopology(\n mpicomm,\n elements.horizontal,\n Rrange,\n boundary = boundary,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = array,\n polynomialorder = (\n polynomialorder.vertical,\n polynomialorder.horizontal,\n ),\n meshwarp = equiangular_cubed_sphere_warp,\n )\n return grid\nend\n\nfunction DiscretizedDomain(\n domain::SphericalShellDomain;\n elements = nothing,\n polynomial_order = nothing,\n overintegration_order = nothing,\n FT = Float64,\n mpicomm = MPI.COMM_WORLD,\n array = ClimateMachine.array_type(),\n topology = StackedBrickTopology,\n brick_builder = uniform_brick_builder,\n)\n new_polynomial_order = convention(polynomial_order, Val(2))\n new_polynomial_order =\n new_polynomial_order .+ convention(overintegration_order, Val(2))\n vertical, horizontal = new_polynomial_order\n grid = DiscontinuousSpectralElementGrid(\n domain,\n elements = elements,\n polynomialorder = (; vertical, horizontal),\n FT = FT,\n mpicomm = mpicomm,\n array = array,\n )\n return DiscretizedDomain(\n domain,\n (; elements, polynomial_order, overintegration_order),\n grid,\n )\nend\n\n# extensions\ncoordinates(grid::DiscretizedDomain) = coordinates(grid.numerical)\npolynomialorders(grid::DiscretizedDomain) = polynomialorders(grid.numerical)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/sphere_helper_functions.jl", "content": "rad(x, y, z) = sqrt(x^2 + y^2 + z^2)\nlat(x, y, z) = asin(z / rad(x, y, z)) # ϕ ∈ [-π/2, π/2] \nlon(x, y, z) = atan(y, x) # λ ∈ [-π, π) \n\nr̂ⁿᵒʳᵐ(x, y, z) = norm([x, y, z]) ≈ 0 ? 1 : norm([x, y, z])^(-1)\nϕ̂ⁿᵒʳᵐ(x, y, z) =\n norm([x, y, 0]) ≈ 0 ? 1 : (norm([x, y, z]) * norm([x, y, 0]))^(-1)\nλ̂ⁿᵒʳᵐ(x, y, z) = norm([x, y, 0]) ≈ 0 ? 1 : norm([x, y, 0])^(-1)\n\nr̂(x, y, z) = r̂ⁿᵒʳᵐ(x, y, z) * @SVector([x, y, z])\nϕ̂(x, y, z) = ϕ̂ⁿᵒʳᵐ(x, y, z) * @SVector [x * z, y * z, -(x^2 + y^2)]\nλ̂(x, y, z) = λ̂ⁿᵒʳᵐ(x, y, z) * @SVector [-y, x, 0]\n\nrfunc(p, x...) = ρuʳᵃᵈ(p, lon(x...), lat(x...), rad(x...)) * r̂(x...)\nϕfunc(p, x...) = ρuˡᵃᵗ(p, lon(x...), lat(x...), rad(x...)) * ϕ̂(x...)\nλfunc(p, x...) = ρuˡᵒⁿ(p, lon(x...), lat(x...), rad(x...)) * λ̂(x...)\n\nρu⃗(p, x...) = rfunc(p, x...) + ϕfunc(p, x...) + λfunc(p, x...)\n\n\nfunction vector_field_representation(\n u⃗,\n grid::DiscontinuousSpectralElementGrid,\n rep::AbstractRepresentation,\n)\n n_ijk, _, n_e = size(u⃗)\n uᴿ = copy(u⃗)\n v⃗ =\n VectorField(data = (u⃗[:, 1, :], u⃗[:, 2, :], u⃗[:, 3, :]), grid = grid)\n for ijk in 1:n_ijk, e in 1:n_e, s in 1:3\n uᴿ[ijk, s, e] = v⃗(rep)[ijk, e, verbose = false][s]\n end\n return uᴿ\nend\n\nvector_field_representation(simulation::Simulation, rep) =\n vector_field_representation(simulation.state.ρu, simulation.model.grid, rep)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_heat_equation.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"../three_dimensional/ThreeDimensionalCompressibleNavierStokesEquations.jl\")\ninclude(\"sphere_helper_functions.jl\")\n\nClimateMachine.init()\n\n#! format: off\nrefVals = (\n[\n [ \"state\", \"ρ\", 9.99999999999998778755e-01, 1.00000000000000111022e+00, 1.00000000000000000000e+00, 4.03070839356576998483e-16 ],\n [ \"state\", \"ρu[1]\", -6.93674606129580362185e-18, 1.19513727285696382782e-17, 9.77991462267996185035e-19, 2.78015892816367900233e-18 ],\n [ \"state\", \"ρu[2]\", -1.21340560458612968582e-17, 1.22630936333135748291e-17, -1.30065979768613433360e-21, 2.78919703417897612909e-18 ],\n [ \"state\", \"ρu[3]\", -1.00545092917566730233e-17, 1.00043577527427175213e-17, 3.76131929192945135398e-19, 2.47195369518944690265e-18 ],\n [ \"state\", \"ρθ\", -3.00959038541689510859e-02, -3.00957518135526284897e-02, -3.00958465374152675520e-02, 3.12498428220800724718e-08 ],\n],\n[\n [ \"state\", \"ρ\", 12, 12, 12, 0 ],\n [ \"state\", \"ρu[1]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[2]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[3]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρθ\", 12, 12, 12, 0 ],\n],\n)\n#! format: on\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n ########\n # Setup physical and numerical domains\n ########\n Ω = AtmosDomain(radius = 1, height = 0.2)\n\n grid = DiscretizedDomain(\n Ω;\n elements = (vertical = 1, horizontal = 5),\n polynomial_order = (vertical = 1 + 0, horizontal = 3 + 0),\n overintegration_order = (vertical = 1, horizontal = 1),\n )\n\n ########\n # Define physical parameters and parameterizations\n ########\n parameters = (\n ρₒ = 1, # reference density\n cₛ = 1e-2, # sound speed\n R = Ω.radius, # [m]\n ω = 0.0, # [s⁻¹]\n K = 7.848e-6, # [s⁻¹]\n n = 4, # dimensionless\n Ω = 0.0, # [s⁻¹] 2π/86400\n )\n\n physics = FluidPhysics(;\n orientation = SphericalOrientation(),\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0.0, κ = 1e-3),\n coriolis = nothing,\n buoyancy = Buoyancy{Float64}(α = 2e-4, g = 0),\n )\n\n ########\n # Define timestepping parameters\n ########\n Δt = 0.1 * min_node_distance(grid.numerical)^2 / physics.dissipation.κ\n start_time = 0\n end_time = 20.0 * Δt\n\n method = SSPRK22Heuns\n\n timestepper = TimeStepper(method = method, timestep = Δt)\n\n callbacks = (Info(), StateCheck(10))\n\n ########\n # Define boundary conditions (west east are the ones that are enforced for a sphere)\n ########\n ρu_bcs = (bottom = Impenetrable(FreeSlip()), top = Impenetrable(FreeSlip()))\n ρθ_bcs = (\n bottom = TemperatureFlux(\n flux = (p, state, aux, t) -> p.Q,\n params = (Q = 1e-3,), # positive means removing heat from system\n ),\n top = TemperatureFlux(\n flux = (p, state, aux, t) -> p.Q,\n params = (Q = 1e-3,), # positive means removing heat from system\n ),\n )\n BC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n ########\n # Define initial conditions\n ########\n\n # Earth Spherical Representation\n # longitude: λ ∈ [-π, π), λ = 0 is the Greenwich meridian\n # latitude: ϕ ∈ [-π/2, π/2], ϕ = 0 is the equator\n # radius: r ∈ [Rₑ - hᵐⁱⁿ, Rₑ + hᵐᵃˣ]\n # Rₑ = Radius of sphere; hᵐⁱⁿ, hᵐᵃˣ ≥ 0\n\n ρ₀(p, λ, ϕ, r) = p.ρₒ\n ρuʳᵃᵈ(p, λ, ϕ, r) = 0.0\n ρuˡᵃᵗ(p, λ, ϕ, r) = 0.0\n ρuˡᵒⁿ(p, λ, ϕ, r) = 0.0\n ρθ₀(p, λ, ϕ, r) = 0.0\n\n # Cartesian Representation (boiler plate really)\n ρ₀ᶜᵃʳᵗ(p, x...) = ρ₀(p, lon(x...), lat(x...), rad(x...))\n ρu⃗₀ᶜᵃʳᵗ(p, x...) = (\n ρuʳᵃᵈ(p, lon(x...), lat(x...), rad(x...)) * r̂(x...) +\n ρuˡᵃᵗ(p, lon(x...), lat(x...), rad(x...)) * ϕ̂(x...) +\n ρuˡᵒⁿ(p, lon(x...), lat(x...), rad(x...)) * λ̂(x...)\n )\n ρθ₀ᶜᵃʳᵗ(p, x...) = ρθ₀(p, lon(x...), lat(x...), rad(x...))\n\n initial_conditions = (ρ = ρ₀ᶜᵃʳᵗ, ρu = ρu⃗₀ᶜᵃʳᵗ, ρθ = ρθ₀ᶜᵃʳᵗ)\n\n ########\n # Create the things\n ########\n\n model = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n tic = Base.time()\n\n @testset \"State Check\" begin\n evolve!(simulation, model, refDat = refVals)\n end\n\n toc = Base.time()\n time = toc - tic\n println(time)\n\n ## Check the budget\n θ̄ = weightedsum(simulation.state, 5)\n\n @testset \"Exact Geometry\" begin\n # Exact (with exact geometry)\n R = grid.domain.radius\n H = grid.domain.height\n\n θ̄ᴱ = 0\n θ̄ᴱ -= 4π * R^2 * end_time * BC.ρθ.bottom.params.Q\n θ̄ᴱ -= 4π * (R + H)^2 * end_time * BC.ρθ.top.params.Q\n\n δθ̄ᴱ = abs(θ̄ - θ̄ᴱ) / abs(θ̄ᴱ)\n println(\"The relative error w.r.t. exact geometry \", δθ̄ᴱ)\n\n @test isapprox(0, δθ̄ᴱ; atol = 1e-9)\n end\n\n # TODO: make Approx Geom test MPI safe\n #=\n @testset \"Approx Geometry\" begin\n # Exact (with approximated geometry)\n n_e = convention(grid.resolution.elements, Val(3))\n n_ijk = convention(grid.resolution.polynomial_order, Val(3))\n n_ijk =\n n_ijk .+ convention(grid.resolution.overintegration_order, Val(3))\n n_ijk = n_ijk .+ (1, 1, 1)\n\n Mᴴ = reshape(\n grid.numerical.vgeo[:, grid.numerical.MHid, :],\n (n_ijk..., n_e[3], 6 * n_e[1] * n_e[2]),\n )\n Aᴮᵒᵗ = sum(Mᴴ[:, :, 1, 1, :])\n Aᵀᵒᵖ = sum(Mᴴ[:, :, end, end, :])\n\n θ̄ᴬ = 0\n θ̄ᴬ -= Aᴮᵒᵗ * end_time * BC.ρθ.bottom.params.Q\n θ̄ᴬ -= Aᵀᵒᵖ * end_time * BC.ρθ.top.params.Q\n\n δθ̄ᴬ = abs(θ̄ - θ̄ᴬ) / abs(θ̄ᴬ)\n println(\"The relative error w.r.t. approximated geometry \", δθ̄ᴬ)\n\n @test isapprox(0, δθ̄ᴬ; atol = 1e-15)\n end\n =#\n\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_hydrostatic_balance.jl", "content": "#!/usr/bin/env julia --project\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"../three_dimensional/ThreeDimensionalCompressibleNavierStokesEquations.jl\")\ninclude(\"sphere_helper_functions.jl\")\n\nClimateMachine.init()\n\n#! format: off\nrefVals = (\n[\n [ \"state\", \"ρ\", 1.99999999999600046319e-02, 1.00000000000004396483e+00, 5.10000000000001008083e-01, 2.91613390275248574035e-01 ],\n [ \"state\", \"ρu[1]\", -3.13311246519756673238e-11, 3.34851852426747048856e-11, 1.50582245544617117858e-13, 4.32064975692569532007e-12 ],\n [ \"state\", \"ρu[2]\", -3.98988708940419829624e-11, 4.12795751448340701883e-11, -1.66842288197274268178e-14, 4.34989190229689036913e-12 ],\n [ \"state\", \"ρu[3]\", -4.16490644835465690341e-11, 4.18355107496431387673e-11, 5.18313726923951858909e-14, 4.62898170296138141882e-12 ],\n [ \"state\", \"ρθ\", -4.90000000000050732751e+01, -9.79999999998041548821e-01, -2.49900000000000446221e+01, 1.42890561234872102148e+01 ],\n],\n[\n [ \"state\", \"ρ\", 12, 12, 12, 12 ],\n [ \"state\", \"ρu[1]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[2]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[3]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρθ\", 12, 12, 12, 12 ],\n],\n)\n#! format: on\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n ########\n # Setup physical and numerical domains\n ########\n Ω = AtmosDomain(radius = 6e6, height = 1e5)\n grid = DiscretizedDomain(\n Ω;\n elements = (vertical = 4, horizontal = 4),\n polynomial_order = (vertical = 1, horizontal = 3),\n overintegration_order = (vertical = 1, horizontal = 1),\n )\n\n ########\n # Define physical parameters and parameterizations\n ########\n parameters = (\n ρₒ = 1, # reference density\n cₛ = 100.0, # sound speed\n ν = 1e-5,\n ∂θ = 0.98 / 1e5,\n α = 2e-4,\n g = 10.0,\n power = 1,\n )\n\n ########\n # Define timestepping parameters\n ########\n Δtᴬ = min_node_distance(grid.numerical) / parameters.cₛ * 0.25\n Δtᴰ = min_node_distance(grid.numerical)^2 / parameters.ν * 0.25\n Δt = minimum([Δtᴬ, Δtᴰ])\n start_time = 0\n end_time = 86400 * 0.5\n\n method = SSPRK22Heuns\n timestepper = TimeStepper(method = method, timestep = Δt)\n callbacks = (Info(), StateCheck(10))# , VTKState(iteration = 2880, filepath = \".\"))\n\n physics = FluidPhysics(;\n orientation = SphericalOrientation(),\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(\n μ = 0,\n ν = parameters.ν,\n κ = 0.0,\n ),\n coriolis = SphereCoriolis{Float64}(Ω = 2π / 86400),\n buoyancy = Buoyancy{Float64}(α = parameters.α, g = parameters.g),\n )\n\n ########\n # Define boundary conditions (west east are the ones that are enforced for a sphere)\n ########\n ρu_bcs = (bottom = Impenetrable(NoSlip()), top = Impenetrable(NoSlip()))\n ρθ_bcs = (bottom = Insulating(), top = Insulating())\n BC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n ########\n # Define initial conditions\n ########\n # Earth Spherical Representation\n # longitude: λ ∈ [-π, π), λ = 0 is the Greenwich meridian\n # latitude: ϕ ∈ [-π/2, π/2], ϕ = 0 is the equator\n # radius: r ∈ [Rₑ - hᵐⁱⁿ, Rₑ + hᵐᵃˣ], Rₑ = Radius of sphere; hᵐⁱⁿ, hᵐᵃˣ ≥ 0\n\n ρ₀(p, λ, ϕ, r) =\n (1 - p.∂θ * (r - 6e6)^p.power / p.power * 1e5^(1 - p.power)) * p.ρₒ\n ρuʳᵃᵈ(p, λ, ϕ, r) = 0.0\n ρuˡᵃᵗ(p, λ, ϕ, r) = 0.0\n ρuˡᵒⁿ(p, λ, ϕ, r) = 0.0\n ρθ₀(p, λ, ϕ, r) =\n -ρ₀(p, λ, ϕ, r) *\n p.∂θ *\n (r - 6e6)^(p.power - 1) *\n 1e5^(1 - p.power) *\n (p.cₛ)^2 / (p.α * p.g)\n\n # Cartesian Representation (boiler plate really)\n ρ₀ᶜᵃʳᵗ(p, x...) = ρ₀(p, lon(x...), lat(x...), rad(x...))\n ρu⃗₀ᶜᵃʳᵗ(p, x...) = (\n ρuʳᵃᵈ(p, lon(x...), lat(x...), rad(x...)) * r̂(x...) +\n ρuˡᵃᵗ(p, lon(x...), lat(x...), rad(x...)) * ϕ̂(x...) +\n ρuˡᵒⁿ(p, lon(x...), lat(x...), rad(x...)) * λ̂(x...)\n )\n ρθ₀ᶜᵃʳᵗ(p, x...) = ρθ₀(p, lon(x...), lat(x...), rad(x...))\n\n initial_conditions = (ρ = ρ₀ᶜᵃʳᵗ, ρu = ρu⃗₀ᶜᵃʳᵗ, ρθ = ρθ₀ᶜᵃʳᵗ)\n\n ########\n # Create the things\n ########\n model = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n tic = Base.time()\n evolve!(simulation, model, refDat = refVals)\n toc = Base.time()\n time = toc - tic\n println(time)\n\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/sphere/test_sphere.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"../three_dimensional/ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n#! format: off\nrefVals = (\n[\n [ \"state\", \"ρ\", 9.99999999999999000799e-01, 1.00000000000000066613e+00, 1.00000000000000000000e+00, 3.49433608871729305050e-16 ],\n [ \"state\", \"ρu[1]\", -1.03159890504820785885e-17, 9.46342010159792552240e-18, -5.38891312633679647218e-19, 2.05248736591058928424e-18 ],\n [ \"state\", \"ρu[2]\", -1.12556410356665774824e-17, 1.11125313793689735132e-17, -2.17938113394074827612e-22, 2.14911076688864467562e-18 ],\n [ \"state\", \"ρu[3]\", -1.12757405914899880325e-17, 9.66517285069783167720e-18, -2.95076506235937354400e-19, 2.21386287201671452859e-18 ],\n [ \"state\", \"ρθ\", 9.99999999999999000799e-01, 1.00000000000000066613e+00, 1.00000000000000000000e+00, 3.49433608871729305050e-16 ],\n],\n[\n [ \"state\", \"ρ\", 12, 12, 12, 0 ],\n [ \"state\", \"ρu[1]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[2]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρu[3]\", 0, 0, 0, 0 ],\n [ \"state\", \"ρθ\", 12, 12, 12, 0 ],\n],\n)\n#! format: on\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n ########\n # Setup physical and numerical domains\n ########\n Ω = AtmosDomain(radius = 1, height = 0.02)\n grid = DiscretizedDomain(\n Ω;\n elements = (vertical = 1, horizontal = 4),\n polynomial_order = (vertical = 1, horizontal = 4),\n overintegration_order = 1,\n )\n\n ########\n # Define timestepping parameters\n ########\n start_time = 0\n end_time = 2.0\n Δt = 0.05\n method = SSPRK22Heuns\n timestepper = TimeStepper(method = method, timestep = Δt)\n callbacks = (Info(), StateCheck(10))\n\n ########\n # Define physical parameters and parameterizations\n ########\n parameters = (\n ρₒ = 1, # reference density\n cₛ = 1e-2, # sound speed\n )\n\n physics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0.0, κ = 0.0),\n coriolis = nothing,\n buoyancy = Buoyancy{Float64}(α = 2e-4, g = 0),\n )\n\n ########\n # Define boundary conditions (west east are the ones that are enforced for a sphere)\n ########\n ρu_bcs = (bottom = Impenetrable(NoSlip()), top = Impenetrable(NoSlip()))\n ρθ_bcs = (bottom = Insulating(), top = Insulating())\n BC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n ########\n # Define initial conditions\n ########\n ρ₀(p, x, y, z) = p.ρₒ\n ρu₀(p, x...) = ρ₀(p, x...) * -0\n ρv₀(p, x...) = ρ₀(p, x...) * -0\n ρw₀(p, x...) = ρ₀(p, x...) * -0\n ρθ₀(p, x...) = ρ₀(p, x...) * 1.0\n\n ρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\n initial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n ########\n # Create the things\n ########\n model = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n tic = Base.time()\n evolve!(simulation, model; refDat = refVals)\n toc = Base.time()\n time = toc - tic\n println(time)\n\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/ThreeDimensionalCompressibleNavierStokesEquations.jl", "content": "include(\"../shared_source/boilerplate.jl\")\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_conditions,\n boundary_state!\nimport ClimateMachine.DGMethods: DGModel\nimport ClimateMachine.NumericalFluxes: numerical_flux_first_order!\nimport ClimateMachine.Orientations: vertical_unit_vector\n\n\"\"\"\n ThreeDimensionalCompressibleNavierStokesEquations <: BalanceLaw\nA `BalanceLaw` for shallow water modeling.\nwrite out the equations here\n# Usage\n ThreeDimensionalCompressibleNavierStokesEquations()\n\"\"\"\nabstract type AbstractFluid3D <: AbstractFluid end\nstruct Fluid3D <: AbstractFluid3D end\n\nstruct ThreeDimensionalCompressibleNavierStokesEquations{\n I,\n D,\n O,\n A,\n T,\n C,\n F,\n BC,\n FT,\n} <: AbstractFluid3D\n initial_value_problem::I\n domain::D\n orientation::O\n advection::A\n turbulence::T\n coriolis::C\n forcing::F\n boundary_conditions::BC\n cₛ::FT\n ρₒ::FT\n function ThreeDimensionalCompressibleNavierStokesEquations{FT}(\n initial_value_problem::I,\n domain::D,\n orientation::O,\n advection::A,\n turbulence::T,\n coriolis::C,\n forcing::F,\n boundary_conditions::BC;\n cₛ = FT(sqrt(10)), # m/s\n ρₒ = FT(1), #kg/m³\n ) where {FT <: AbstractFloat, I, D, O, A, T, C, F, BC}\n return new{I, D, O, A, T, C, F, BC, FT}(\n initial_value_problem,\n domain,\n orientation,\n advection,\n turbulence,\n coriolis,\n forcing,\n boundary_conditions,\n cₛ,\n ρₒ,\n )\n end\nend\n\nCNSE3D = ThreeDimensionalCompressibleNavierStokesEquations\n\nfunction vars_state(m::CNSE3D, ::Prognostic, T)\n @vars begin\n ρ::T\n ρu::SVector{3, T}\n ρθ::T\n end\nend\n\nfunction init_state_prognostic!(m::CNSE3D, state::Vars, aux::Vars, localgeo, t)\n cnse_init_state!(m, state, aux, localgeo, t)\nend\n\n# default initial state if IVP == nothing\nfunction cnse_init_state!(model::CNSE3D, state, aux, localgeo, t)\n\n x = aux.x\n y = aux.y\n z = aux.z\n\n ρ = model.ρₒ\n state.ρ = ρ\n state.ρu = ρ * @SVector [-0, -0, -0]\n state.ρθ = ρ\n\n return nothing\nend\n\n# user defined initial state\nfunction cnse_init_state!(\n model::CNSE3D{<:InitialValueProblem},\n state,\n aux,\n localgeo,\n t,\n)\n x = aux.x\n y = aux.y\n z = aux.z\n\n params = model.initial_value_problem.params\n ic = model.initial_value_problem.initial_conditions\n\n state.ρ = ic.ρ(params, x, y, z)\n state.ρu = ic.ρu(params, x, y, z)\n state.ρθ = ic.ρθ(params, x, y, z)\n\n return nothing\nend\n\nfunction vars_state(m::CNSE3D, st::Auxiliary, T)\n @vars begin\n x::T\n y::T\n z::T\n orientation::vars_state(m.orientation, st, T)\n end\nend\n\nfunction init_state_auxiliary!(\n m::CNSE3D,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n # update the geopotential Φ in state_auxiliary.orientation.Φ\n init_state_auxiliary!(\n m,\n (m, aux, tmp, geom) ->\n orientation_nodal_init_aux!(m.orientation, m.domain, aux, geom),\n state_auxiliary,\n grid,\n direction,\n )\n\n # update ∇Φ in state_auxiliary.orientation.∇Φ\n orientation_gradient(m, m.orientation, state_auxiliary, grid, direction)\n\n # store coordinates and potentially other stuff\n init_state_auxiliary!(\n m,\n (m, aux, tmp, geom) -> cnse_init_aux!(m, aux, geom),\n state_auxiliary,\n grid,\n direction,\n )\n\n return nothing\nend\n\nfunction orientation_gradient(\n model::CNSE3D,\n ::Orientation,\n state_auxiliary,\n grid,\n direction,\n)\n auxiliary_field_gradient!(\n model,\n state_auxiliary,\n (\"orientation.∇Φ\",),\n state_auxiliary,\n (\"orientation.Φ\",),\n grid,\n direction,\n )\n\n return nothing\nend\n\nfunction orientation_gradient(::CNSE3D, ::NoOrientation, _...)\n return nothing\nend\n\nfunction orientation_nodal_init_aux!(\n ::SphericalOrientation,\n domain::Tuple,\n aux::Vars,\n geom::LocalGeometry,\n)\n norm_R = norm(geom.coord)\n @inbounds aux.orientation.Φ = norm_R - domain[1]\n\n return nothing\nend\n\n\"\"\"\nfunction orientation_nodal_init_aux!(\n ::SuperSphericalOrientation,\n domain::Tuple,\n aux::Vars,\n geom::LocalGeometry,\n)\n norm_R = norm(geom.coord)\n @inbounds aux.orientation.Φ = 1 / norm_R^2 \nend\n\"\"\"\n\nfunction orientation_nodal_init_aux!(\n ::FlatOrientation,\n domain::Tuple,\n aux::Vars,\n geom::LocalGeometry,\n)\n @inbounds aux.orientation.Φ = geom.coord[3]\n\n return nothing\nend\n\nfunction orientation_nodal_init_aux!(\n ::NoOrientation,\n domain::Tuple,\n aux::Vars,\n geom::LocalGeometry,\n)\n return nothing\nend\n\nfunction cnse_init_aux!(::CNSE3D, aux, geom)\n @inbounds begin\n aux.x = geom.coord[1]\n aux.y = geom.coord[2]\n aux.z = geom.coord[3]\n end\n\n return nothing\nend\n\nfunction vars_state(m::CNSE3D, ::Gradient, T)\n @vars begin\n ∇ρ::T\n ∇u::SVector{3, T}\n ∇θ::T\n ∇p::T\n end\nend\n\nfunction compute_gradient_argument!(\n model::CNSE3D,\n grad::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ρ = state.ρ\n cₛ = model.cₛ\n ρₒ = model.ρₒ\n\n grad.∇p = (cₛ * ρ)^2 / (2 * ρₒ)\n\n compute_gradient_argument!(model.turbulence, grad, state, aux, t)\nend\n\n@inline function compute_gradient_argument!(\n ::ConstantViscosity,\n grad::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρθ = state.ρθ\n\n u = ρu\n θ = ρθ\n\n grad.∇ρ = ρ\n grad.∇u = u\n grad.∇θ = θ\n\n return nothing\nend\n\nfunction vars_state(m::CNSE3D, ::GradientFlux, T)\n @vars begin\n μ∇ρ::SVector{3, T}\n ν∇u::SMatrix{3, 3, T, 9}\n κ∇θ::SVector{3, T}\n ∇p::SVector{3, T}\n end\nend\n\nfunction compute_gradient_flux!(\n model::CNSE3D,\n gradflux::Vars,\n grad::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n\n gradflux.∇p = grad.∇p\n\n compute_gradient_flux!(\n model,\n model.turbulence,\n gradflux,\n grad,\n state,\n aux,\n t,\n )\nend\n\n@inline function compute_gradient_flux!(\n ::CNSE3D,\n turb::ConstantViscosity,\n gradflux::Vars,\n grad::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n μ = turb.μ * I\n ν = turb.ν * I\n κ = turb.κ * I\n\n gradflux.μ∇ρ = -μ * grad.∇ρ\n gradflux.ν∇u = -ν * grad.∇u\n gradflux.κ∇θ = -κ * grad.∇θ\n\n return nothing\nend\n\n@inline function flux_first_order!(\n model::CNSE3D,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρθ = state.ρθ\n\n cₛ = model.cₛ\n ρₒ = model.ρₒ\n\n flux.ρ += ρu\n # flux.ρu += (cₛ * ρ)^2 / (2 * ρₒ) * I\n\n advective_flux!(model, model.advection, flux, state, aux, t)\n\n return nothing\nend\n\nadvective_flux!(::CNSE3D, ::Nothing, _...) = nothing\n\n@inline function advective_flux!(\n ::CNSE3D,\n ::NonLinearAdvectionTerm,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρθ = state.ρθ\n\n flux.ρu += ρu ⊗ ρu / ρ\n flux.ρθ += ρu * ρθ / ρ\n\n return nothing\nend\n\nfunction flux_second_order!(\n model::CNSE3D,\n flux::Grad,\n state::Vars,\n gradflux::Vars,\n ::Vars,\n aux::Vars,\n t::Real,\n)\n flux_second_order!(model, model.turbulence, flux, state, gradflux, aux, t)\nend\n\n@inline function flux_second_order!(\n ::CNSE3D,\n ::ConstantViscosity,\n flux::Grad,\n state::Vars,\n gradflux::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρ += gradflux.μ∇ρ\n flux.ρu += gradflux.ν∇u\n flux.ρθ += gradflux.κ∇θ\n\n return nothing\nend\n\n@inline function source!(\n model::CNSE3D,\n source::Vars,\n state::Vars,\n gradflux::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n\n source.ρu -= gradflux.∇p\n\n coriolis_force!(model, model.coriolis, source, state, aux, t)\n forcing_term!(model, model.forcing, source, state, aux, t)\n\n return nothing\nend\n\ncoriolis_force!(::CNSE3D, ::Nothing, _...) = nothing\n\n@inline function coriolis_force!(\n model::CNSE3D,\n coriolis::fPlaneCoriolis,\n source,\n state,\n aux,\n t,\n)\n # f × u\n f = [-0, -0, coriolis_parameter(model, coriolis, aux.coords)]\n ρu = state.ρu\n\n source.ρu -= f × ρu\n\n return nothing\nend\n\n@inline function coriolis_force!(\n model::CNSE3D,\n coriolis::SphereCoriolis,\n source,\n state,\n aux,\n t,\n)\n # f × u\n Ω = coriolis.Ω\n f = @SVector [-0, -0, 2Ω]\n ρu = state.ρu\n source.ρu -= f × ρu\n return nothing\nend\n\n\nforcing_term!(::CNSE3D, ::Nothing, _...) = nothing\n\n@inline function forcing_term!(\n model::CNSE3D,\n buoy::Buoyancy,\n source,\n state,\n aux,\n t,\n)\n α = buoy.α\n g = buoy.g\n ρθ = state.ρθ\n\n # as temperature increase, density decreases\n B = -α * g * ρθ\n k̂ = vertical_unit_vector(model.orientation, aux)\n\n # gravity points downward\n source.ρu -= B * k̂\nend\n\n@inline vertical_unit_vector(::Orientation, aux) = aux.orientation.∇Φ\n@inline vertical_unit_vector(::NoOrientation, aux) = @SVector [0, 0, 1]\n\n@inline wavespeed(m::CNSE3D, _...) = m.cₛ\n\nroe_average(ρ⁻, ρ⁺, var⁻, var⁺) =\n (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))\n\nfunction numerical_flux_first_order!(\n ::RoeNumericalFlux,\n model::CNSE3D,\n fluxᵀn::Vars{S},\n n⁻::SVector,\n state⁻::Vars{S},\n aux⁻::Vars{A},\n state⁺::Vars{S},\n aux⁺::Vars{A},\n t,\n direction,\n) where {S, A}\n numerical_flux_first_order!(\n CentralNumericalFluxFirstOrder(),\n model,\n fluxᵀn,\n n⁻,\n state⁻,\n aux⁻,\n state⁺,\n aux⁺,\n t,\n direction,\n )\n\n FT = eltype(fluxᵀn)\n\n # constants and normal vectors\n cₛ = model.cₛ\n ρₒ = model.ρₒ\n\n # - states\n ρ⁻ = state⁻.ρ\n ρu⁻ = state⁻.ρu\n ρθ⁻ = state⁻.ρθ\n\n # constructed states\n u⁻ = ρu⁻ / ρ⁻\n θ⁻ = ρθ⁻ / ρ⁻\n uₙ⁻ = u⁻' * n⁻\n\n # in general thermodynamics\n p⁻ = (cₛ * ρ⁻)^2 / (2 * ρₒ)\n c⁻ = cₛ * sqrt(ρ⁻ / ρₒ)\n\n # + states\n ρ⁺ = state⁺.ρ\n ρu⁺ = state⁺.ρu\n ρθ⁺ = state⁺.ρθ\n\n # constructed states\n u⁺ = ρu⁺ / ρ⁺\n θ⁺ = ρθ⁺ / ρ⁺\n uₙ⁺ = u⁺' * n⁻\n\n # in general thermodynamics\n p⁺ = (cₛ * ρ⁺)^2 / (2 * ρₒ)\n c⁺ = cₛ * sqrt(ρ⁺ / ρₒ)\n\n # construct roe averges\n ρ = sqrt(ρ⁻ * ρ⁺)\n u = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)\n θ = roe_average(ρ⁻, ρ⁺, θ⁻, θ⁺)\n c = roe_average(ρ⁻, ρ⁺, c⁻, c⁺)\n\n # construct normal velocity\n uₙ = u' * n⁻\n\n # differences\n Δρ = ρ⁺ - ρ⁻\n Δp = p⁺ - p⁻\n Δu = u⁺ - u⁻\n Δρθ = ρθ⁺ - ρθ⁻\n Δuₙ = Δu' * n⁻\n\n # constructed values\n c⁻² = 1 / c^2\n w1 = abs(uₙ - c) * (Δp - ρ * c * Δuₙ) * 0.5 * c⁻²\n w2 = abs(uₙ + c) * (Δp + ρ * c * Δuₙ) * 0.5 * c⁻²\n w3 = abs(uₙ) * (Δρ - Δp * c⁻²)\n w4 = abs(uₙ) * ρ\n w5 = abs(uₙ) * (Δρθ - θ * Δp * c⁻²)\n\n # fluxes!!!\n fluxᵀn.ρ -= (w1 + w2 + w3) * 0.5\n fluxᵀn.ρu -=\n (\n w1 * (u - c * n⁻) +\n w2 * (u + c * n⁻) +\n w3 * u +\n w4 * (Δu - Δuₙ * n⁻)\n ) * 0.5\n fluxᵀn.ρθ -= ((w1 + w2) * θ + w5) * 0.5\n\n return nothing\nend\n\nboundary_conditions(model::CNSE3D) = model.boundary_conditions\n\n\"\"\"\n boundary_state!(nf, ::CNSE3D, args...)\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in CNSEBoundaryConditions\n\"\"\"\n@inline function boundary_state!(nf, bc, model::CNSE3D, args...)\n return _cnse_boundary_state!(nf, bc, model, args...)\nend\n\n\"\"\"\n cnse_boundary_state!(nf, bc::FluidBC, ::CNSE3D)\nsplits boundary condition application into velocity\n\"\"\"\n@inline function cnse_boundary_state!(nf, bc::FluidBC, m::CNSE3D, args...)\n cnse_boundary_state!(nf, bc.momentum, m, m.turbulence, args...)\n cnse_boundary_state!(nf, bc.temperature, m, args...)\n\n return nothing\nend\n\ninclude(\"bc_momentum.jl\")\ninclude(\"bc_temperature.jl\")\n\n\"\"\"\nSTUFF FOR ANDRE'S WRAPPERS\n\"\"\"\n\nfunction get_boundary_conditions(\n model::SpatialModel{BL},\n) where {BL <: AbstractFluid3D}\n bcs = model.boundary_conditions\n\n west_east = (check_bc(bcs, :west), check_bc(bcs, :east))\n south_north = (check_bc(bcs, :south), check_bc(bcs, :north))\n bottom_top = (check_bc(bcs, :bottom), check_bc(bcs, :top))\n\n return (west_east..., south_north..., bottom_top...)\nend\n\nfunction DGModel(\n model::SpatialModel{BL};\n initial_conditions = nothing,\n) where {BL <: AbstractFluid3D}\n params = model.parameters\n physics = model.physics\n\n Lˣ, Lʸ, Lᶻ = length(model.grid.domain)\n bcs = get_boundary_conditions(model)\n FT = eltype(model.grid.numerical.vgeo)\n\n if !isnothing(initial_conditions)\n initial_conditions = InitialValueProblem(params, initial_conditions)\n end\n\n balance_law = CNSE3D{FT}(\n initial_conditions,\n (Lˣ, Lʸ, Lᶻ),\n physics.orientation,\n physics.advection,\n physics.dissipation,\n physics.coriolis,\n physics.buoyancy,\n bcs,\n ρₒ = params.ρₒ,\n cₛ = params.cₛ,\n )\n\n numerical_flux_first_order = model.numerics.flux # should be a function\n\n rhs = DGModel(\n balance_law,\n model.grid.numerical,\n numerical_flux_first_order,\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n return rhs\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/bc_momentum.jl", "content": "\"\"\"\n cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::CNSE3D)\napply free slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{FreeSlip},\n ::CNSE3D,\n ::TurbulenceClosure,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n\n ρu⁻ = state⁻.ρu\n state⁺.ρu = ρu⁻ - 2 * n⁻ ⋅ ρu⁻ .* SVector(n⁻)\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::CNSE3D)\napply free slip boundary condition for velocity\nsets non-reflective ghost point\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{FreeSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n\n ρu⁻ = state⁻.ρu\n state⁺.ρu = ρu⁻ - n⁻ ⋅ ρu⁻ .* SVector(n⁻)\n\n return nothing\nend\n\n\"\"\"\n shallow_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::CNSE3D)\napply free slip boundary condition for velocity\napply zero numerical flux in the normal direction\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{FreeSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0, -0])'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::CNSE3D)\napply no slip boundary condition for velocity\nsets reflective ghost point\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{NoSlip},\n ::CNSE3D,\n ::TurbulenceClosure,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n state⁺.ρu = -state⁻.ρu\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::CNSE3D)\napply no slip boundary condition for velocity\nset numerical flux to zero for U\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{NoSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n FT = eltype(state⁺)\n state⁺.ρu = @SVector zeros(FT, 3)\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::CNSE3D)\napply no slip boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for U\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{NoSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = -state⁻.ρu\n gradflux⁺.ν∇u = gradflux⁻.ν∇u\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::CNSE3D)\nno mass boundary condition for penetrable\n\"\"\"\ncnse_boundary_state!(\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{FreeSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n _...,\n) = nothing\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::CNSE3D)\napply free slip boundary condition for velocity\napply zero numerical flux in the normal direction\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Penetrable{FreeSlip},\n ::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0, -0])'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{MomentumFlux}, ::HBModel)\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Impenetrable{<:MomentumFlux},\n model::CNSE3D,\n turb::TurbulenceClosure,\n args...,\n)\n return cnse_boundary_state!(\n nf,\n Impenetrable(FreeSlip()),\n model,\n turb,\n args...,\n )\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{MomentumFlux}, ::HBModel)\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n bc::Impenetrable{<:MomentumFlux},\n model::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{MomentumFlux}, ::HBModel)\napply kinematic stress boundary condition for velocity\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{<:MomentumFlux},\n model::CNSE3D,\n turb::TurbulenceClosure,\n args...,\n)\n return cnse_boundary_state!(\n nf,\n Penetrable(FreeSlip()),\n model,\n turb,\n args...,\n )\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{MomentumFlux}, ::HBModel)\napply kinematic stress boundary condition for velocity\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n bc::Penetrable{<:MomentumFlux},\n shallow::CNSE3D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'\n\n return nothing\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/bc_temperature.jl", "content": "\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Insulating, ::HBModel)\n\napply insulating boundary condition for temperature\nsets transmissive ghost point\n\"\"\"\nfunction cnse_boundary_state!(\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Insulating,\n ::CNSE3D,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Insulating, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for κ∇θ\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Insulating,\n ::CNSE3D,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n gradflux⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::TemperatureFlux, ::HBModel)\n\napply temperature flux boundary condition for velocity\napplies insulating conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::TemperatureFlux,\n model::CNSE3D,\n args...,\n)\n return cnse_boundary_state!(nf, Insulating(), model, args...)\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::TemperatureFlux, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have specified flux on the boundary for κ∇θ\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n bc::TemperatureFlux,\n ::CNSE3D,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n gradflux⁺.κ∇θ = n⁻ * bc(state⁻, aux⁻, t)\n\n return nothing\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/config_sphere.jl", "content": "include(\"../CNSE.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nfunction Config(\n name,\n resolution,\n domain,\n params;\n numerical_flux_first_order = RoeNumericalFlux(),\n Nover = 0,\n boundary = (1, 1),\n boundary_conditons = (FluidBC(Impenetrable(FreeSlip()), Insulating()),),\n)\n mpicomm = MPI.COMM_WORLD\n ArrayType = ClimateMachine.array_type()\n\n println(string(resolution.Nᶻ) * \" elems in the vertical\")\n\n vert_range =\n grid1d(domain.min_height, domain.max_height, nelem = resolution.Nᶻ)\n\n println(\n string(resolution.Nʰ) * \"x\" * string(resolution.Nʰ) * \" elems per face\",\n )\n\n topology = StackedCubedSphereTopology(\n mpicomm,\n resolution.Nʰ,\n vert_range;\n boundary = boundary,\n )\n\n println(\"poly order is \" * string(resolution.N))\n println(\"OI order is \" * string(Nover))\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = resolution.N + Nover,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n\n model = CNSE3D{FT}(\n nothing,\n (domain.min_height, domain.max_height),\n ClimateMachine.Orientations.SphericalOrientation(),\n NonLinearAdvectionTerm(),\n ConstantViscosity{FT}(μ = params.μ, ν = params.ν, κ = params.κ),\n nothing,\n nothing,\n boundary_conditons;\n cₛ = params.cₛ,\n ρₒ = params.ρₒ,\n )\n\n dg = DGModel(\n model,\n grid,\n numerical_flux_first_order,\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n return Config(name, dg, Nover, mpicomm, ArrayType)\nend\n\nfunction cnse_init_aux!(::CNSE3D, aux, geom)\n @inbounds begin\n aux.x = geom.coord[1]\n aux.y = geom.coord[2]\n aux.z = geom.coord[3]\n end\n\n return nothing\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/refvals_bickley_jet.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\n#! format: off\nfirst_order = (\n[\n [ \"state\", \"ρ\", 9.88641470218911577739e-01, 1.00735927139044068035e+00, 1.00000000000000244249e+00, 1.32618819574356725140e-03 ],\n [ \"state\", \"ρu[1]\", -2.42850976605026747102e-01, 5.73364756714398238202e-01, 1.59153832831308655882e-01, 9.67668639083398146594e-02 ],\n [ \"state\", \"ρu[2]\", -4.12933704690272018745e-01, 4.50172304099278386413e-01, -4.55511721449193612529e-14, 1.48850741581273815495e-01 ],\n [ \"state\", \"ρu[3]\", -3.70076433121673098459e-01, 3.26009261713208653433e-01, -6.46953234297288930041e-14, 8.65065044705036617634e-02 ],\n [ \"state\", \"ρθ\", -1.71139241110591289186e+00, 1.79353274817685659492e+00, -7.54540463032282362914e-16, 3.87893456476833764501e-01 ],\n],\n[\n [\"state\", \"ρ\", 12, 12, 12, 12],\n [\"state\", \"ρu[1]\", 11, 11, 12, 12],\n [\"state\", \"ρu[2]\", 12, 11, 0, 12],\n [\"state\", \"ρu[3]\", 12, 11, 0, 12],\n [\"state\", \"ρθ\", 11, 10, 0, 12],\n],\n)\n\nfourth_order = (\n[\n [ \"state\", \"ρ\", 9.79756567265035793746e-01, 1.00986882291859325633e+00, 9.99972252557182250676e-01, 1.27738464370501453651e-03 ],\n [ \"state\", \"ρu[1]\", -3.85746589538976891731e-01, 7.12308328222086784010e-01, 1.58403896799433618892e-01, 1.03266735858658170732e-01 ],\n [ \"state\", \"ρu[2]\", -5.86525244955121705104e-01, 6.33258737050533038193e-01, 2.62244534853688972004e-04, 1.37491363204364197559e-01 ],\n [ \"state\", \"ρu[3]\", -4.58727280284793481613e-01, 4.82564240905019259387e-01, -2.81636420608788126414e-04, 8.90227260571099660025e-02 ],\n [ \"state\", \"ρθ\", -4.51696345879803118351e+00, 4.14217705122501378412e+00, -5.08780751729477078403e-04, 3.32659156438805059253e-01 ],\n],\n[\n [\"state\", \"ρ\", 8, 8, 10, 7],\n [\"state\", \"ρu[1]\", 5, 5, 8, 7],\n [\"state\", \"ρu[2]\", 6, 5, 0, 6],\n [\"state\", \"ρu[3]\", 6, 5, 0, 7],\n [\"state\", \"ρθ\", 5, 4, 0, 6],\n],\n)\n\n#! format: on\n\nrefVals = (; first_order, fourth_order)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/refvals_buoyancy.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\n#! format: off\n\nparr = [\n [\"state\", \"ρ\", 12, 12, 12, 12],\n [\"state\", \"ρu[1]\", 0, 0, 0, 0],\n [\"state\", \"ρu[2]\", 0, 0, 0, 0],\n [\"state\", \"ρu[3]\", 12, 8, 12, 12],\n [\"state\", \"ρθ\", 12, 10, 12, 12],\n]\n\nsecond_order_flat = (\n[\n [ \"state\", \"ρ\", 9.95252314022507689195e-01, 9.99992856011554298590e-01, 9.98330419819817738158e-01, 1.48639562654353791886e-03 ],\n [ \"state\", \"ρu[1]\", -6.36093867411581375298e-15, 8.28629359748938561747e-15, 5.16470747993988398929e-16, 1.17216775586980107502e-15 ],\n [ \"state\", \"ρu[2]\", -3.70797682276813316617e-15, 9.42818503691220985643e-15, 6.07164354178963622023e-16, 1.23143801722553508699e-15 ],\n [ \"state\", \"ρu[3]\", -1.65133743521589989450e-03, 5.29367075491514133403e-09, -8.40309050919211576736e-04, 4.66618546037470796790e-04 ],\n [ \"state\", \"ρθ\", -9.95249493245247940365e+00, 1.99973806108046075331e-05, -4.98740538889376860965e+00, 2.91965149708573168397e+00 ],\n],\n parr,\n)\n\nsecond_order = (\n[\n [ \"state\", \"ρ\", 9.95252314022507689195e-01, 9.99992856011554298590e-01, 9.98330419819817738158e-01, 1.48639562654353791886e-03 ],\n [ \"state\", \"ρu[1]\", -6.36142448250610767485e-15, 8.28656656189152269018e-15, 5.17120118507278235699e-16, 1.17302082244252352324e-15 ],\n [ \"state\", \"ρu[2]\", -3.70772154095122499196e-15, 9.42925797553789671672e-15, 6.06575884435408245721e-16, 1.23026363135339100359e-15 ],\n [ \"state\", \"ρu[3]\", -1.65133743521588883564e-03, 5.29367075398690075732e-09, -8.40309050919211468315e-04, 4.66618546037470417320e-04 ],\n [ \"state\", \"ρθ\", -9.95249493245247940365e+00, 1.99973806108054952236e-05, -4.98740538889376860965e+00, 2.91965149708573168397e+00 ],\n],\n parr,\n)\n\nfourth_order_flat = (\n[\n [ \"state\", \"ρ\", 9.95377495534709000324e-01, 9.99992951378667060958e-01, 9.98321272635789513927e-01, 1.50722816639464090097e-03 ],\n [ \"state\", \"ρu[1]\", -1.22866254290737312252e-14, 1.95946879007561421956e-14, 1.73741725325975513295e-15, 4.23083086870931785383e-15 ],\n [ \"state\", \"ρu[2]\", -1.63058197687820133665e-14, 2.71134092821328864807e-14, 1.75302916450504673999e-15, 4.33921909208334743316e-15 ],\n [ \"state\", \"ρu[3]\", -1.66252523985498286418e-03, 5.55021884647818839491e-08, -8.13842490777055664608e-04, 4.76414940515918383830e-04 ],\n [ \"state\", \"ρθ\", -9.95373884733410818626e+00, -4.05672848247591079102e-07, -4.98722877855243940104e+00, 2.97859288054288384728e+00 ],\n],\n parr,\n)\n\nfourth_order = (\n[\n [ \"state\", \"ρ\", 9.95377495534709000324e-01, 9.99992951378667060958e-01, 9.98321272635789513927e-01, 1.50722816639464068413e-03 ],\n [ \"state\", \"ρu[1]\", -1.22778750186973996879e-14, 1.95893093780676565332e-14, 1.73744809621522011733e-15, 4.23044883848200063451e-15 ],\n [ \"state\", \"ρu[2]\", -1.62568523313266733927e-14, 2.71071963864711276060e-14, 1.75189805678757112174e-15, 4.33763908149213610479e-15 ],\n [ \"state\", \"ρu[3]\", -1.66252523985503642377e-03, 5.55021884639224168229e-08, -8.13842490777055447768e-04, 4.76414940515918817511e-04 ],\n [ \"state\", \"ρθ\", -9.95373884733410818626e+00, -4.05672848142820462126e-07, -4.98722877855243940104e+00, 2.97859288054288384728e+00 ],\n],\n parr,\n)\n\n#! format: on\n\nrefVals = (; second_order, fourth_order, second_order_flat, fourth_order_flat)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/run_bickley_jet.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n########\n# Setup physical and numerical domains\n########\nΩˣ = IntervalDomain(-2π, 2π, periodic = true)\nΩʸ = IntervalDomain(-2π, 2π, periodic = true)\nΩᶻ = IntervalDomain(-2π, 2π, periodic = true)\n\ngrid = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 13,\n polynomial_order = 4,\n overintegration_order = 1,\n)\n\n########\n# Define timestepping parameters\n########\nstart_time = 0\nend_time = 200.0\nΔt = 0.004\nmethod = SSPRK22Heuns\n\ntimestepper = TimeStepper(method = method, timestep = Δt)\n\ncallbacks = (Info(), StateCheck(10))\n\n########\n# Define physical parameters and parameterizations\n########\nparameters = (\n ϵ = 0.1, # perturbation size for initial condition\n l = 0.5, # Gaussian width\n k = 0.5, # Sinusoidal wavenumber\n ρₒ = 1, # reference density\n cₛ = sqrt(10), # sound speed\n)\n\nphysics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0, κ = 0),\n coriolis = nothing,\n buoyancy = nothing,\n)\n\n########\n# Define initial conditions\n########\n\n# The Bickley jet\nU₀(p, x, y, z) = cosh(y)^(-2)\nV₀(p, x, y, z) = 0\nW₀(p, x, y, z) = 0\n\n# Slightly off-center vortical perturbations\nΨ₁(p, x, y, z) =\n exp(-(y + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * x) * cos(p.k * y)\nΨ₂(p, x, y, z) =\n exp(-(z + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * y) * cos(p.k * z)\n\n# Vortical velocity fields (u, v, w) = (-∂ʸ, +∂ˣ, 0) Ψ₁ + (0, -∂ᶻ, +∂ʸ)Ψ₂ \nu₀(p, x, y, z) =\n Ψ₁(p, x, y, z) * (p.k * tan(p.k * y) + y / (p.l^2) + 1 / (10 * p.l))\nv₀(p, x, y, z) =\n Ψ₂(p, x, y, z) * (p.k * tan(p.k * z) + z / (p.l^2) + 1 / (10 * p.l)) -\n Ψ₁(p, x, y, z) * p.k * tan(p.k * x)\nw₀(p, x, y, z) = -Ψ₂(p, x, y, z) * p.k * tan(p.k * y)\nθ₀(p, x, y, z) = sin(p.k * y)\n\nρ₀(p, x, y, z) = p.ρₒ\nρu₀(p, x...) = ρ₀(p, x...) * (p.ϵ * u₀(p, x...) + U₀(p, x...))\nρv₀(p, x...) = ρ₀(p, x...) * (p.ϵ * v₀(p, x...) + V₀(p, x...))\nρw₀(p, x...) = ρ₀(p, x...) * (p.ϵ * w₀(p, x...) + W₀(p, x...))\nρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\nρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\ninitial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n########\n# Create the things\n########\n\nmodel = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = NamedTuple(),\n parameters = parameters,\n)\n\nsimulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n)\n\n########\n# Run the model\n########\n\ntic = Base.time()\n\nevolve!(simulation, model)\n\ntoc = Base.time()\ntime = toc - tic\nprintln(time)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/run_box.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n########\n# Setup physical and numerical domains\n########\nΩˣ = IntervalDomain(-2π, 2π, periodic = true)\nΩʸ = IntervalDomain(-2π, 2π, periodic = true)\nΩᶻ = IntervalDomain(-2π, 2π, periodic = false)\n\ngrid = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 8,\n polynomial_order = 1,\n overintegration_order = 1,\n)\n\n########\n# Define timestepping parameters\n########\nstart_time = 0\nend_time = 200.0\nΔt = 0.05\nmethod = SSPRK22Heuns\n\ntimestepper = TimeStepper(method = method, timestep = Δt)\n\ncallbacks = (Info(), StateCheck(10))\n\n########\n# Define physical parameters and parameterizations\n########\nparameters = (\n ρₒ = 1, # reference density\n cₛ = sqrt(10), # sound speed\n)\n\nphysics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 1e-2, κ = 1e-2),\n coriolis = nothing,\n buoyancy = Buoyancy{Float64}(α = 2e-4, g = 10),\n)\n\n########\n# Define boundary conditions\n########\nρu_bcs = (\n bottom = Impenetrable(NoSlip()),\n top = Impenetrable(MomentumFlux(\n flux = (p, state, aux, t) -> (@SVector [p.τ / state.ρ, -0, -0]),\n params = (τ = 0.01,),\n )),\n)\nρθ_bcs = (\n bottom = Insulating(),\n top = TemperatureFlux(flux = (p, state, aux, t) -> (p.Q)),\n params = (Q = 0.1,), # positive means removing heat\n)\nBC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n########\n# Define initial conditions\n########\n\nρ₀(p, x, y, z) = p.ρₒ\nρu₀(p, x...) = ρ₀(p, x...) * -0\nρv₀(p, x...) = ρ₀(p, x...) * -0\nρw₀(p, x...) = ρ₀(p, x...) * -0\nρθ₀(p, x...) = ρ₀(p, x...) * 5\n\nρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\ninitial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n########\n# Create the things\n########\n\nmodel = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n)\n\nsimulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n)\n\n########\n# Run the model\n########\n\ntic = Base.time()\n\nevolve!(simulation, model)\n\ntoc = Base.time()\ntime = toc - tic\nprintln(time)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/run_taylor_green_vortex.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n########\n# Setup physical and numerical domains\n########\nΩˣ = IntervalDomain(-2π, 2π, periodic = true)\nΩʸ = IntervalDomain(-2π, 2π, periodic = true)\nΩᶻ = IntervalDomain(-2π, 2π, periodic = true)\n\ngrid = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 16,\n polynomial_order = 1,\n overintegration_order = 1,\n)\n\n########\n# Define timestepping parameters\n########\nstart_time = 0\nend_time = 200.0\nΔt = 0.01\nmethod = SSPRK22Heuns\n\ntimestepper = TimeStepper(method = method, timestep = Δt)\n\ncallbacks = (Info(), StateCheck(10))\n\n########\n# Define physical parameters and parameterizations\n########\nparameters = (\n Uₒ = 1, # reference velocity\n ρₒ = 1, # reference density\n cₛ = sqrt(10), # sound speed\n)\n\nphysics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 1e-3, κ = 1e-3),\n coriolis = nothing,\n buoyancy = nothing,\n)\n\n########\n# Define initial conditions\n########\n\nu₀(p, x, y, z) = p.Uₒ * sin(x) * cos(y) * cos(z)\nv₀(p, x, y, z) = -p.Uₒ * cos(x) * sin(y) * cos(z)\nw₀(p, x, y, z) = -0\nθ₀(p, x, y, z) = sin(0.5 * z)\n\nρ₀(p, x, y, z) = p.ρₒ\nρu₀(p, x...) = ρ₀(p, x...) * u₀(p, x...)\nρv₀(p, x...) = ρ₀(p, x...) * v₀(p, x...)\nρw₀(p, x...) = ρ₀(p, x...) * w₀(p, x...)\nρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\nρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\ninitial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n########\n# Create the things\n########\n\nmodel = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = NamedTuple(),\n parameters = parameters,\n)\n\nsimulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n)\n\n########\n# Run the model\n########\n\ntic = Base.time()\n\nevolve!(simulation, model)\n\ntoc = Base.time()\ntime = toc - tic\nprintln(time)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/test_bickley_jet.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"refvals_bickley_jet.jl\")\n\n ########\n # Setup physical and numerical domains\n ########\n Ωˣ = IntervalDomain(-2π, 2π, periodic = true)\n Ωʸ = IntervalDomain(-2π, 2π, periodic = true)\n Ωᶻ = IntervalDomain(-2π, 2π, periodic = true)\n\n first_order = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 32,\n polynomial_order = 1,\n overintegration_order = 1,\n )\n\n fourth_order = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 13,\n polynomial_order = 4,\n overintegration_order = 1,\n )\n\n grids = Dict(\"first_order\" => first_order, \"fourth_order\" => fourth_order)\n\n ########\n # Define timestepping parameters\n ########\n start_time = 0\n end_time = 100.0\n Δt = 0.004\n method = SSPRK22Heuns\n\n timestepper = TimeStepper(method = method, timestep = Δt)\n\n callbacks = (Info(), StateCheck(10))\n\n ########\n # Define physical parameters and parameterizations\n ########\n parameters = (\n ϵ = 0.1, # perturbation size for initial condition\n l = 0.5, # Gaussian width\n k = 0.5, # Sinusoidal wavenumber\n ρₒ = 1, # reference density\n cₛ = sqrt(10), # sound speed\n )\n\n physics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0, κ = 0),\n coriolis = nothing,\n buoyancy = nothing,\n )\n\n ########\n # Define initial conditions\n ########\n\n # The Bickley jet\n U₀(p, x, y, z) = cosh(y)^(-2)\n V₀(p, x, y, z) = 0\n W₀(p, x, y, z) = 0\n\n # Slightly off-center vortical perturbations\n Ψ₁(p, x, y, z) =\n exp(-(y + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * x) * cos(p.k * y)\n Ψ₂(p, x, y, z) =\n exp(-(z + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * y) * cos(p.k * z)\n\n # Vortical velocity fields (u, v, w) = (-∂ʸ, +∂ˣ, 0) Ψ₁ + (0, -∂ᶻ, +∂ʸ)Ψ₂ \n u₀(p, x, y, z) =\n Ψ₁(p, x, y, z) * (p.k * tan(p.k * y) + y / (p.l^2) + 1 / (10 * p.l))\n v₀(p, x, y, z) =\n Ψ₂(p, x, y, z) * (p.k * tan(p.k * z) + z / (p.l^2) + 1 / (10 * p.l)) -\n Ψ₁(p, x, y, z) * p.k * tan(p.k * x)\n w₀(p, x, y, z) = -Ψ₂(p, x, y, z) * p.k * tan(p.k * y)\n θ₀(p, x, y, z) = sin(p.k * y)\n\n ρ₀(p, x, y, z) = p.ρₒ\n ρu₀(p, x...) = ρ₀(p, x...) * (p.ϵ * u₀(p, x...) + U₀(p, x...))\n ρv₀(p, x...) = ρ₀(p, x...) * (p.ϵ * v₀(p, x...) + V₀(p, x...))\n ρw₀(p, x...) = ρ₀(p, x...) * (p.ϵ * w₀(p, x...) + W₀(p, x...))\n ρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\n ρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\n initial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n for (key, grid) in grids\n @testset \"$(key)\" begin\n model = SpatialModel(\n balance_law = Fluid3D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = NamedTuple(),\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n evolve!(\n simulation,\n model;\n refDat = getproperty(refVals, Symbol(key)),\n )\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/three_dimensional/test_buoyancy.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"ThreeDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"refvals_buoyancy.jl\")\n\n ########\n # Setup physical and numerical domains\n ########\n Ωˣ = IntervalDomain(-2π, 2π, periodic = true)\n Ωʸ = IntervalDomain(-2π, 2π, periodic = true)\n Ωᶻ = IntervalDomain(0, 4π, periodic = false)\n\n second_order = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 5,\n polynomial_order = 2,\n overintegration_order = 1,\n )\n\n fourth_order = DiscretizedDomain(\n Ωˣ × Ωʸ × Ωᶻ;\n elements = 3,\n polynomial_order = 4,\n overintegration_order = 1,\n )\n\n grids = Dict(\"second_order\" => second_order, \"fourth_order\" => fourth_order)\n\n ########\n # Define timestepping parameters\n ########\n start_time = 0\n end_time = 0.1\n Δt = 0.001\n method = SSPRK22Heuns\n\n timestepper = TimeStepper(method = method, timestep = Δt)\n\n callbacks = (\n Info(),\n # VTKState(; iteration = 1, filepath = \"output/buoyancy\"),\n StateCheck(10),\n )\n\n ########\n # Define physical parameters and parameterizations\n ########\n Lˣ, Lʸ, Lᶻ = length(Ωˣ × Ωʸ × Ωᶻ)\n parameters = (\n ρₒ = 1, # reference density\n cₛ = sqrt(10), # sound speed\n α = 1e-4, # thermal expansion coefficient\n g = 10, # gravity\n θₒ = 10, # initial temperature value\n Lˣ = Lˣ,\n Lʸ = Lʸ,\n H = Lᶻ,\n )\n\n physics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0, κ = 0),\n coriolis = nothing,\n buoyancy = Buoyancy{Float64}(α = parameters.α, g = parameters.g),\n )\n\n ########\n # Define initial conditions\n ########\n\n u₀(p, x, y, z) = -0\n v₀(p, x, y, z) = -0\n w₀(p, x, y, z) = -0\n θ₀(p, x, y, z) = -p.θₒ * (1 - z / 4π)\n\n ρ₀(p, x, y, z) =\n p.ρₒ * (1 - (p.α * p.g / p.cₛ^2) / 2 * (-p.θₒ * (1 - z / 4π))^2)\n ρu₀(p, x...) = ρ₀(p, x...) * u₀(p, x...)\n ρv₀(p, x...) = ρ₀(p, x...) * v₀(p, x...)\n ρw₀(p, x...) = ρ₀(p, x...) * w₀(p, x...)\n ρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\n ρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...), ρw₀(p, x...)]\n initial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n orientations = Dict(\n \"\" => ClimateMachine.Orientations.NoOrientation(),\n \"_flat\" => ClimateMachine.Orientations.FlatOrientation(),\n )\n\n for (key1, grid) in grids\n for (key2, orientation) in orientations\n key = key1 * key2\n\n println(\"running \", key)\n\n local_physics = FluidPhysics(;\n orientation = orientation,\n advection = physics.advection,\n dissipation = physics.dissipation,\n coriolis = physics.coriolis,\n buoyancy = physics.buoyancy,\n )\n\n @testset \"$(key)\" begin\n model = SpatialModel(\n balance_law = Fluid3D(),\n physics = local_physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = NamedTuple(),\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n evolve!(\n simulation,\n model,\n refDat = getproperty(refVals, Symbol(key)),\n )\n end\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/TwoDimensionalCompressibleNavierStokesEquations.jl", "content": "include(\"../shared_source/boilerplate.jl\")\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n init_state_prognostic!,\n init_state_auxiliary!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_conditions,\n boundary_state!\nimport ClimateMachine.DGMethods: DGModel\nimport ClimateMachine.NumericalFluxes: numerical_flux_first_order!\n\n\"\"\"\n TwoDimensionalCompressibleNavierStokesEquations <: BalanceLaw\n\nA `BalanceLaw` for shallow water modeling.\n\nwrite out the equations here\n\n# Usage\n\n TwoDimensionalCompressibleNavierStokesEquations()\n\n\"\"\"\nabstract type AbstractFluid2D <: AbstractFluid end\nstruct Fluid2D <: AbstractFluid2D end\nstruct TwoDimensionalCompressibleNavierStokesEquations{\n I,\n D,\n A,\n T,\n C,\n F,\n BC,\n FT,\n} <: AbstractFluid2D\n initial_value_problem::I\n domain::D\n advection::A\n turbulence::T\n coriolis::C\n forcing::F\n boundary_conditions::BC\n g::FT\n c::FT\n function TwoDimensionalCompressibleNavierStokesEquations{FT}(\n initial_value_problem::I,\n domain::D,\n advection::A,\n turbulence::T,\n coriolis::C,\n forcing::F,\n boundary_conditions::BC;\n g = FT(10), # m/s²\n c = FT(0), #m/s\n ) where {FT <: AbstractFloat, I, D, A, T, C, F, BC}\n return new{I, D, A, T, C, F, BC, FT}(\n initial_value_problem,\n domain,\n advection,\n turbulence,\n coriolis,\n forcing,\n boundary_conditions,\n g,\n c,\n )\n end\nend\nCNSE2D = TwoDimensionalCompressibleNavierStokesEquations\n\nfunction vars_state(m::CNSE2D, ::Prognostic, T)\n @vars begin\n ρ::T\n ρu::SVector{2, T}\n ρθ::T\n end\nend\n\nfunction init_state_prognostic!(m::CNSE2D, state::Vars, aux::Vars, localgeo, t)\n cnse_init_state!(m, state, aux, localgeo, t)\nend\n\n# default initial state if IVP == nothing\nfunction cnse_init_state!(model::CNSE2D, state, aux, localgeo, t)\n ρ = 1\n\n state.ρ = ρ\n state.ρu = ρ * @SVector [-0, -0]\n state.ρθ = ρ\n\n return nothing\nend\n\n# user defined initial state\nfunction cnse_init_state!(\n model::CNSE2D{<:InitialValueProblem},\n state,\n aux,\n localgeo,\n t,\n)\n x = aux.x\n y = aux.y\n z = aux.z\n\n params = model.initial_value_problem.params\n ic = model.initial_value_problem.initial_conditions\n\n state.ρ = ic.ρ(params, x, y, z)\n state.ρu = ic.ρu(params, x, y, z)\n state.ρθ = ic.ρθ(params, x, y, z)\n\n return nothing\nend\n\nfunction vars_state(m::CNSE2D, ::Auxiliary, T)\n @vars begin\n x::T\n y::T\n z::T\n end\nend\n\nfunction init_state_auxiliary!(\n model::CNSE2D,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_state_auxiliary!(\n model,\n (model, aux, tmp, geom) -> cnse_init_aux!(model, aux, geom),\n state_auxiliary,\n grid,\n direction,\n )\nend\n\nfunction cnse_init_aux!(::CNSE2D, aux, geom)\n @inbounds begin\n aux.x = geom.coord[1]\n aux.y = geom.coord[2]\n aux.z = geom.coord[3]\n end\n\n return nothing\nend\n\nfunction vars_state(m::CNSE2D, ::Gradient, T)\n @vars begin\n ∇u::SVector{2, T}\n ∇θ::T\n end\nend\n\nfunction compute_gradient_argument!(\n model::CNSE2D,\n grad::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n compute_gradient_argument!(model.turbulence, grad, state, aux, t)\nend\n\ncompute_gradient_argument!(::LinearDrag, _...) = nothing\n\n@inline function compute_gradient_argument!(\n ::ConstantViscosity,\n grad::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρθ = state.ρθ\n\n u = ρu / ρ\n θ = ρθ / ρ\n\n grad.∇u = u\n grad.∇θ = θ\n\n return nothing\nend\n\nfunction vars_state(m::CNSE2D, ::GradientFlux, T)\n @vars begin\n ν∇u::SMatrix{3, 2, T, 6}\n κ∇θ::SVector{3, T}\n end\nend\n\nfunction compute_gradient_flux!(\n model::CNSE2D,\n gradflux::Vars,\n grad::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n compute_gradient_flux!(\n model,\n model.turbulence,\n gradflux,\n grad,\n state,\n aux,\n t,\n )\nend\n\ncompute_gradient_flux!(::CNSE2D, ::LinearDrag, _...) = nothing\n\n@inline function compute_gradient_flux!(\n ::CNSE2D,\n turb::ConstantViscosity,\n gradflux::Vars,\n grad::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ν = Diagonal(@SVector [turb.ν, turb.ν, -0])\n κ = Diagonal(@SVector [turb.κ, turb.κ, -0])\n\n gradflux.ν∇u = -ν * grad.∇u\n gradflux.κ∇θ = -κ * grad.∇θ\n\n return nothing\nend\n\n@inline function flux_first_order!(\n model::CNSE2D,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n\n ρ = state.ρ\n ρu = @SVector [state.ρu[1], state.ρu[2], -0]\n ρθ = state.ρθ\n\n ρₜ = flux.ρ\n ρuₜ = flux.ρu\n θₜ = flux.ρθ\n\n g = model.g\n\n Iʰ = @SMatrix [\n 1 -0\n -0 1\n -0 -0\n ]\n\n flux.ρ += ρu\n flux.ρu += g * ρ^2 * Iʰ / 2\n\n advective_flux!(model, model.advection, flux, state, aux, t)\n\n return nothing\nend\n\nadvective_flux!(::CNSE2D, ::Nothing, _...) = nothing\n\n@inline function advective_flux!(\n ::CNSE2D,\n ::NonLinearAdvectionTerm,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρv = @SVector [state.ρu[1], state.ρu[2], -0]\n ρθ = state.ρθ\n\n flux.ρu += ρv ⊗ ρu / ρ\n flux.ρθ += ρv * ρθ / ρ\n\n return nothing\nend\n\nfunction flux_second_order!(\n model::CNSE2D,\n flux::Grad,\n state::Vars,\n gradflux::Vars,\n ::Vars,\n aux::Vars,\n t::Real,\n)\n flux_second_order!(model, model.turbulence, flux, state, gradflux, aux, t)\nend\n\nflux_second_order!(::CNSE2D, ::LinearDrag, _...) = nothing\n\n@inline function flux_second_order!(\n ::CNSE2D,\n ::ConstantViscosity,\n flux::Grad,\n state::Vars,\n gradflux::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρu += gradflux.ν∇u\n flux.ρθ += gradflux.κ∇θ\n\n return nothing\nend\n\n@inline function source!(\n model::CNSE2D,\n source::Vars,\n state::Vars,\n gradflux::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n coriolis_force!(model, model.coriolis, source, state, aux, t)\n forcing_term!(model, model.forcing, source, state, aux, t)\n linear_drag!(model, model.turbulence, source, state, aux, t)\n\n return nothing\nend\n\ncoriolis_force!(::CNSE2D, ::Nothing, _...) = nothing\n\n@inline function coriolis_force!(\n model::CNSE2D,\n coriolis::fPlaneCoriolis,\n source,\n state,\n aux,\n t,\n)\n ρu = @SVector [state.ρu[1], state.ρu[2], -0]\n\n # f × u\n f = [-0, -0, coriolis_parameter(model, coriolis, aux.coords)]\n id = @SVector [1, 2]\n fxρu = (f × ρu)[id]\n\n source.ρu -= fxρu\n\n return nothing\nend\n\nforcing_term!(::CNSE2D, ::Nothing, _...) = nothing\n\n@inline function forcing_term!(\n model::CNSE2D,\n forcing::KinematicStress,\n source,\n state,\n aux,\n t,\n)\n source.ρu += kinematic_stress(model, forcing, aux.coords)\n\n return nothing\nend\n\nlinear_drag!(::CNSE2D, ::ConstantViscosity, _...) = nothing\n\n@inline function linear_drag!(::CNSE2D, turb::LinearDrag, source, state, aux, t)\n source.ρu -= turb.λ * state.ρu\n\n return nothing\nend\n\n@inline wavespeed(m::CNSE2D, _...) = m.c\n\nroe_average(ρ⁻, ρ⁺, var⁻, var⁺) =\n (sqrt(ρ⁻) * var⁻ + sqrt(ρ⁺) * var⁺) / (sqrt(ρ⁻) + sqrt(ρ⁺))\n\nfunction numerical_flux_first_order!(\n ::RoeNumericalFlux,\n model::CNSE2D,\n fluxᵀn::Vars{S},\n n⁻::SVector,\n state⁻::Vars{S},\n aux⁻::Vars{A},\n state⁺::Vars{S},\n aux⁺::Vars{A},\n t,\n direction,\n) where {S, A}\n numerical_flux_first_order!(\n CentralNumericalFluxFirstOrder(),\n model,\n fluxᵀn,\n n⁻,\n state⁻,\n aux⁻,\n state⁺,\n aux⁺,\n t,\n direction,\n )\n\n FT = eltype(fluxᵀn)\n\n # constants and normal vectors\n g = model.g\n @inbounds nˣ = n⁻[1]\n @inbounds nʸ = n⁻[2]\n\n # get minus side states\n ρ⁻ = state⁻.ρ\n @inbounds ρu⁻ = state⁻.ρu[1]\n @inbounds ρv⁻ = state⁻.ρu[2]\n ρθ⁻ = state⁻.ρθ\n\n u⁻ = ρu⁻ / ρ⁻\n v⁻ = ρv⁻ / ρ⁻\n θ⁻ = ρθ⁻ / ρ⁻\n\n # get plus side states\n ρ⁺ = state⁺.ρ\n @inbounds ρu⁺ = state⁺.ρu[1]\n @inbounds ρv⁺ = state⁺.ρu[2]\n ρθ⁺ = state⁺.ρθ\n\n u⁺ = ρu⁺ / ρ⁺\n v⁺ = ρv⁺ / ρ⁺\n θ⁺ = ρθ⁺ / ρ⁺\n\n # averages for roe fluxes\n ρ = (ρ⁺ + ρ⁻) / 2\n ρu = (ρu⁺ + ρu⁻) / 2\n ρv = (ρv⁺ + ρv⁻) / 2\n ρθ = (ρθ⁺ + ρθ⁻) / 2\n\n u = roe_average(ρ⁻, ρ⁺, u⁻, u⁺)\n v = roe_average(ρ⁻, ρ⁺, v⁻, v⁺)\n θ = roe_average(ρ⁻, ρ⁺, θ⁻, θ⁺)\n\n # normal and tangent velocities\n uₙ = nˣ * u + nʸ * v\n uₚ = nˣ * v - nʸ * u\n\n # differences for difference vector\n Δρ = ρ⁺ - ρ⁻\n Δρu = ρu⁺ - ρu⁻\n Δρv = ρv⁺ - ρv⁻\n Δρθ = ρθ⁺ - ρθ⁻\n\n Δφ = @SVector [Δρ, Δρu, Δρv, Δρθ]\n\n \"\"\"\n # jacobian\n ∂F∂φ = [\n 0 nˣ nʸ 0\n (nˣ * c^2 - u * uₙ) (uₙ + nˣ * u) (nʸ * u) 0\n (nʸ * c^2 - v * uₙ) (nˣ * v) (uₙ + nʸ * v) 0\n (-θ * uₙ) (nˣ * θ) (nʸ * θ) uₙ\n ]\n\n # eigen decomposition\n λ, R = eigen(∂F∂φ)\n \"\"\"\n\n # eigen values matrix\n c = sqrt(g * ρ)\n λ = @SVector [uₙ, uₙ + c, uₙ - c, uₙ]\n Λ = Diagonal(abs.(λ))\n\n\n # eigenvector matrix\n R = @SMatrix [\n 0 1 1 0\n -nʸ (u+nˣ * c) (u-nˣ * c) 0\n nˣ (v+nʸ * c) (v-nʸ * c) 0\n 0 θ θ 1\n ]\n\n # inverse of eigenvector matrix\n R⁻¹ = @SMatrix [\n -uₚ -nʸ nˣ 0\n (c - uₙ)/(2c) nˣ/(2c) nʸ/(2c) 0\n (c + uₙ)/(2c) -nˣ/(2c) -nʸ/(2c) 0\n -θ 0 0 1\n ]\n\n # @test norm(R⁻¹ * R - I) ≈ 0\n\n # actually calculate flux\n # parent(fluxᵀn) .-= R * Λ * (R \\ Δφ) / 2\n parent(fluxᵀn) .-= R * Λ * R⁻¹ * Δφ / 2\n\n return nothing\nend\n\nboundary_conditions(model::CNSE2D) = model.boundary_conditions\n\n\"\"\"\n boundary_state!(nf, ::CNSE2D, args...)\n\napplies boundary conditions for the hyperbolic fluxes\ndispatches to a function in CNSEBoundaryConditions\n\"\"\"\n@inline function boundary_state!(nf, bc, model::CNSE2D, args...)\n return _cnse_boundary_state!(nf, bc, model, args...)\nend\n\n\"\"\"\n CNSE_boundary_state!(nf, bc::FluidBC, ::CNSE2D)\n\nsplits boundary condition application into velocity\n\"\"\"\n@inline function cnse_boundary_state!(nf, bc::FluidBC, m::CNSE2D, args...)\n return cnse_boundary_state!(nf, bc.momentum, m, m.turbulence, args...)\n return cnse_boundary_state!(nf, bc.temperature, m, args...)\nend\n\ninclude(\"bc_momentum.jl\")\ninclude(\"bc_tracer.jl\")\n\n\"\"\"\nSTUFF FOR ANDRE'S WRAPPERS\n\"\"\"\n\nfunction get_boundary_conditions(\n model::SpatialModel{BL},\n) where {BL <: AbstractFluid2D}\n bcs = model.boundary_conditions\n\n west_east = (check_bc(bcs, :west), check_bc(bcs, :east))\n south_north = (check_bc(bcs, :south), check_bc(bcs, :north))\n\n return (west_east..., south_north...)\nend\n\nfunction DGModel(\n model::SpatialModel{BL};\n initial_conditions = nothing,\n) where {BL <: AbstractFluid2D}\n params = model.parameters\n physics = model.physics\n\n Lˣ, Lʸ = length(model.grid.domain)\n bcs = get_boundary_conditions(model)\n FT = eltype(model.grid.numerical.vgeo)\n\n if !isnothing(initial_conditions)\n initial_conditions = InitialValueProblem(params, initial_conditions)\n end\n\n balance_law = CNSE2D{FT}(\n initial_conditions,\n (Lˣ, Lʸ),\n physics.advection,\n physics.dissipation,\n physics.coriolis,\n nothing,\n bcs,\n c = params.c,\n g = params.g,\n )\n\n numerical_flux_first_order = model.numerics.flux # should be a function\n\n rhs = DGModel(\n balance_law,\n model.grid.numerical,\n numerical_flux_first_order,\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n return rhs\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/bc_momentum.jl", "content": "\"\"\"\n cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{FreeSlip}, ::CNSE2D)\n\napply free slip boundary condition for momentum\nsets reflective ghost point\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{FreeSlip},\n ::CNSE2D,\n ::TurbulenceClosure,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n\n ρu⁻ = @SVector [state⁻.ρu[1], state⁻.ρu[2], -0]\n ρu⁺ = ρu⁻ - 2 * n⁻ ⋅ ρu⁻ .* SVector(n⁻)\n\n state⁺.ρu = @SVector [ρu⁺[1], ρu⁺[2]]\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxGradient, NumericalFluxSecondOrder}, ::Impenetrable{FreeSlip}, ::CNSE2D)\n\nno second order flux computed for linear drag\n\"\"\"\ncnse_boundary_state!(\n ::Union{NumericalFluxGradient, NumericalFluxSecondOrder},\n ::MomentumBC,\n ::CNSE2D,\n ::LinearDrag,\n _...,\n) = nothing\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{FreeSlip}, ::CNSE2D)\n\napply free slip boundary condition for momentum\nsets non-reflective ghost point\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{FreeSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n\n ρu⁻ = @SVector [state⁻.ρu[1], state⁻.ρu[2], -0]\n ρu⁺ = ρu⁻ - n⁻ ⋅ ρu⁻ .* SVector(n⁻)\n\n state⁺.ρu = @SVector [ρu⁺[1], ρu⁺[2]]\n\n return nothing\nend\n\n\"\"\"\n shallow_normal_boundary_flux_second_order!(::NumericalFluxSecondOrder, ::Impenetrable{FreeSlip}, ::CNSE2D)\n\napply free slip boundary condition for momentum\napply zero numerical flux in the normal direction\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{FreeSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxFirstOrder, ::Impenetrable{NoSlip}, ::CNSE2D)\n\napply no slip boundary condition for momentum\nsets reflective ghost point\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxFirstOrder,\n ::Impenetrable{NoSlip},\n ::CNSE2D,\n ::TurbulenceClosure,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρ = state⁻.ρ\n state⁺.ρu = -state⁻.ρu\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxGradient, ::Impenetrable{NoSlip}, ::CNSE2D)\n\napply no slip boundary condition for momentum\nset numerical flux to zero for U\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxGradient,\n ::Impenetrable{NoSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n args...,\n)\n FT = eltype(state⁺)\n state⁺.ρu = @SVector zeros(FT, 2)\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{NoSlip}, ::CNSE2D)\n\napply no slip boundary condition for momentum\nsets ghost point to have no numerical flux on the boundary for U\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{NoSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = -state⁻.ρu\n gradflux⁺.ν∇u = gradflux⁻.ν∇u\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{FreeSlip}, ::CNSE2D)\n\nno mass boundary condition for penetrable\n\"\"\"\ncnse_boundary_state!(\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{FreeSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n _...,\n) = nothing\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{FreeSlip}, ::CNSE2D)\n\napply free slip boundary condition for momentum\napply zero numerical flux in the normal direction\n\"\"\"\nfunction cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Penetrable{FreeSlip},\n ::CNSE2D,\n ::ConstantViscosity,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n args...,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * (@SVector [-0, -0])'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Impenetrable{MomentumFlux}, ::HBModel)\n\napply kinematic stress boundary condition for momentum\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Impenetrable{<:MomentumFlux},\n model::CNSE2D,\n turb::TurbulenceClosure,\n args...,\n)\n return cnse_boundary_state!(\n nf,\n Impenetrable(FreeSlip()),\n model,\n turb,\n args...,\n )\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Impenetrable{MomentumFlux}, ::HBModel)\n\napply kinematic stress boundary condition for momentum\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Impenetrable{<:MomentumFlux},\n model::CNSE2D,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Penetrable{MomentumFlux}, ::HBModel)\n\napply kinematic stress boundary condition for momentum\napplies free slip conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Penetrable{<:MomentumFlux},\n model::CNSE2D,\n turb::TurbulenceClosure,\n args...,\n)\n return cnse_boundary_state!(\n nf,\n Penetrable(FreeSlip()),\n model,\n turb,\n args...,\n )\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Penetrable{MomentumFlux}, ::HBModel)\n\napply kinematic stress boundary condition for momentum\nsets ghost point to have specified flux on the boundary for ν∇u\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n bc::Penetrable{<:MomentumFlux},\n shallow::CNSE2D,\n ::TurbulenceClosure,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρu = state⁻.ρu\n gradflux⁺.ν∇u = n⁻ * bc.drag(state⁻, aux⁻, t)'\n\n return nothing\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/bc_tracer.jl", "content": "\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::Insulating, ::HBModel)\n\napply insulating boundary condition for temperature\nsets transmissive ghost point\n\"\"\"\nfunction cnse_boundary_state!(\n ::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::Insulating,\n ::CNSE2D,\n state⁺,\n aux⁺,\n n⁻,\n state⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::Insulating, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have no numerical flux on the boundary for κ∇θ\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n ::Insulating,\n ::CNSE2D,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n gradflux⁺.κ∇θ = n⁻ * -0\n\n return nothing\nend\n\n\"\"\"\n cnse_boundary_state!(::Union{NumericalFluxFirstOrder, NumericalFluxGradient}, ::TemperatureFlux, ::HBModel)\n\napply temperature flux boundary condition for velocity\napplies insulating conditions for first-order and gradient fluxes\n\"\"\"\nfunction cnse_boundary_state!(\n nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},\n ::TemperatureFlux,\n model::CNSE2D,\n args...,\n)\n return cnse_boundary_state!(nf, Insulating(), model, args...)\nend\n\n\"\"\"\n cnse_boundary_state!(::NumericalFluxSecondOrder, ::TemperatureFlux, ::HBModel)\n\napply insulating boundary condition for velocity\nsets ghost point to have specified flux on the boundary for κ∇θ\n\"\"\"\n@inline function cnse_boundary_state!(\n ::NumericalFluxSecondOrder,\n bc::TemperatureFlux,\n model::CNSE2D,\n state⁺,\n gradflux⁺,\n aux⁺,\n n⁻,\n state⁻,\n gradflux⁻,\n aux⁻,\n t,\n)\n state⁺.ρθ = state⁻.ρθ\n gradflux⁺.κ∇θ = n⁻ * bc(state⁻, aux⁻, t)\n\n return nothing\nend\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/refvals_bickley_jet.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\nparr = [\n [\"state\", :ρ, 12, 12, 12, 12],\n [\"state\", \"ρu[1]\", 12, 11, 12, 12],\n [\"state\", \"ρu[2]\", 12, 12, 11, 12],\n [\"state\", :ρθ, 11, 12, 10, 12],\n]\n\n#! format: off\nrusanov_periodic = (\n[\n [ \"state\", \"ρ\", 9.76557021043823247908e-01, 1.00785052430958410596e+00, 1.00000023257151737788e+00, 5.40633258005515995870e-03 ],\n [ \"state\", \"ρu[1]\", -2.08235364819134516345e-01, 5.42021330588998817568e-01, 1.59191748432813973135e-01, 1.71679977050671228600e-01 ],\n [ \"state\", \"ρu[2]\", -4.44598929671790876750e-01, 3.96193527502323006306e-01, -9.19857600654994111977e-06, 2.23835345394622464710e-01 ],\n [ \"state\", \"ρθ\", -1.05926174767863745529e+00, 1.78285318416958671328e+00, 3.38466687536443516793e-03, 2.92797512839404083795e-01 ],\n],\n parr,\n)\n\nroeflux_periodic = (\n[\n [ \"state\", \"ρ\", 9.76408216346011381681e-01, 1.00685186827167338919e+00, 1.00000020920188337215e+00, 5.28437765352953326553e-03 ],\n [ \"state\", \"ρu[1]\", -2.51377173896010663867e-01, 5.33033893482195431091e-01, 1.59154423274343037598e-01, 1.96922130305308085152e-01 ],\n [ \"state\", \"ρu[2]\", -4.04321943123139071474e-01, 3.81969301036260144855e-01, -2.87040331153030903177e-05, 2.02392226536863034658e-01 ],\n [ \"state\", \"ρθ\", -1.21186199231530267184e+00, 9.21722158564563742722e-01, 2.34662081315961668082e-03, 2.61922031216240580598e-01 ],\n],\n parr,\n)\n\nrusanov = (\n[\n [ \"state\", \"ρ\", 9.76844648039728813416e-01, 1.00662311214286837036e+00, 1.00000043240984126669e+00, 6.05641827304116090597e-03 ],\n [ \"state\", \"ρu[1]\", -6.93213807814652027695e-01, 5.89769972499923134102e-01, 1.42589259313275984464e-01, 2.32576446220999155656e-01 ],\n [ \"state\", \"ρu[2]\", -5.73828377062444716650e-01, 5.07664260030296077275e-01, -4.23962044125134052130e-04, 1.96899492339057097245e-01 ],\n [ \"state\", \"ρθ\", -3.65527466807722500874e+00, 4.07350442876234186684e+00, 1.27099405499731775426e-02, 5.02926135539447205502e-01 ],\n],\n parr,\n)\n\nroeflux = (\n[\n [ \"state\", \"ρ\", 9.72455511248961124160e-01, 1.00730359035806360524e+00, 1.00000004243418749716e+00, 6.29619947553510493632e-03 ],\n [ \"state\", \"ρu[1]\", -3.40062758027173561715e-01, 5.98382642258275421199e-01, 1.61534160931526560301e-01, 1.83683987730847486652e-01 ],\n [ \"state\", \"ρu[2]\", -4.78833188511707363855e-01, 5.60276621431795018857e-01, -5.51384446753656887186e-05, 2.12893889070269098918e-01 ],\n [ \"state\", \"ρθ\", -1.44374005663885895956e+01, 2.86753217014698513765e+00, 1.48814717217727286724e-03, 5.22593283079209602882e-01 ],\n],\n parr,\n)\n\nrusanov_overintegration = (\n[\n [ \"state\", \"ρ\", 9.71215931957259970275e-01, 1.00598091306596182370e+00, 1.00000018331969764418e+00, 6.61958163862229592017e-03 ],\n [ \"state\", \"ρu[1]\", -3.70465052437365993665e-01, 7.21959883918644518275e-01, 1.59085030658594694941e-01, 2.39392692254175726285e-01 ],\n [ \"state\", \"ρu[2]\", -4.40021840550348763976e-01, 4.17225753221437845042e-01, 7.34809696136107232513e-05, 1.50458151629451780673e-01 ],\n [ \"state\", \"ρθ\", -3.05207824503129643290e+00, 1.86668064854026383159e+00, -1.05021657861818027563e-02, 4.60981422911087401761e-01 ],\n],\n parr,\n)\n\nroeflux_overintegration = (\n[\n [ \"state\", \"ρ\", 9.69901051313856066294e-01, 1.00695111308847851106e+00, 1.00000002953972089159e+00, 6.94722590819101589593e-03 ],\n [ \"state\", \"ρu[1]\", -4.21771768146742886962e-01, 6.34790159642324547384e-01, 1.59384804664454482470e-01, 2.22844788492525480716e-01 ],\n [ \"state\", \"ρu[2]\", -4.37678540756429201863e-01, 4.53959896024816400573e-01, -1.22554654623556756642e-04, 1.73930782418124457722e-01 ],\n [ \"state\", \"ρθ\", -1.13636401173616619076e+00, 1.51145239791767727056e+00, 1.96870293018723595616e-03, 3.82573373360323043535e-01 ],\n],\n parr,\n)\n#! format: on\n\nrefVals = (;\n rusanov_periodic,\n roeflux_periodic,\n rusanov,\n roeflux,\n rusanov_overintegration,\n roeflux_overintegration,\n)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/run_bickley_jet.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"TwoDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\n########\n# Setup physical and numerical domains\n########\nΩˣ = IntervalDomain(-2π, 2π, periodic = true)\nΩʸ = IntervalDomain(-2π, 2π, periodic = false)\n\ngrid = DiscretizedDomain(\n Ωˣ × Ωʸ;\n elements = 8,\n polynomial_order = 3,\n overintegration_order = 1,\n)\n\n########\n# Define timestepping parameters\n########\nstart_time = 0\nend_time = 200.0\nΔt = 0.02\nmethod = SSPRK22Heuns\n\ntimestepper = TimeStepper(method = method, timestep = Δt)\n\ncallbacks = (Info(), StateCheck(10))\n\n########\n# Define physical parameters and parameterizations\n########\nparameters = (\n ϵ = 0.1, # perturbation size for initial condition\n l = 0.5, # Gaussian width\n k = 0.5, # Sinusoidal wavenumber\n ρₒ = 1.0, # reference density\n c = 2,\n g = 10,\n)\n\nphysics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0, κ = 0),\n coriolis = nothing,\n buoyancy = nothing,\n)\n\n########\n# Define boundary conditions\n########\nρu_bcs = (south = Impenetrable(FreeSlip()), north = Impenetrable(FreeSlip()))\nρθ_bcs = (south = Insulating(), north = Insulating())\nBC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n########\n# Define initial conditions\n########\n\n# The Bickley jet\nU₀(p, x, y, z) = cosh(y)^(-2)\n\n# Slightly off-center vortical perturbations\nΨ₀(p, x, y, z) =\n exp(-(y + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * x) * cos(p.k * y)\n\n# Vortical velocity fields (ũ, ṽ) = (-∂ʸ, +∂ˣ) ψ̃\nu₀(p, x, y, z) = Ψ₀(p, x, y, z) * (p.k * tan(p.k * y) + y / (p.l^2))\nv₀(p, x, y, z) = -Ψ₀(p, x, y, z) * p.k * tan(p.k * x)\nθ₀(p, x, y, z) = sin(p.k * y)\n\nρ₀(p, x, y, z) = p.ρₒ\nρu₀(p, x...) = ρ₀(p, x...) * (p.ϵ * u₀(p, x...) + U₀(p, x...))\nρv₀(p, x...) = ρ₀(p, x...) * p.ϵ * v₀(p, x...)\nρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\nρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...)]\ninitial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n########\n# Create the things\n########\n\nmodel = SpatialModel(\n balance_law = Fluid2D(),\n physics = physics,\n numerics = (flux = RoeNumericalFlux(),),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n)\n\nsimulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n)\n\n########\n# Run the model\n########\n\ntic = Base.time()\n\nevolve!(simulation, model)\n\ntoc = Base.time()\ntime = toc - tic\nprintln(time)\n" }, { "path": "test/Numerics/DGMethods/compressible_navier_stokes_equations/two_dimensional/test_bickley_jet.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../shared_source/boilerplate.jl\")\ninclude(\"TwoDimensionalCompressibleNavierStokesEquations.jl\")\n\nClimateMachine.init()\n\nsetups = [\n (;\n name = \"rusanov_periodic\",\n flux = RusanovNumericalFlux(),\n periodicity = true,\n Nover = 0,\n ),\n (;\n name = \"roeflux_periodic\",\n flux = RoeNumericalFlux(),\n periodicity = true,\n Nover = 0,\n ),\n (;\n name = \"rusanov\",\n flux = RusanovNumericalFlux(),\n periodicity = false,\n Nover = 0,\n ),\n (;\n name = \"roeflux\",\n flux = RoeNumericalFlux(),\n periodicity = false,\n Nover = 0,\n ),\n (;\n name = \"rusanov_overintegration\",\n flux = RusanovNumericalFlux(),\n periodicity = false,\n Nover = 1,\n ),\n (;\n name = \"roeflux_overintegration\",\n flux = RoeNumericalFlux(),\n periodicity = false,\n Nover = 1,\n ),\n]\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"refvals_bickley_jet.jl\")\n\n ########\n # Define timestepping parameters\n ########\n start_time = 0\n end_time = 200.0\n Δt = 0.02\n method = LSRK54CarpenterKennedy\n\n timestepper = TimeStepper(method = method, timestep = Δt)\n\n callbacks = (Info(), StateCheck(10))\n\n ########\n # Define physical parameters and parameterizations\n ########\n parameters = (\n ϵ = 0.1, # perturbation size for initial condition\n l = 0.5, # Gaussian width\n k = 0.5, # Sinusoidal wavenumber\n ρₒ = 1.0, # reference density\n c = 2,\n g = 10,\n )\n\n physics = FluidPhysics(;\n advection = NonLinearAdvectionTerm(),\n dissipation = ConstantViscosity{Float64}(μ = 0, ν = 0, κ = 0),\n coriolis = nothing,\n buoyancy = nothing,\n )\n\n ########\n # Define boundary conditions\n ########\n ρu_bc = Impenetrable(FreeSlip())\n ρθ_bc = Insulating()\n ρu_bcs = (south = ρu_bc, north = ρu_bc)\n ρθ_bcs = (south = ρθ_bc, north = ρθ_bc)\n BC = (ρθ = ρθ_bcs, ρu = ρu_bcs)\n\n ########\n # Define initial conditions\n ########\n\n # The Bickley jet\n U₀(p, x, y, z) = cosh(y)^(-2)\n\n # Slightly off-center vortical perturbations\n Ψ₀(p, x, y, z) =\n exp(-(y + p.l / 10)^2 / (2 * (p.l^2))) * cos(p.k * x) * cos(p.k * y)\n\n # Vortical velocity fields (ũ, ṽ) = (-∂ʸ, +∂ˣ) ψ̃\n u₀(p, x, y, z) = Ψ₀(p, x, y, z) * (p.k * tan(p.k * y) + y / (p.l^2))\n v₀(p, x, y, z) = -Ψ₀(p, x, y, z) * p.k * tan(p.k * x)\n θ₀(p, x, y, z) = sin(p.k * y)\n\n ρ₀(p, x, y, z) = p.ρₒ\n ρu₀(p, x...) = ρ₀(p, x...) * (p.ϵ * u₀(p, x...) + U₀(p, x...))\n ρv₀(p, x...) = ρ₀(p, x...) * p.ϵ * v₀(p, x...)\n ρθ₀(p, x...) = ρ₀(p, x...) * θ₀(p, x...)\n\n ρu⃗₀(p, x...) = @SVector [ρu₀(p, x...), ρv₀(p, x...)]\n initial_conditions = (ρ = ρ₀, ρu = ρu⃗₀, ρθ = ρθ₀)\n\n for setup in setups\n @testset \"$(setup.name)\" begin\n Ωˣ = Periodic(-2π, 2π)\n Ωʸ = IntervalDomain(-2π, 2π, periodic = setup.periodicity)\n\n grid = DiscretizedDomain(\n Ωˣ × Ωʸ,\n elements = 16,\n polynomial_order = 3,\n overintegration_order = setup.Nover,\n )\n\n model = SpatialModel(\n balance_law = Fluid2D(),\n physics = physics,\n numerics = (flux = setup.flux,),\n grid = grid,\n boundary_conditions = BC,\n parameters = parameters,\n )\n\n simulation = Simulation(\n model = model,\n initial_conditions = initial_conditions,\n timestepper = timestepper,\n callbacks = callbacks,\n time = (; start = start_time, finish = end_time),\n )\n\n ########\n # Run the model\n ########\n evolve!(\n simulation,\n model;\n refDat = getproperty(refVals, Symbol(setup.name)),\n )\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/conservation/sphere.jl", "content": "#=\nHere we solve the equation:\n```math\n q + dot(∇, uq) = 0\n p - dot(∇, up) = 0\n```\non a sphere to test the conservation of the numerics\n\nThe boundary conditions are `p = q` when `dot(n, u) > 0` and\n`q = p` when `dot(n, u) < 0` (i.e., `p` flows into `q` and vice-sersa).\n=#\n\nusing MPI\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.BalanceLaws:\n Prognostic, Auxiliary, AbstractStateType, BalanceLaw\n\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\nusing Random\n\nusing ClimateMachine.VariableTemplates\nimport ClimateMachine.BalanceLaws: vars_state\n\nimport ClimateMachine.DGMethods:\n flux_first_order!,\n flux_second_order!,\n source!,\n boundary_conditions,\n boundary_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!\n\nimport ClimateMachine.DGMethods: init_ode_state\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nimport ClimateMachine.DGMethods.NumericalFluxes:\n NumericalFluxFirstOrder,\n numerical_flux_first_order!,\n numerical_boundary_flux_first_order!\n\nstruct ConservationTestModel <: BalanceLaw end\n\nvars_state(::ConservationTestModel, ::Auxiliary, T) = @vars(vel::SVector{3, T})\nvars_state(::ConservationTestModel, ::Prognostic, T) = @vars(q::T, p::T)\n\nvars_state(::ConservationTestModel, ::AbstractStateType, T) = @vars()\n\nfunction nodal_init_state_auxiliary!(\n ::ConservationTestModel,\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n)\n x, y, z = g.coord\n r = x^2 + y^2 + z^2\n aux.vel = SVector(\n cos(10 * π * x) * sin(10 * π * y) + cos(20 * π * z),\n exp(sin(π * r)),\n sin(π * (x + y + z)),\n )\nend\n\nfunction init_state_prognostic!(\n ::ConservationTestModel,\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n)\n state.q = rand()\n state.p = rand()\nend\n\nfunction flux_first_order!(\n ::ConservationTestModel,\n flux::Grad,\n state::Vars,\n auxstate::Vars,\n t::Real,\n direction,\n)\n vel = auxstate.vel\n flux.q = state.q .* vel\n flux.p = -state.p .* vel\nend\n\nflux_second_order!(::ConservationTestModel, _...) = nothing\n\nsource!(::ConservationTestModel, _...) = nothing\n\nstruct ConservationTestModelNumFlux <: NumericalFluxFirstOrder end\n\nboundary_conditions(::ConservationTestModel) = (nothing,)\nboundary_state!(\n ::CentralNumericalFluxSecondOrder,\n bc,\n ::ConservationTestModel,\n _...,\n) = nothing\n\nfunction numerical_flux_first_order!(\n ::ConservationTestModelNumFlux,\n bl::BalanceLaw,\n fluxᵀn::Vars{S},\n n::SVector,\n state⁻::Vars{S},\n aux⁻::Vars{A},\n state⁺::Vars{S},\n aux⁺::Vars{A},\n t,\n direction,\n) where {S, A}\n un⁻ = dot(n, aux⁻.vel)\n un⁺ = dot(n, aux⁺.vel)\n un = (un⁺ + un⁻) / 2\n\n if un > 0\n fluxᵀn.q = un * state⁻.q\n fluxᵀn.p = -un * state⁺.p\n else\n fluxᵀn.q = un * state⁺.q\n fluxᵀn.p = -un * state⁻.p\n end\nend\n\nfunction numerical_boundary_flux_first_order!(\n ::ConservationTestModelNumFlux,\n bctype,\n bl::BalanceLaw,\n fluxᵀn::Vars{S},\n n::SVector,\n state⁻::Vars{S},\n aux⁻::Vars{A},\n state⁺::Vars{S},\n aux⁺::Vars{A},\n t,\n direction,\n state1⁻::Vars{S},\n aux1⁻::Vars{A},\n) where {S, A}\n un = dot(n, aux⁻.vel)\n\n if un > 0\n fluxᵀn.q = un * state⁻.q\n fluxᵀn.p = -un * state⁻.q\n else\n fluxᵀn.q = un * state⁻.p\n fluxᵀn.p = -un * state⁻.p\n end\nend\n\nfunction test_run(mpicomm, ArrayType, N, Nhorz, Rrange, timeend, FT, dt)\n\n topl = StackedCubedSphereTopology(mpicomm, Nhorz, Rrange)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = Topologies.equiangular_cubed_sphere_warp,\n )\n dg = DGModel(\n ConservationTestModel(),\n grid,\n ConservationTestModelNumFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0); init_on_cpu = true)\n\n lsrk = LSRK54CarpenterKennedy(dg, Q; dt = dt, t0 = 0)\n\n eng0 = norm(Q)\n sum0 = weightedsum(Q)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n sum(Q₀) = %.16e\"\"\" eng0 sum0\n\n max_mass_loss = FT(0)\n max_mass_gain = FT(0)\n cbmass = GenericCallbacks.EveryXSimulationSteps(1) do\n cbsum = weightedsum(Q)\n max_mass_loss = max(max_mass_loss, sum0 - cbsum)\n max_mass_gain = max(max_mass_gain, cbsum - sum0)\n end\n solve!(Q, lsrk; timeend = timeend, callbacks = (cbmass,))\n\n # Print some end of the simulation information\n engf = norm(Q)\n sumf = weightedsum(Q)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n max mass loss = %.16e\n max mass gain = %.16e\n initial mass = %.16e\n \"\"\" engf engf / eng0 engf - eng0 max_mass_loss max_mass_gain sum0\n max(max_mass_loss, max_mass_gain) / sum0\nend\n\nusing Test\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n\n Nhorz = 4\n\n tolerance = Dict(Float64 => 1e-15, Float32 => 1e-7)\n\n @testset \"$(@__FILE__)\" for FT in (Float64, Float32)\n dt = FT(1e-4)\n timeend = 100 * dt\n Rrange = range(FT(1), stop = FT(2), step = FT(1 // 4))\n\n @info (ArrayType, FT)\n delta_mass = test_run(\n mpicomm,\n ArrayType,\n polynomialorder,\n Nhorz,\n Rrange,\n timeend,\n FT,\n dt,\n )\n @test abs(delta_mass) < tolerance[FT]\n end\nend\n" }, { "path": "test/Numerics/DGMethods/courant.jl", "content": "using Test\nusing LinearAlgebra\nusing MPI\nusing StaticArrays\n\nusing ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Courant\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Orientations\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: kappa_d\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nconst p∞ = 10^5\nconst T∞ = 300.0\n\nfunction initialcondition!(problem, bl, state, aux, localgeo, t)\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n coord = localgeo.coord\n\n translation_speed::FT = 150\n translation_angle::FT = pi / 4\n α = translation_angle\n u∞ = SVector(\n FT(translation_speed * coord[1]),\n FT(translation_speed * coord[1]),\n FT(0),\n )\n _kappa_d::FT = kappa_d(param_set)\n\n u = u∞\n T = FT(T∞)\n # adiabatic/isentropic relation\n p = FT(p∞) * (T / FT(T∞))^(FT(1) / _kappa_d)\n ρ = air_density(param_set, T, p)\n\n state.ρ = ρ\n state.ρu = ρ * u\n e_kin = u' * u / 2\n state.energy.ρe = ρ * total_energy(param_set, e_kin, FT(0), T)\n\n nothing\nend\n\n\nlet\n # boiler plate MPI stuff\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n # Mesh generation parameters\n N = 4\n Nq = N + 1\n Neh = 10\n Nev = 4\n\n @testset \"$(@__FILE__) DGModel matrix\" begin\n for FT in (Float64, Float32)\n for dim in (2, 3)\n connectivity = dim == 3 ? :full : :face\n if dim == 2\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n elseif dim == 3\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n end\n μ = FT(2)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n connectivity = connectivity,\n )\n\n\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n problem = AtmosProblem(\n boundaryconditions = (),\n init_state_prognostic = initialcondition!,\n )\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(μ, WithDivergence()),\n moisture = DryModel(),\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n source = (Gravity(),),\n )\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Δt = FT(1 // 200)\n\n Q = init_ode_state(dg, FT(0))\n\n Δx = min_node_distance(grid, EveryDirection())\n Δx_v = min_node_distance(grid, VerticalDirection())\n Δx_h = min_node_distance(grid, HorizontalDirection())\n\n translation_speed = FT(norm([150.0, 150.0, 0.0]))\n\n diff_speed_h = FT(μ / air_density(param_set, FT(T∞), FT(p∞)))\n diff_speed_v = FT(μ / air_density(param_set, FT(T∞), FT(p∞)))\n c_h =\n Δt *\n (translation_speed + soundspeed_air(param_set, FT(T∞))) /\n Δx_h\n c_v = Δt * (soundspeed_air(param_set, FT(T∞))) / Δx_v\n d_h = Δt * diff_speed_h / Δx_h^2\n d_v = Δt * diff_speed_v / Δx_v^2\n simtime = FT(0)\n\n # tests for non diffusive courant number\n rtol = FT === Float64 ? 1e-4 : 1f-2\n @test courant(\n nondiffusive_courant,\n dg,\n model,\n Q,\n Δt,\n simtime,\n HorizontalDirection(),\n ) ≈ c_h rtol = rtol\n @test courant(\n nondiffusive_courant,\n dg,\n model,\n Q,\n Δt,\n simtime,\n VerticalDirection(),\n ) ≈ c_v rtol = rtol\n\n # tests for diffusive courant number\n @test courant(\n diffusive_courant,\n dg,\n model,\n Q,\n Δt,\n simtime,\n HorizontalDirection(),\n ) ≈ d_h\n @test courant(\n diffusive_courant,\n dg,\n model,\n Q,\n Δt,\n simtime,\n VerticalDirection(),\n ) ≈ d_v\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/custom_filter.jl", "content": "using Test\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates: @vars, Vars\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.DGMethods: AbstractCustomFilter, apply!\nimport ClimateMachine\nimport ClimateMachine.BalanceLaws:\n vars_state, init_state_auxiliary!, init_state_prognostic!\nusing MPI\nusing LinearAlgebra\n\nstruct CustomFilterTestModel <: BalanceLaw end\nstruct CustomTestFilter <: AbstractCustomFilter end\n\nvars_state(::CustomFilterTestModel, ::Auxiliary, FT) = @vars()\nvars_state(::CustomFilterTestModel, ::Prognostic, FT) where {N} = @vars(q::FT)\n\ninit_state_auxiliary!(::CustomFilterTestModel, _...) = nothing\n\nfunction init_state_prognostic!(\n ::CustomFilterTestModel,\n state::Vars,\n aux::Vars,\n localgeo,\n)\n coord = localgeo.coord\n state.q = hypot(coord[1], coord[2])\nend\n\n@testset \"Test custom filter\" begin\n let\n ClimateMachine.init()\n N = 4\n Ne = (2, 2, 2)\n\n function ClimateMachine.DGMethods.custom_filter!(\n ::CustomTestFilter,\n bl::CustomFilterTestModel,\n state::Vars,\n aux::Vars,\n )\n state.q = state.q^2\n end\n\n @testset for FT in (Float64, Float32)\n dim = 2\n brickrange =\n ntuple(j -> range(FT(-1); length = Ne[j] + 1, stop = 1), dim)\n topl = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_WORLD,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n )\n\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = N,\n )\n\n model = CustomFilterTestModel()\n dg = ClimateMachine.DGMethods.DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n Q = ClimateMachine.DGMethods.init_ode_state(dg)\n data = Array(Q.realdata)\n\n apply!(CustomTestFilter(), grid, model, Q, dg.state_auxiliary)\n\n @test all(Array(Q.realdata) .== data .^ 2)\n end\n end\nend\n" }, { "path": "test/Numerics/DGMethods/fv_reconstruction_test.jl", "content": "using Test\nusing StaticArrays\nusing ClimateMachine.Atmos:\n AtmosProblem,\n NoReferenceState,\n AtmosPhysics,\n AtmosModel,\n DryModel,\n ConstantDynamicViscosity,\n AtmosLESConfigType,\n HBFVReconstruction\nimport ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear, width\nusing ClimateMachine.Orientations\nimport StaticArrays: SUnitRange\nimport ClimateMachine.BalanceLaws:\n Primitive, Prognostic, vars_state, number_states\nusing ClimateMachine.VariableTemplates: Vars\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n\n\n# lin_func, quad_func, third_func, fourth_func\n#\n# ```\n# num_state_primitive::Int64\n# pointwise values::Array{FT, num_state_primitive by length(ξ)}\n# integration values::Array{FT, num_state_primitive by length(ξ)}\n# ```\n\nfunction lin_func(ξ)\n return 1, [(2 * ξ .+ 1)';], [(ξ .^ 2 .+ ξ)';]\nend\n\nfunction quad_func(ξ)\n return 2,\n [(3 * ξ .^ 2 .+ 1)'; (2 * ξ .+ 1)'],\n [(ξ .^ 3 .+ ξ)'; (ξ .^ 2 .+ ξ)']\nend\n\nfunction third_func(ξ)\n return 1, [(4 * ξ .^ 3 .+ 1)';], [(ξ .^ 4 .+ ξ)';]\nend\n\nfunction fourth_func(ξ)\n return 2,\n [(5 * ξ .^ 4 .+ 1)'; (3 * ξ .^ 2 .+ 1)'],\n [(ξ .^ 5 .+ ξ)'; (ξ .^ 3 .+ ξ)']\nend\n\n\n@testset \"Hydrostatic balanced linear reconstruction test\" begin\n\n function initialcondition!(problem, bl, state, aux, coords, t, args...) end\n\n fv_recon! = FVLinear()\n stencil_width = width(fv_recon!)\n stencil_center = stencil_width + 1\n stencil_diameter = 2stencil_width + 1\n @test stencil_width == 1\n func = lin_func\n\n for FT in (Float64,)\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = AtmosProblem(;\n physics = physics,\n init_state_prognostic = initialcondition!,\n ),\n orientation = FlatOrientation(),\n )\n\n vars_prim = Vars{vars_state(model, Primitive(), FT)}\n\n hb_recon! = HBFVReconstruction(model, fv_recon!)\n _grav = FT(grav(param_set))\n\n @test width(hb_recon!) == 1\n\n num_state_prognostic = number_states(model, Prognostic())\n num_state_primitive = number_states(model, Primitive())\n local_state_face_primitive = ntuple(Val(2)) do _\n MArray{Tuple{num_state_primitive}, FT}(undef)\n end\n\n\n\n # interior point reconstruction test\n grid = FT[0; 1; 3; 6]\n grid_c = (grid[2:end] + grid[1:(end - 1)]) / 2\n local_cell_weights =\n MArray{Tuple{stencil_diameter}, FT}(grid[2:end] - grid[1:(end - 1)])\n\n # linear profile for all variables expect pressure\n local_state_primitive = SVector(ntuple(Val(stencil_diameter)) do _\n MArray{Tuple{num_state_primitive}, FT}(undef)\n end...)\n\n # values at the cell centers 1 2 3\n _, uc, _ = func(grid_c)\n # values at the cell faces 0.5 1.5* 2.5* 3.5\n _, uf, _ = func(grid)\n for i_d in 1:stencil_diameter\n for i_p in 1:num_state_prognostic\n local_state_primitive[i_d][i_p] = uc[i_d]\n end\n end\n\n # pressure profile is updated to satisfy the discrete hydrostatic balance \n p_surf = FT(100) # at the bottom wall\n p_ref = p_surf\n for i_d in 1:stencil_diameter\n p_ref -=\n vars_prim(local_state_primitive[i_d]).ρ *\n _grav *\n local_cell_weights[i_d] / 2\n vars_prim(local_state_primitive[i_d]).p = p_ref\n p_ref -=\n vars_prim(local_state_primitive[i_d]).ρ *\n _grav *\n local_cell_weights[i_d] / 2\n end\n\n local_state_primitive_hb = copy(local_state_primitive)\n\n # interior point reconstruction\n hb_recon!(\n local_state_face_primitive[1],\n local_state_face_primitive[2],\n local_state_primitive,\n local_cell_weights,\n )\n # bottom face\n @test vars_prim(local_state_face_primitive[1]).ρ ≈ uf[2]\n @test vars_prim(local_state_face_primitive[1]).p ≈\n p_surf -\n vars_prim(local_state_primitive[1]).ρ *\n _grav *\n local_cell_weights[1]\n # top face\n @test vars_prim(local_state_face_primitive[2]).ρ ≈ uf[3]\n @test vars_prim(local_state_face_primitive[2]).p ≈\n p_surf -\n vars_prim(local_state_primitive[1]).ρ *\n _grav *\n local_cell_weights[1] -\n vars_prim(local_state_primitive[2]).ρ *\n _grav *\n local_cell_weights[2]\n\n # make sure the \n for i_d in 1:stencil_diameter\n @test all(\n local_state_primitive[i_d] ≈ local_state_primitive_hb[i_d],\n )\n end\n\n @info \"Start boundary test\"\n\n\n # boundary point reconstruction \n rng = SUnitRange(stencil_center, stencil_center)\n\n hb_recon!(\n local_state_face_primitive[1],\n local_state_face_primitive[2],\n local_state_primitive[rng],\n local_cell_weights[rng],\n )\n\n # bottom face\n @test vars_prim(local_state_face_primitive[1]).ρ ≈ uc[stencil_center]\n @test vars_prim(local_state_face_primitive[1]).p ≈\n vars_prim(local_state_primitive[stencil_center]).p +\n uc[stencil_center] * _grav * local_cell_weights[stencil_center] /\n 2\n\n # top face\n @test vars_prim(local_state_face_primitive[2]).ρ ≈ uc[stencil_center]\n @test vars_prim(local_state_face_primitive[2]).p ≈\n vars_prim(local_state_primitive[stencil_center]).p -\n uc[stencil_center] * _grav * local_cell_weights[stencil_center] /\n 2\n\n end\n\n\nend\n" }, { "path": "test/Numerics/DGMethods/grad_test.jl", "content": "using MPI\nusing StaticArrays\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing Printf\n\nusing ClimateMachine.BalanceLaws\n\nimport ClimateMachine.BalanceLaws: vars_state, nodal_init_state_auxiliary!\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nstruct GradTestModel{dim, dir} <: BalanceLaw end\n\nvars_state(m::GradTestModel, ::Auxiliary, T) = @vars begin\n a::T\n ∇a::SVector{3, T}\nend\nvars_state(::GradTestModel, ::Prognostic, T) = @vars()\n\nfunction nodal_init_state_auxiliary!(\n ::GradTestModel{dim, dir},\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n) where {dim, dir}\n x, y, z = g.coord\n if dim == 2\n aux.a = x^2 + y^3 - x * y\n if dir isa EveryDirection\n aux.∇a = SVector(2 * x - y, 3 * y^2 - x, 0)\n elseif dir isa HorizontalDirection\n aux.∇a = SVector(2 * x - y, 0, 0)\n elseif dir isa VerticalDirection\n aux.∇a = SVector(0, 3 * y^2 - x, 0)\n end\n else\n aux.a = x^2 + y^3 + z^2 * y^2 - x * y * z\n if dir isa EveryDirection\n aux.∇a = SVector(\n 2 * x - y * z,\n 3 * y^2 + 2 * z^2 * y - x * z,\n 2 * z * y^2 - x * y,\n )\n elseif dir isa HorizontalDirection\n aux.∇a = SVector(2 * x - y * z, 3 * y^2 + 2 * z^2 * y - x * z, 0)\n elseif dir isa VerticalDirection\n aux.∇a = SVector(0, 0, 2 * z * y^2 - x * y)\n end\n end\nend\n\nusing Test\nfunction test_run(mpicomm, dim, direction, Ne, N, FT, ArrayType)\n connectivity = dim == 3 ? :full : :face\n brickrange = ntuple(j -> range(FT(0); length = Ne[j] + 1, stop = 3), dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = ntuple(j -> false, dim),\n connectivity = connectivity,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n model = GradTestModel{dim, direction}()\n dg = DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n exact_aux = copy(dg.state_auxiliary)\n\n auxiliary_field_gradient!(\n model,\n dg.state_auxiliary,\n (:∇a,),\n dg.state_auxiliary,\n (:a,),\n grid,\n direction,\n )\n\n # Wrapping in Array ensure both GPU and CPU code use same approx\n approx = Array(dg.state_auxiliary.∇a) ≈ Array(exact_aux.∇a)\n err = euclidean_distance(exact_aux, dg.state_auxiliary)\n return approx, err\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n numelem = (5, 5, 5)\n\n expected_result = Dict()\n expected_result[Float64, 2, 1] = 6.2135207410935696e+00\n expected_result[Float64, 2, 2] = 2.3700094936518794e+00\n expected_result[Float64, 2, 3] = 8.7105082013050261e-01\n expected_result[Float64, 2, 4] = 3.1401219279927611e-01\n\n expected_result[Float64, 3, 1] = 7.9363467666175236e+00\n expected_result[Float64, 3, 2] = 2.8059223082616098e+00\n expected_result[Float64, 3, 3] = 9.9204334582727094e-01\n expected_result[Float64, 3, 4] = 3.5074028853276679e-01\n\n expected_result[Float32, 2, 1] = 6.2135176658630371e+00\n expected_result[Float32, 2, 2] = 2.3700129985809326e+00\n expected_result[Float32, 3, 1] = 7.9363408088684082e+00\n expected_result[Float32, 3, 2] = 2.8059237003326416e+00\n\n\n @testset for FT in (Float64, Float32)\n lvls =\n integration_testing || ClimateMachine.Settings.integration_testing ?\n (FT === Float32 ? 2 : 4) : 1\n\n @testset for polynomialorder in ((4, 4), (4, 0))\n @testset for dim in 2:3\n @testset for direction in (\n EveryDirection(),\n HorizontalDirection(),\n VerticalDirection(),\n )\n err = zeros(FT, lvls)\n for l in 1:lvls\n approx, err[l] = test_run(\n mpicomm,\n dim,\n direction,\n ntuple(j -> 2^(l - 1) * numelem[j], dim),\n polynomialorder,\n FT,\n ArrayType,\n )\n if polynomialorder[end] != 0 ||\n direction isa HorizontalDirection\n @test approx\n else\n @test err[l] ≈ expected_result[FT, dim, l]\n end\n end\n if polynomialorder[end] != 0 ||\n direction isa HorizontalDirection\n @info begin\n msg = \"Polynomial order = $polynomialorder, direction = $direction\\n\"\n for l in 1:(lvls - 1)\n rate = log2(err[l]) - log2(err[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/grad_test_sphere.jl", "content": "using MPI\nusing StaticArrays\nusing LinearAlgebra\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing Printf\n\nusing ClimateMachine.BalanceLaws\n\nimport ClimateMachine.BalanceLaws: vars_state, nodal_init_state_auxiliary!\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nstruct GradSphereTestModel{dir} <: BalanceLaw end\n\nvars_state(m::GradSphereTestModel, ::Auxiliary, T) = @vars begin\n a::T\n ∇a::SVector{3, T}\nend\nvars_state(::GradSphereTestModel, ::Prognostic, T) = @vars()\n\nfunction nodal_init_state_auxiliary!(\n ::GradSphereTestModel{dir},\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n) where {dir}\n x, y, z = g.coord\n r = hypot(x, y, z)\n aux.a = r^3\n if !(dir isa HorizontalDirection)\n aux.∇a = 3 * r^2 * g.coord / r\n else\n aux.∇a = SVector(0, 0, 0)\n end\nend\n\nusing Test\nfunction test_run(mpicomm, Ne_horz, Ne_vert, N, FT, ArrayType, direction)\n\n Rrange = range(FT(1 // 2); length = Ne_vert + 1, stop = FT(1))\n topl = StackedCubedSphereTopology(mpicomm, Ne_horz, Rrange)\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = Topologies.equiangular_cubed_sphere_warp,\n )\n\n model = GradSphereTestModel{direction}()\n dg = DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n exact_aux = copy(dg.state_auxiliary)\n\n auxiliary_field_gradient!(\n model,\n dg.state_auxiliary,\n (:∇a,),\n dg.state_auxiliary,\n (:a,),\n grid,\n direction,\n )\n\n err = euclidean_distance(exact_aux, dg.state_auxiliary)\n return err\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n base_Nhorz = 4\n base_Nvert = 2\n\n expected_result = Dict()\n expected_result[(4, 4), 1] = 4.3759489495202896e-04\n expected_result[(4, 4), 2] = 2.9065372851175251e-05\n expected_result[(4, 4), 3] = 1.8457379514995729e-06\n expected_result[(4, 4), 4] = 1.1582093840084037e-07\n\n expected_result[(4, 0), 1] = 1.1070045305138052e+00\n expected_result[(4, 0), 2] = 4.2750547196265593e-01\n expected_result[(4, 0), 3] = 1.6041478911787385e-01\n expected_result[(4, 0), 4] = 5.8570590697850776e-02\n\n lvls =\n integration_testing || ClimateMachine.Settings.integration_testing ? 4 :\n 1\n\n @testset for FT in (Float64,)\n @testset for polynomialorder in ((4, 4), (4, 0))\n @testset for direction in (\n EveryDirection(),\n HorizontalDirection(),\n VerticalDirection(),\n )\n err = zeros(FT, lvls)\n @testset for l in 1:lvls\n Ne_horz = 2^(l - 1) * base_Nhorz\n Ne_vert = 2^(l - 1) * base_Nvert\n\n err[l] = test_run(\n mpicomm,\n Ne_horz,\n Ne_vert,\n polynomialorder,\n FT,\n ArrayType,\n direction,\n )\n if !(direction isa HorizontalDirection)\n @test err[l] ≈ expected_result[polynomialorder, l]\n else\n @test abs(err[l]) < FT(1.3e-13)\n end\n end\n if !(direction isa HorizontalDirection)\n @info begin\n msg = \"Polynomial order = $polynomialorder, direction = $direction\\n\"\n for l in 1:(lvls - 1)\n rate = log2(err[l]) - log2(err[l + 1])\n msg *= @sprintf(\n \"\\n rate for level %d = %e\\n\",\n l,\n rate\n )\n end\n msg\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/horizontal_integral_test.jl", "content": "using MPI\nusing StaticArrays\nusing Logging\nusing Printf\nusing LinearAlgebra\nusing Test\nimport KernelAbstractions: CPU\n\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.VariableTemplates\nusing Thermodynamics\n\nfunction run_test1(mpicomm, dim, Ne, N, FT, ArrayType)\n warpfun =\n (ξ1, ξ2, ξ3) -> begin\n x1 = ξ1 + (ξ1 - 1 / 2) * cos(2 * π * ξ2 * ξ3) / 4\n x2 = ξ2 + (ξ2 - 1 / 2) * cos(2 * π * ξ2 * ξ3) / 4\n x3 = ξ3 + ξ1 / 4 + sin(2 * π * ξ1) / 16\n return (x1, x2, x3)\n end\n\n brickrange = (\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = 2, stop = 1),\n )\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false, false),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n nrealelem = length(topl.realelems)\n Nq = N + 1\n Nqk = dimensionality(grid) == 2 ? 1 : Nq\n vgeo = grid.vgeo\n localvgeo = array_device(vgeo) isa CPU ? vgeo : Array(vgeo)\n\n S = zeros(Nqk)\n S1 = zeros(Nqk)\n\n for e in 1:nrealelem\n for k in 1:Nqk\n for j in 1:Nq\n for i in 1:Nq\n ijk = i + Nq * ((j - 1) + Nq * (k - 1))\n S[k] +=\n localvgeo[ijk, grid.x1id, e] *\n localvgeo[ijk, grid.MHid, e]\n S1[k] += localvgeo[ijk, grid.MHid, e]\n end\n end\n end\n end\n\n Stot = zeros(Nqk)\n S1tot = zeros(Nqk)\n Err = 0.0\n\n for k in 1:Nqk\n Stot[k] = MPI.Reduce(S[k], +, 0, mpicomm)\n S1tot[k] = MPI.Reduce(S1[k], +, 0, mpicomm)\n Err += (0.5 - Stot[k] / S1tot[k])^2\n end\n Err = sqrt(Err / Nqk)\n @test Err < 2e-15\nend\n\nfunction run_test2(mpicomm, dim, Ne, N, FT, ArrayType)\n warpfun =\n (ξ1, ξ2, ξ3) -> begin\n x1 = sin(2 * π * ξ3) / 16 + ξ1\n x2 = ξ2 + (ξ2 - 1 / 2) * cos(2 * π * ξ2 * ξ3) / 4\n x3 = ξ3 + sin(2π * (ξ1)) / 20\n return (x1, x2, x3)\n end\n\n brickrange = (\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = 2, stop = 1),\n )\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false, false),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n nrealelem = length(topl.realelems)\n Nq = N + 1\n Nqk = dimensionality(grid) == 2 ? 1 : Nq\n vgeo = grid.vgeo\n localvgeo = array_device(vgeo) isa CPU ? vgeo : Array(vgeo)\n\n S = zeros(Nqk)\n S1 = zeros(Nqk)\n K = zeros(Nqk)\n\n for e in 1:nrealelem\n for k in 1:Nqk\n for j in 1:Nq\n for i in 1:Nq\n ijk = i + Nq * ((j - 1) + Nq * (k - 1))\n S[k] +=\n localvgeo[ijk, grid.x1id, e] *\n localvgeo[ijk, grid.MHid, e]\n S1[k] += localvgeo[ijk, grid.MHid, e]\n K[k] = localvgeo[ijk, grid.x3id, e]\n end\n end\n end\n end\n\n Stot = zeros(Nqk)\n S1tot = zeros(Nqk)\n Err = 0.0\n\n for k in 1:Nqk\n Stot[k] = MPI.Reduce(S[k], +, 0, mpicomm)\n S1tot[k] = MPI.Reduce(S1[k], +, 0, mpicomm)\n Err += (0.5 + sin(2 * π * K[k]) / 16 - Stot[k] / S1tot[k])^2\n end\n Err = sqrt(Err / Nqk)\n @test 2e-15 > Err\nend\n\nfunction run_test3(mpicomm, dim, Ne, N, FT, ArrayType)\n base_Nhorz = 4\n base_Nvert = 2\n Rinner = 1 // 2\n Router = 1\n expected_result = [\n -4.5894269717905445e-8 -1.1473566985387151e-8\n -2.0621904184281448e-10 -5.155431637149377e-11\n -8.72191208145523e-13 -2.1715962361668062e-13\n ]\n for l in 1:3\n Nhorz = 2^(l - 1) * base_Nhorz\n Nvert = 2^(l - 1) * base_Nvert\n Rrange = grid1d(FT(Rinner), FT(Router); nelem = Nvert)\n topl = StackedCubedSphereTopology(mpicomm, Nhorz, Rrange)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = Topologies.equiangular_cubed_sphere_warp,\n )\n\n nrealelem = length(topl.realelems)\n Nq = N + 1\n Nqk = dimensionality(grid) == 2 ? 1 : Nq\n vgeo = grid.vgeo\n localvgeo = array_device(vgeo) isa CPU ? vgeo : Array(vgeo)\n\n topology = grid.topology\n nvertelem = topology.stacksize\n nhorzelem = div(nrealelem, nvertelem)\n\n Surfout = 0\n Surfin = 0\n\n for ev in 1:nvertelem\n for eh in 1:nhorzelem\n e = ev + (eh - 1) * nvertelem\n for i in 1:Nq\n for j in 1:Nq\n for k in 1:Nqk\n ijk = i + Nq * ((j - 1) + Nqk * (k - 1))\n if (k == Nqk && ev == nvertelem)\n Surfout += localvgeo[ijk, grid.MHid, e]\n end\n if (k == 1 && ev == 1)\n Surfin += localvgeo[ijk, grid.MHid, e]\n end\n end\n end\n end\n end\n end\n Surfouttot = MPI.Reduce(Surfout, +, 0, MPI.COMM_WORLD)\n Surfintot = MPI.Reduce(Surfin, +, 0, MPI.COMM_WORLD)\n @test (4 * π * Router^2 - Surfouttot) ≈ expected_result[l, 1] rtol =\n 1e-3 atol = eps(FT) * 4 * π * Router^2\n @test (4 * π * Rinner^2 - Surfintot) ≈ expected_result[l, 2] rtol = 1e-3 atol =\n eps(FT) * 4 * π * Rinner^2\n end\nend\n\n# Test for 2D integral\nfunction run_test4(mpicomm, dim, Ne, N, FT, ArrayType)\n brickrange = ntuple(j -> range(FT(0); length = Ne + 1, stop = 1), 2)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = ntuple(j -> true, 2),\n connectivity = :face,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n nrealelem = length(topl.realelems)\n Nq = N + 1\n vgeo = grid.vgeo\n localvgeo = array_device(vgeo) isa CPU ? vgeo : Array(vgeo)\n\n S = zeros(Nq)\n\n for e in 1:nrealelem\n for i in 1:Nq\n for j in 1:Nq\n ij = i + Nq * (j - 1)\n S[j] +=\n localvgeo[ij, grid.x1id, e] * localvgeo[ij, grid.MHid, e]\n end\n end\n end\n\n Err = 0\n for j in 1:Nq\n Err += (S[j] - 0.5)^2\n end\n\n Err = sqrt(Err / Nq)\n @test Err <= 1e-15\nend\n\nfunction run_test5(mpicomm, dim, Ne, N, FT, ArrayType)\n warpfun = (ξ1, ξ2, ξ3) -> begin\n x1 = cos(π * ξ2) / 16 + abs(ξ1)\n x2 = ξ2\n x3 = ξ3\n return (x1, x2, x3)\n end\n\n brickrange = ntuple(j -> range(FT(0); length = Ne + 1, stop = 1), 2)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = ntuple(j -> true, 2),\n connectivity = :face,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n nrealelem = length(topl.realelems)\n Nq = N + 1\n vgeo = grid.vgeo\n localvgeo = array_device(vgeo) isa CPU ? vgeo : Array(vgeo)\n\n S = zeros(Nq)\n J = zeros(Nq)\n\n for e in 1:nrealelem\n for i in 1:Nq\n for j in 1:Nq\n ij = i + Nq * (j - 1)\n S[j] +=\n localvgeo[ij, grid.x1id, e] * localvgeo[ij, grid.MHid, e]\n J[j] = localvgeo[ij, grid.x2id, e]\n end\n end\n end\n\n Err = 0\n for j in 1:Nq\n Err += (S[j] - 0.5 - cos(π * J[j]) / 16)^2\n end\n\n Err = sqrt(Err / Nq)\n @test Err < 1e-15\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n FT = Float64\n dim = 3\n Ne = 1\n polynomialorder = 4\n\n @info (ArrayType, FT, dim)\n\n @testset \"horizontal_integral\" begin\n run_test1(mpicomm, dim, Ne, polynomialorder, FT, ArrayType)\n run_test2(mpicomm, dim, Ne, polynomialorder, FT, ArrayType)\n run_test3(mpicomm, dim, Ne, polynomialorder, FT, ArrayType)\n run_test4(mpicomm, dim, Ne, polynomialorder, FT, ArrayType)\n run_test5(mpicomm, dim, Ne, polynomialorder, FT, ArrayType)\n end\nend\n" }, { "path": "test/Numerics/DGMethods/integral_test.jl", "content": "using MPI\nusing StaticArrays\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing Printf\nusing LinearAlgebra\nusing Logging\n\nusing ClimateMachine.BalanceLaws\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n boundary_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n update_auxiliary_state!,\n indefinite_stack_integral!,\n reverse_indefinite_stack_integral!,\n integral_load_auxiliary_state!,\n integral_set_auxiliary_state!,\n reverse_integral_load_auxiliary_state!,\n reverse_integral_set_auxiliary_state!\n\nimport ClimateMachine.DGMethods: init_ode_state\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\n\nstruct IntegralTestModel{dim} <: BalanceLaw end\n\nvars_state(::IntegralTestModel, ::DownwardIntegrals, T) = @vars(a::T, b::T)\nvars_state(::IntegralTestModel, ::UpwardIntegrals, T) = @vars(a::T, b::T)\nvars_state(m::IntegralTestModel, ::Auxiliary, T) = @vars(\n int::vars_state(m, UpwardIntegrals(), T),\n rev_int::vars_state(m, DownwardIntegrals(), T),\n coord::SVector{3, T},\n a::T,\n b::T,\n rev_a::T,\n rev_b::T\n)\n\nvars_state(::IntegralTestModel, ::AbstractStateType, T) = @vars()\n\nflux_first_order!(::IntegralTestModel, _...) = nothing\nflux_second_order!(::IntegralTestModel, _...) = nothing\nsource!(::IntegralTestModel, _...) = nothing\nboundary_state!(_, ::IntegralTestModel, _...) = nothing\ninit_state_prognostic!(::IntegralTestModel, _...) = nothing\nwavespeed(::IntegralTestModel, _...) = 1\n\nfunction nodal_init_state_auxiliary!(\n ::IntegralTestModel{dim},\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n) where {dim}\n x, y, z = aux.coord = g.coord\n if dim == 2\n aux.a = x * y\n aux.b = 2 * x * y + sin(x) * y^2 / 2 - (z - 1)^2 * y^3 / 3\n y_top = 3\n a_top = x * y_top\n b_top = 2 * x * y_top + sin(x) * y_top^2 / 2 - (z - 1)^2 * y_top^3 / 3\n aux.rev_a = a_top - aux.a\n aux.rev_b = b_top - aux.b\n else\n aux.a = x * z + y * z\n aux.b = 2 * x * z + sin(x) * y * z - (1 + (z - 1)^3) * y^2 / 3\n zz_top = 3\n a_top = x * zz_top + y * zz_top\n b_top =\n 2 * x * zz_top + sin(x) * y * zz_top -\n (1 + (zz_top - 1)^3) * y^2 / 3\n aux.rev_a = a_top - aux.a\n aux.rev_b = b_top - aux.b\n end\nend\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n m::IntegralTestModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n indefinite_stack_integral!(dg, m, Q, dg.state_auxiliary, t, elems)\n reverse_indefinite_stack_integral!(dg, m, Q, dg.state_auxiliary, t, elems)\n\n return true\nend\n\n@inline function integral_load_auxiliary_state!(\n m::IntegralTestModel{dim},\n integrand::Vars,\n state::Vars,\n aux::Vars,\n) where {dim}\n x, y, z = aux.coord\n integrand.a = x + (dim == 3 ? y : 0)\n integrand.b = 2 * x + sin(x) * y - (z - 1)^2 * y^2\nend\n\n@inline function integral_set_auxiliary_state!(\n m::IntegralTestModel,\n aux::Vars,\n integral::Vars,\n)\n aux.int.a = integral.a\n aux.int.b = integral.b\nend\n\n@inline function reverse_integral_load_auxiliary_state!(\n m::IntegralTestModel,\n integral::Vars,\n state::Vars,\n aux::Vars,\n)\n integral.a = aux.int.a\n integral.b = aux.int.b\nend\n\n@inline function reverse_integral_set_auxiliary_state!(\n m::IntegralTestModel,\n aux::Vars,\n integral::Vars,\n)\n aux.rev_int.a = integral.a\n aux.rev_int.b = integral.b\nend\n\nusing Test\nfunction test_run(mpicomm, dim, Ne, N, FT, ArrayType)\n connectivity = dim == 3 ? :full : :face\n brickrange = ntuple(j -> range(FT(0); length = Ne[j] + 1, stop = 3), dim)\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n connectivity = connectivity,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n dg = DGModel(\n IntegralTestModel{dim}(),\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n dQdt = similar(Q)\n\n dg(dQdt, Q, nothing, 0.0)\n\n # Wrapping in Array ensure both GPU and CPU code use same approx\n if N[end] > 0\n # Forward integral a\n @test Array(dg.state_auxiliary.data[:, 1, :]) ≈\n Array(dg.state_auxiliary.data[:, 8, :])\n # Forward integral b\n @test Array(dg.state_auxiliary.data[:, 2, :]) ≈\n Array(dg.state_auxiliary.data[:, 9, :])\n # Reverse integral a\n @test Array(dg.state_auxiliary.data[:, 3, :]) ≈\n Array(dg.state_auxiliary.data[:, 10, :])\n # Reverse integral b\n @test Array(dg.state_auxiliary.data[:, 4, :]) ≈\n Array(dg.state_auxiliary.data[:, 11, :])\n else\n # For N = 0 we only compare the first integral which is the integral of\n # a vertical constant function; N = 0 can also integrate linears exactly\n # since we use the midpoint rule to compute the face value, but the\n # averaging procedure we use below does not work in this case\n Nq = polynomialorders(grid) .+ 1\n\n # Reshape the data array to be (dofs, vertical elem, horizontal elem)\n A_faces = reshape(\n Array(dg.state_auxiliary.data[:, 1, :]),\n prod(Nq),\n Ne[end], # Vertical element is fastest on stacked meshes\n prod(Ne[1:(end - 1)]), # Horiztonal elements\n )\n A_center_exact = reshape(\n Array(dg.state_auxiliary.data[:, 8, :]),\n prod(Nq),\n Ne[end], # Vertical element is fastest on stacked meshes\n prod(Ne[1:(end - 1)]), # Horiztonal elements\n )\n # With N = 0, the integral will return the values in the faces. Namely,\n # verical element index `eV` will be the value of the integral on the\n # face ABOVE element `eV`. Namely,\n # A_faces[n, eV, eH]\n # will be degree of freedom `n`, in horizontal element stack `eH`, and\n # face `eV + 1/2`.\n #\n # The exact values stored in `A_center_exact` are actually at the cell\n # centers because these are computed using the `init_state_auxiliary!`\n # which evaluates using the cell centers.\n #\n # This mismatch means we need to convert from faces to cell centers for\n # comparison, and we do this using averaging to go from faces to cell\n # centers.\n\n # Storage for the averaging\n A_center = similar(A_faces)\n\n # Bottom cell value is average of 0 and top face of cell\n A_center[:, 1, :] .= A_faces[:, 1, :] / 2\n\n # Remaining cells are average of the two faces\n A_center[:, 2:end, :, :] .=\n (A_faces[:, 1:(end - 1), :] + A_faces[:, 2:end, :]) / 2\n\n # Compare the exact and computed\n @test A_center ≈ A_center_exact\n\n # We do the same things for the reverse integral, the only difference is\n # now the values\n # RA_faces[n, eV, eH]\n # will be degree of freedom `n`, in horizontal element stack `eH`, and\n # face `eV - 1/2` (e.g., the face below element `eV`\n # Reshape the data array to be (dofs, vertical elm, horizontal elm)\n RA_faces = reshape(\n Array(dg.state_auxiliary.data[:, 3, :]),\n prod(Nq),\n Ne[end], # Vertical element is fastest on stacked meshes\n prod(Ne[1:(end - 1)]), # Horiztonal elements\n )\n RA_center_exact = reshape(\n Array(dg.state_auxiliary.data[:, 10, :]),\n prod(Nq),\n Ne[end], # Vertical element is fastest on stacked meshes\n prod(Ne[1:(end - 1)]), # Horiztonal elements\n )\n\n # Storage for the averaging\n RA_center = similar(RA_faces)\n\n # Top cell value is average of 0 and top face of cell\n RA_center[:, end, :] .= RA_faces[:, end, :] / 2\n\n # Remaining cells are average of the two faces\n RA_center[:, 1:(end - 1), :, :] .=\n (RA_faces[:, 1:(end - 1), :] + RA_faces[:, 2:end, :]) / 2\n\n # Compare the exact and computed\n @test RA_center ≈ RA_center_exact\n end\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n numelem = (5, 6, 7)\n lvls = 1\n\n for polynomialorder in ((4, 4), (4, 3), (4, 0))\n for FT in (Float64,)\n for dim in 2:3\n err = zeros(FT, lvls)\n for l in 1:lvls\n @info (ArrayType, FT, dim, polynomialorder)\n test_run(\n mpicomm,\n dim,\n ntuple(j -> 2^(l - 1) * numelem[j], dim),\n polynomialorder,\n FT,\n ArrayType,\n )\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/integral_test_sphere.jl", "content": "using MPI\nusing StaticArrays\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing Printf\nusing LinearAlgebra\nusing Logging\n\nusing ClimateMachine.BalanceLaws:\n BalanceLaw,\n Prognostic,\n Auxiliary,\n GradientFlux,\n UpwardIntegrals,\n DownwardIntegrals\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n integral_load_auxiliary_state!,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n update_auxiliary_state!,\n indefinite_stack_integral!,\n reverse_indefinite_stack_integral!,\n boundary_conditions,\n boundary_state!,\n compute_gradient_argument!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n integral_set_auxiliary_state!,\n reverse_integral_load_auxiliary_state!,\n reverse_integral_set_auxiliary_state!\n\nimport ClimateMachine.DGMethods: init_ode_state\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nstruct IntegralTestSphereModel{T} <: BalanceLaw\n Rinner::T\n Router::T\nend\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n m::IntegralTestSphereModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n indefinite_stack_integral!(dg, m, Q, dg.state_auxiliary, t, elems)\n reverse_indefinite_stack_integral!(dg, m, Q, dg.state_auxiliary, t, elems)\n\n return true\nend\n\nvars_state(::IntegralTestSphereModel, ::UpwardIntegrals, T) = @vars(v::T, r::T)\nvars_state(::IntegralTestSphereModel, ::DownwardIntegrals, T) =\n @vars(v::T, r::T)\nvars_state(m::IntegralTestSphereModel, ::Auxiliary, T) = @vars(\n int::vars_state(m, UpwardIntegrals(), T),\n rev_int::vars_state(m, DownwardIntegrals(), T),\n r::T,\n v::T\n)\n\nvars_state(::IntegralTestSphereModel, ::Prognostic, T) = @vars()\nvars_state(::IntegralTestSphereModel, ::GradientFlux, T) = @vars()\n\nflux_first_order!(::IntegralTestSphereModel, _...) = nothing\nflux_second_order!(::IntegralTestSphereModel, _...) = nothing\nsource!(::IntegralTestSphereModel, _...) = nothing\nboundary_conditions(::IntegralTestSphereModel) = (nothing,)\nboundary_state!(_, ::Nothing, ::IntegralTestSphereModel, _...) = nothing\ninit_state_prognostic!(::IntegralTestSphereModel, _...) = nothing\nwavespeed(::IntegralTestSphereModel, _...) = 1\n\nfunction nodal_init_state_auxiliary!(\n m::IntegralTestSphereModel,\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n)\n x, y, z = g.coord\n aux.r = hypot(x, y, z)\n θ = atan(y, x)\n ϕ = asin(z / aux.r)\n # Exact integral\n aux.v = 1 + cos(ϕ)^2 * sin(θ)^2 + sin(ϕ)^2\n aux.int.v = exp(-aux.v * aux.r^2) - exp(-aux.v * m.Rinner^2)\n aux.int.r = aux.r - m.Rinner\n aux.rev_int.v = exp(-aux.v * m.Router^2) - exp(-aux.v * aux.r^2)\n aux.rev_int.r = m.Router - aux.r\nend\n\n@inline function integral_load_auxiliary_state!(\n m::IntegralTestSphereModel,\n integrand::Vars,\n state::Vars,\n aux::Vars,\n)\n integrand.v = -2 * aux.r * aux.v * exp(-aux.v * aux.r^2)\n integrand.r = 1\nend\n\n@inline function integral_set_auxiliary_state!(\n m::IntegralTestSphereModel,\n aux::Vars,\n integral::Vars,\n)\n aux.int.v = integral.v\n aux.int.r = integral.r\nend\n\n@inline function reverse_integral_load_auxiliary_state!(\n m::IntegralTestSphereModel,\n integral::Vars,\n state::Vars,\n aux::Vars,\n)\n integral.v = aux.int.v\n integral.r = aux.int.r\nend\n\n@inline function reverse_integral_set_auxiliary_state!(\n m::IntegralTestSphereModel,\n aux::Vars,\n integral::Vars,\n)\n aux.rev_int.v = integral.v\n aux.rev_int.r = integral.r\nend\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nusing Test\nfunction test_run(mpicomm, topl, ArrayType, N, FT, Rinner, Router)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = Topologies.equiangular_cubed_sphere_warp,\n )\n dg = DGModel(\n IntegralTestSphereModel(Rinner, Router),\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n dQdt = similar(Q)\n\n exact_aux = copy(dg.state_auxiliary)\n dg(dQdt, Q, nothing, 0.0)\n (int_r_ind, rev_int_r_ind) = varsindices(\n vars_state(dg.balance_law, Auxiliary(), FT),\n (\"int.r\", \"rev_int.r\"),\n )\n\n # Since N = 0 integrals live at the faces we need to average values to the\n # faces for comparison\n if N == 0\n\n nvertelem = topl.stacksize\n nhorzelem = div(length(topl.elems), nvertelem)\n naux = size(exact_aux, 2)\n ndof = size(exact_aux, 1)\n\n # Reshape the data array to be (dofs, naux, vertical elm, horizontal elm)\n aux =\n reshape(dg.state_auxiliary.data, (ndof, naux, nvertelem, nhorzelem))\n\n # average forward integrals to cell centers\n for ind in varsindices(\n vars_state(dg.balance_law, Auxiliary(), FT),\n (\"int.r\", \"int.v\"),\n )\n\n # Store the computed face values\n A_faces = aux[:, ind, :, :]\n\n # Bottom cell value is average of 0 and top face of cell\n aux[:, ind, 1, :] .= A_faces[:, 1, :] / 2\n\n # Remaining cells are average of the two faces\n aux[:, ind, 2:end, :] .=\n (A_faces[:, 1:(end - 1), :] + A_faces[:, 2:end, :]) / 2\n end\n\n # average reverse integrals to cell centers\n for ind in varsindices(\n vars_state(dg.balance_law, Auxiliary(), FT),\n (\"rev_int.r\", \"rev_int.v\"),\n )\n\n # Store the computed face values\n RA_faces = aux[:, ind, :, :]\n\n # Bottom cell value is average of 0 and top face of cell\n aux[:, ind, end, :] .= RA_faces[:, end, :] / 2\n\n # Remaining cells are average of the two faces\n aux[:, ind, 1:(end - 1), :] .=\n (RA_faces[:, 1:(end - 1), :] + RA_faces[:, 2:end, :]) / 2\n end\n end\n # We should be exact for the integral of ∫_{R_{inner}}^{r} 1\n @test exact_aux[:, int_r_ind, :] ≈ dg.state_auxiliary[:, int_r_ind, :]\n @test exact_aux[:, rev_int_r_ind, :] ≈\n dg.state_auxiliary[:, rev_int_r_ind, :]\n euclidean_distance(exact_aux, dg.state_auxiliary)\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n base_Nhorz = 4\n base_Nvert = 2\n Rinner = 1 // 2\n Router = 1\n\n expected_result = Dict()\n expected_result[0] = [\n 2.2259657670562167e-02\n 5.6063943176909315e-03\n 1.4042479532664005e-03\n 3.5122834187695408e-04\n ]\n expected_result[1] = [\n 1.5934735012225074e-02\n 4.0030667455285352e-03\n 1.0020652111566574e-03\n 2.5059856392475887e-04\n ]\n expected_result[4] = [\n 4.662884229467401e-7,\n 7.218989778540723e-9,\n 1.1258613174916711e-10,\n 1.7587739986848968e-12,\n ]\n lvls = integration_testing ? length(expected_result[4]) : 1\n\n for N in (0, 1, 4)\n for FT in (Float64,)\n err = zeros(FT, lvls)\n for l in 1:lvls\n @info (ArrayType, FT, \"sphere\", N, l)\n Nhorz = 2^(l - 1) * base_Nhorz\n Nvert = 2^(l - 1) * base_Nvert\n Rrange = grid1d(FT(Rinner), FT(Router); nelem = Nvert)\n topl = StackedCubedSphereTopology(mpicomm, Nhorz, Rrange)\n err[l] = test_run(\n mpicomm,\n topl,\n ArrayType,\n N,\n FT,\n FT(Rinner),\n FT(Router),\n )\n @test expected_result[N][l] ≈ err[l] rtol = 1e-3 atol = eps(FT)\n end\n if lvls > 1\n @info begin\n msg = \"polynomialorder order $N\"\n for l in 1:(lvls - 1)\n rate = log2(err[l]) - log2(err[l + 1])\n msg *= @sprintf(\"\\n rate for level %d = %e\\n\", l, rate)\n end\n msg\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/DGMethods/remainder_model.jl", "content": "using Test\nusing LinearAlgebra\nusing MPI\nusing Random\nusing StaticArrays\n\nusing ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Orientations\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: planet_radius\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n\"\"\"\n main()\n\nRun this test problem\n\"\"\"\nfunction main()\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 5\n numelem_horz = 10\n numelem_vert = 5\n\n @testset \"remainder model\" begin\n for FT in (Float64,)# Float32)\n result = test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n ArrayType,\n FT,\n )\n end\n end\nend\n\nfunction test_run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n ArrayType,\n FT,\n)\n\n # Structure to pass around to setup the simulation\n setup = RemainderTestSetup{FT}()\n\n # Create the cubed sphere mesh\n _planet_radius::FT = planet_radius(param_set)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + setup.domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = equiangular_cubed_sphere_warp,\n )\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n\n # This is the base model which defines all the data (all other DGModels\n # for substepping components will piggy-back off of this models data)\n fullphysics = AtmosPhysics{FT}(\n param_set;\n ref_state = HydrostaticState(T_profile),\n turbulence = Vreman(FT(0.23)),\n moisture = DryModel(),\n )\n fullmodel = AtmosModel{FT}(\n AtmosLESConfigType,\n fullphysics;\n orientation = SphericalOrientation(),\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n dg = DGModel(\n fullmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n Random.seed!(1235)\n Q = init_ode_state(dg, FT(0); init_on_cpu = true)\n\n acousticmodel = AtmosAcousticGravityLinearModel(fullmodel)\n\n acoustic_dg = DGModel(\n acousticmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = EveryDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n vacoustic_dg = DGModel(\n acousticmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n hacoustic_dg = DGModel(\n acousticmodel,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = HorizontalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n\n # Create some random data to check the wavespeed function with\n nM = rand(3)\n nM /= norm(nM)\n state_prognostic = Vars{vars_state(dg.balance_law, Prognostic(), FT)}(rand(\n FT,\n number_states(dg.balance_law, Prognostic()),\n ))\n state_auxiliary = Vars{vars_state(dg.balance_law, Auxiliary(), FT)}(rand(\n FT,\n number_states(dg.balance_law, Auxiliary()),\n ))\n full_wavespeed = wavespeed(\n dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n )\n acoustic_wavespeed = wavespeed(\n acoustic_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n )\n\n # Evaluate the full tendency\n full_tendency = similar(Q)\n dg(full_tendency, Q, nothing, 0; increment = false)\n\n # Evaluate various splittings\n split_tendency = similar(Q)\n\n # Check pulling acoustic model out\n @testset \"full acoustic\" begin\n rem_dg = remainder_DGModel(\n dg,\n (acoustic_dg,);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (acoustic_dg.numerical_flux_first_order,),\n ),\n )\n rem_dg(split_tendency, Q, nothing, 0; increment = false)\n acoustic_dg(split_tendency, Q, nothing, 0; increment = true)\n A = Array(full_tendency.data)\n B = Array(split_tendency.data)\n # Test that we have a fully discrete splitting\n @test all(isapprox.(\n A,\n B,\n rtol = sqrt(eps(FT)),\n atol = 10 * sqrt(eps(FT)),\n ))\n\n # Test that wavespeeds are split by direction\n every_wavespeed = full_wavespeed .- acoustic_wavespeed\n horz_wavespeed = -zero(FT)\n vert_wavespeed = -zero(FT)\n @test all(\n every_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n ),\n )\n @test all(\n horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (HorizontalDirection(),),\n ),\n )\n @test all(\n vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (VerticalDirection(),),\n ),\n )\n @test all(\n every_wavespeed .+ horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), HorizontalDirection()),\n ),\n )\n @test all(\n every_wavespeed .+ vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), VerticalDirection()),\n ),\n )\n end\n\n # Check pulling acoustic model but as two pieces\n @testset \"horizontal and vertical acoustic\" begin\n rem_dg = remainder_DGModel(\n dg,\n (hacoustic_dg, vacoustic_dg);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (\n hacoustic_dg.numerical_flux_first_order,\n vacoustic_dg.numerical_flux_first_order,\n ),\n ),\n )\n rem_dg(split_tendency, Q, nothing, 0; increment = false)\n vacoustic_dg(split_tendency, Q, nothing, 0; increment = true)\n hacoustic_dg(split_tendency, Q, nothing, 0; increment = true)\n A = Array(full_tendency.data)\n B = Array(split_tendency.data)\n # Test that we have a fully discrete splitting\n @test all(isapprox.(\n A,\n B,\n rtol = sqrt(eps(FT)),\n atol = 10 * sqrt(eps(FT)),\n ))\n\n # Test that wavespeeds are split by direction\n every_wavespeed = full_wavespeed\n horz_wavespeed = -acoustic_wavespeed\n vert_wavespeed = -acoustic_wavespeed\n @test all(\n every_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n ),\n )\n @test all(\n horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (HorizontalDirection(),),\n ),\n )\n @test all(\n vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (VerticalDirection(),),\n ),\n )\n @test all(\n every_wavespeed .+ horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), HorizontalDirection()),\n ),\n )\n @test all(\n every_wavespeed .+ vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), VerticalDirection()),\n ),\n )\n end\n\n # Check pulling horizontal acoustic model\n @testset \"horizontal acoustic\" begin\n rem_dg = remainder_DGModel(\n dg,\n (hacoustic_dg,);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (hacoustic_dg.numerical_flux_first_order,),\n ),\n )\n rem_dg(split_tendency, Q, nothing, 0; increment = false)\n hacoustic_dg(split_tendency, Q, nothing, 0; increment = true)\n A = Array(full_tendency.data)\n B = Array(split_tendency.data)\n # Test that we have a fully discrete splitting\n @test all(isapprox.(\n A,\n B,\n rtol = sqrt(eps(FT)),\n atol = 10 * sqrt(eps(FT)),\n ))\n\n # Test that wavespeeds are split by direction\n every_wavespeed = full_wavespeed\n horz_wavespeed = -acoustic_wavespeed\n vert_wavespeed = -zero(eltype(FT))\n @test all(\n every_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n ),\n )\n @test all(\n horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (HorizontalDirection(),),\n ),\n )\n @test all(\n vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (VerticalDirection(),),\n ),\n )\n @test all(\n every_wavespeed .+ horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), HorizontalDirection()),\n ),\n )\n @test all(\n every_wavespeed .+ vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), VerticalDirection()),\n ),\n )\n end\n\n # Check pulling vertical acoustic model\n @testset \"vertical acoustic\" begin\n rem_dg = remainder_DGModel(\n dg,\n (vacoustic_dg,);\n numerical_flux_first_order = (\n dg.numerical_flux_first_order,\n (vacoustic_dg.numerical_flux_first_order,),\n ),\n )\n rem_dg(split_tendency, Q, nothing, 0; increment = false)\n vacoustic_dg(split_tendency, Q, nothing, 0; increment = true)\n A = Array(full_tendency.data)\n B = Array(split_tendency.data)\n # Test that we have a fully discrete splitting\n @test all(isapprox.(\n A,\n B,\n rtol = sqrt(eps(FT)),\n atol = 10 * sqrt(eps(FT)),\n ))\n\n # Test that wavespeeds are split by direction\n every_wavespeed = full_wavespeed\n horz_wavespeed = -zero(eltype(FT))\n vert_wavespeed = -acoustic_wavespeed\n @test all(\n every_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(),),\n ),\n )\n @test all(\n horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (HorizontalDirection(),),\n ),\n )\n @test all(\n vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (VerticalDirection(),),\n ),\n )\n @test all(\n every_wavespeed .+ horz_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), HorizontalDirection()),\n ),\n )\n @test all(\n every_wavespeed .+ vert_wavespeed .≈ wavespeed(\n rem_dg.balance_law,\n nM,\n state_prognostic,\n state_auxiliary,\n FT(0),\n (EveryDirection(), VerticalDirection()),\n ),\n )\n end\nend\n\nBase.@kwdef struct RemainderTestSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\nend\n\nfunction (setup::RemainderTestSetup)(problem, bl, state, aux, localgeo, t)\n FT = eltype(state)\n\n # Vary around the reference state by 10% and a random velocity field\n state.ρ = (4 + rand(FT)) / 5 + aux.ref_state.ρ\n state.ρu = 2 * (@SVector rand(FT, 3)) .- 1\n state.energy.ρe = (4 + rand(FT)) / 5 + aux.ref_state.ρe\n\n nothing\nend\n\nmain()\n" }, { "path": "test/Numerics/DGMethods/vars_test.jl", "content": "using MPI\nusing StaticArrays\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing Printf\nusing LinearAlgebra\nusing Logging\n\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, GradientFlux\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n wavespeed,\n update_auxiliary_state!,\n boundary_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!\n\nimport ClimateMachine.DGMethods: init_ode_state\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\n\nstruct VarsTestModel{dim} <: BalanceLaw end\n\nvars_state(::VarsTestModel, ::Prognostic, T) = @vars(x::T, coord::SVector{3, T})\nvars_state(m::VarsTestModel, ::Auxiliary, T) =\n @vars(coord::SVector{3, T}, polynomial::T)\nvars_state(m::VarsTestModel, ::GradientFlux, T) = @vars()\n\nflux_first_order!(::VarsTestModel, _...) = nothing\nflux_second_order!(::VarsTestModel, _...) = nothing\nsource!(::VarsTestModel, _...) = nothing\nboundary_state!(_, ::VarsTestModel, _...) = nothing\nwavespeed(::VarsTestModel, _...) = 1\n\nfunction init_state_prognostic!(\n m::VarsTestModel,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n @inbounds state.x = localgeo.coord[1]\n state.coord = localgeo.coord\nend\n\nfunction nodal_init_state_auxiliary!(\n ::VarsTestModel{dim},\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n) where {dim}\n x, y, z = aux.coord = g.coord\n aux.polynomial = x * y + x * z + y * z\nend\n\nusing Test\nfunction test_run(mpicomm, dim, Ne, N, FT, ArrayType)\n\n brickrange = ntuple(j -> range(FT(0); length = Ne[j] + 1, stop = 3), dim)\n connectivity = dim == 3 ? :full : :face\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n connectivity = connectivity,\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n dg = DGModel(\n VarsTestModel{dim}(),\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n @test Array(Q.x)[:, 1, :] == Array(Q.coord)[:, 1, :]\n @test Array(dg.state_auxiliary.coord) == Array(Q.coord)\n x = Array(Q.coord)[:, 1, :]\n y = Array(Q.coord)[:, 2, :]\n z = Array(Q.coord)[:, 3, :]\n @test Array(dg.state_auxiliary.polynomial)[:, 1, :] ≈\n x .* y + x .* z + y .* z\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n numelem = (5, 5, 5)\n lvls = 1\n\n polynomialorder = 4\n\n for FT in (Float64,) #Float32)\n for dim in 2:3\n err = zeros(FT, lvls)\n for l in 1:lvls\n @info (ArrayType, FT, dim)\n test_run(\n mpicomm,\n dim,\n ntuple(j -> 2^(l - 1) * numelem[j], dim),\n polynomialorder,\n FT,\n ArrayType,\n )\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/ESDGMethods/DryAtmos/DryAtmos.jl", "content": "using ClimateMachine.VariableTemplates: Vars, Grad, @vars\nusing ClimateMachine.BalanceLaws\nimport ClimateMachine.BalanceLaws:\n BalanceLaw,\n vars_state,\n state_to_entropy_variables!,\n entropy_variables_to_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n state_to_entropy,\n boundary_conditions,\n boundary_state!,\n wavespeed,\n flux_first_order!,\n source!\n\nusing StaticArrays\nusing LinearAlgebra: dot, I\nusing ClimateMachine.DGMethods.NumericalFluxes: NumericalFluxFirstOrder\nimport ClimateMachine.DGMethods.NumericalFluxes:\n EntropyConservative,\n numerical_volume_conservative_flux_first_order!,\n numerical_volume_fluctuation_flux_first_order!,\n ave,\n logave,\n numerical_flux_first_order!,\n numerical_flux_second_order!,\n numerical_boundary_flux_second_order!\n\nusing ClimateMachine.Orientations:\n Orientation, FlatOrientation, SphericalOrientation\nusing ClimateMachine.Atmos: NoReferenceState\nusing ClimateMachine.Grids\n\nusing CLIMAParameters: AbstractEarthParameterSet\nusing CLIMAParameters.Planet: grav, R_d, cp_d, cv_d, planet_radius, MSLP, Omega\n\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nconst total_energy = false\nconst fluctuation_gravity = false\n\n@inline gamma(ps::EarthParameterSet) = cp_d(ps) / cv_d(ps)\n\nabstract type AbstractDryAtmosProblem end\n\nstruct DryAtmosModel{D, O, P, RS, S, DS} <: BalanceLaw\n orientation::O\n problem::P\n ref_state::RS\n sources::S\n drag_source::DS\nend\nfunction DryAtmosModel{D}(\n orientation,\n problem::AbstractDryAtmosProblem;\n ref_state = NoReferenceState(),\n sources = (),\n drag_source = NoDrag(),\n) where {D}\n O = typeof(orientation)\n P = typeof(problem)\n RS = typeof(ref_state)\n S = typeof(sources)\n DS = typeof(drag_source)\n DryAtmosModel{D, O, P, RS, S, DS}(\n orientation,\n problem,\n ref_state,\n sources,\n drag_source,\n )\nend\n\nboundary_conditions(::DryAtmosModel) = (1, 2)\n# XXX: Hack for Impenetrable.\n# This is NOT entropy stable / conservative!!!!\nfunction boundary_state!(\n ::NumericalFluxFirstOrder,\n bctype,\n ::DryAtmosModel,\n state⁺,\n aux⁺,\n n,\n state⁻,\n aux⁻,\n _...,\n)\n state⁺.ρ = state⁻.ρ\n state⁺.ρu -= 2 * dot(state⁻.ρu, n) .* SVector(n)\n state⁺.ρe = state⁻.ρe\n aux⁺.Φ = aux⁻.Φ\nend\n\nfunction init_state_prognostic!(m::DryAtmosModel, args...)\n init_state_prognostic!(m, m.problem, args...)\nend\n\nfunction nodal_init_state_auxiliary!(\n m::DryAtmosModel,\n state_auxiliary,\n tmp,\n geom,\n)\n init_state_auxiliary!(m, m.orientation, state_auxiliary, geom)\n init_state_auxiliary!(m, m.ref_state, state_auxiliary, geom)\n init_state_auxiliary!(m, m.problem, state_auxiliary, geom)\nend\n\nfunction altitude(::DryAtmosModel{dim}, ::FlatOrientation, geom) where {dim}\n @inbounds geom.coord[dim]\nend\n\nfunction altitude(::DryAtmosModel, ::SphericalOrientation, geom)\n FT = eltype(geom)\n _planet_radius::FT = planet_radius(param_set)\n norm(geom.coord) - _planet_radius\nend\n\n\"\"\"\n init_state_auxiliary!(\n m::DryAtmosModel,\n aux::Vars,\n geom::LocalGeometry\n )\n\nInitialize geopotential for the `DryAtmosModel`.\n\"\"\"\nfunction init_state_auxiliary!(\n ::DryAtmosModel{dim},\n ::FlatOrientation,\n state_auxiliary,\n geom,\n) where {dim}\n FT = eltype(state_auxiliary)\n _grav = FT(grav(param_set))\n @inbounds r = geom.coord[dim]\n state_auxiliary.Φ = _grav * r\n state_auxiliary.∇Φ =\n dim == 2 ? SVector{3, FT}(0, _grav, 0) : SVector{3, FT}(0, 0, _grav)\nend\nfunction init_state_auxiliary!(\n ::DryAtmosModel,\n ::SphericalOrientation,\n state_auxiliary,\n geom,\n)\n FT = eltype(state_auxiliary)\n _grav = FT(grav(param_set))\n r = norm(geom.coord)\n state_auxiliary.Φ = _grav * r\n state_auxiliary.∇Φ = _grav * geom.coord / r\nend\n\nfunction init_state_auxiliary!(\n ::DryAtmosModel,\n ::NoReferenceState,\n state_auxiliary,\n geom,\n) end\n\nfunction init_state_auxiliary!(\n ::DryAtmosModel,\n ::AbstractDryAtmosProblem,\n state_auxiliary,\n geom,\n) end\n\nstruct DryReferenceState{TP}\n temperature_profile::TP\nend\nvars_state(::DryAtmosModel, ::DryReferenceState, ::Auxiliary, FT) =\n @vars(T::FT, p::FT, ρ::FT, ρe::FT)\nvars_state(::DryAtmosModel, ::NoReferenceState, ::Auxiliary, FT) = @vars()\n\nfunction init_state_auxiliary!(\n m::DryAtmosModel,\n ref_state::DryReferenceState,\n state_auxiliary,\n geom,\n)\n FT = eltype(state_auxiliary)\n z = altitude(m, m.orientation, geom)\n T, p = ref_state.temperature_profile(param_set, z)\n\n _R_d::FT = R_d(param_set)\n ρ = p / (_R_d * T)\n Φ = state_auxiliary.Φ\n ρu = SVector{3, FT}(0, 0, 0)\n\n state_auxiliary.ref_state.T = T\n state_auxiliary.ref_state.p = p\n state_auxiliary.ref_state.ρ = ρ\n state_auxiliary.ref_state.ρe = totalenergy(ρ, ρu, p, Φ)\nend\n\n@inline function flux_first_order!(\n m::DryAtmosModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n ρ = state.ρ\n ρinv = 1 / ρ\n ρu = state.ρu\n ρe = state.ρe\n u = ρinv * ρu\n Φ = aux.Φ\n\n p = pressure(ρ, ρu, ρe, Φ)\n\n flux.ρ = ρ * u\n flux.ρu = p * I + ρ * u .* u'\n flux.ρe = u * (state.ρe + p)\nend\n\nfunction wavespeed(\n ::DryAtmosModel,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρe = state.ρe\n Φ = aux.Φ\n p = pressure(ρ, ρu, ρe, Φ)\n\n u = ρu / ρ\n uN = abs(dot(nM, u))\n return uN + soundspeed(ρ, p)\nend\n\n\"\"\"\n pressure(ρ, ρu, ρe, Φ)\n\nCompute the pressure given density `ρ`, momentum `ρu`, total energy `ρe`, and\ngravitational potential `Φ`.\n\"\"\"\nfunction pressure(ρ, ρu, ρe, Φ)\n FT = eltype(ρ)\n γ = FT(gamma(param_set))\n if total_energy\n (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ - ρ * Φ)\n else\n (γ - 1) * (ρe - dot(ρu, ρu) / 2ρ)\n end\nend\n\n\"\"\"\n totalenergy(ρ, ρu, p, Φ)\n\nCompute the total energy given density `ρ`, momentum `ρu`, pressure `p`, and\ngravitational potential `Φ`.\n\"\"\"\nfunction totalenergy(ρ, ρu, p, Φ)\n FT = eltype(ρ)\n γ = FT(gamma(param_set))\n if total_energy\n return p / (γ - 1) + dot(ρu, ρu) / 2ρ + ρ * Φ\n else\n return p / (γ - 1) + dot(ρu, ρu) / 2ρ\n end\nend\n\n\"\"\"\n soundspeed(ρ, p)\n\nCompute the speed of sound from the density `ρ` and pressure `p`.\n\"\"\"\nfunction soundspeed(ρ, p)\n FT = eltype(ρ)\n γ = FT(gamma(param_set))\n sqrt(γ * p / ρ)\nend\n\n\"\"\"\n vars_state(::DryAtmosModel, ::Prognostic, FT)\n\nThe prognostic state variables for the `DryAtmosModel` are density `ρ`, momentum `ρu`,\nand total energy `ρe`\n\"\"\"\nfunction vars_state(::DryAtmosModel, ::Prognostic, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n end\nend\n\n\"\"\"\n vars_state(::DryAtmosModel, ::Auxiliary, FT)\n\nThe auxiliary variables for the `DryAtmosModel` is gravitational potential\n`Φ`\n\"\"\"\nfunction vars_state(m::DryAtmosModel, st::Auxiliary, FT)\n @vars begin\n Φ::FT\n ∇Φ::SVector{3, FT} # TODO: only needed for the linear model\n ref_state::vars_state(m, m.ref_state, st, FT)\n problem::vars_state(m, m.problem, st, FT)\n end\nend\nvars_state(::DryAtmosModel, ::AbstractDryAtmosProblem, ::Auxiliary, FT) =\n @vars()\n\n\"\"\"\n vars_state(::DryAtmosModel, ::Entropy, FT)\n\nThe entropy variables for the `DryAtmosModel` correspond to the state\nvariables density `ρ`, momentum `ρu`, and total energy `ρe` as well as the\nauxiliary variable gravitational potential `Φ`\n\"\"\"\nfunction vars_state(::DryAtmosModel, ::Entropy, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n Φ::FT\n end\nend\n\n\"\"\"\n state_to_entropy_variables!(\n ::DryAtmosModel,\n entropy::Vars,\n state::Vars,\n aux::Vars,\n )\n\nSee [`BalanceLaws.state_to_entropy_variables!`](@ref)\n\"\"\"\nfunction state_to_entropy_variables!(\n ::DryAtmosModel,\n entropy::Vars,\n state::Vars,\n aux::Vars,\n)\n ρ, ρu, ρe, Φ = state.ρ, state.ρu, state.ρe, aux.Φ\n\n FT = eltype(state)\n γ = FT(gamma(param_set))\n\n p = pressure(ρ, ρu, ρe, Φ)\n s = log(p / ρ^γ)\n b = ρ / 2p\n u = ρu / ρ\n\n if total_energy\n entropy.ρ = (γ - s) / (γ - 1) - (dot(u, u) - 2Φ) * b\n else\n entropy.ρ = (γ - s) / (γ - 1) - (dot(u, u)) * b\n end\n entropy.ρu = 2b * u\n entropy.ρe = -2b\n entropy.Φ = 2ρ * b\nend\n\n\"\"\"\n entropy_variables_to_state!(\n ::DryAtmosModel,\n state::Vars,\n aux::Vars,\n entropy::Vars,\n )\n\nSee [`BalanceLaws.entropy_variables_to_state!`](@ref)\n\"\"\"\nfunction entropy_variables_to_state!(\n ::DryAtmosModel,\n state::Vars,\n aux::Vars,\n entropy::Vars,\n)\n FT = eltype(state)\n β = entropy\n γ = FT(gamma(param_set))\n\n b = -β.ρe / 2\n ρ = β.Φ / (2b)\n ρu = ρ * β.ρu / (2b)\n\n p = ρ / (2b)\n s = log(p / ρ^γ)\n Φ = dot(ρu, ρu) / (2 * ρ^2) - ((γ - s) / (γ - 1) - β.ρ) / (2b)\n\n ρe = p / (γ - 1) + dot(ρu, ρu) / (2ρ) + ρ * Φ\n\n state.ρ = ρ\n state.ρu = ρu\n state.ρe = ρe\n aux.Φ = Φ\nend\n\nfunction state_to_entropy(::DryAtmosModel, state::Vars, aux::Vars)\n FT = eltype(state)\n ρ, ρu, ρe, Φ = state.ρ, state.ρu, state.ρe, aux.Φ\n p = pressure(ρ, ρu, ρe, Φ)\n γ = FT(gamma(param_set))\n s = log(p / ρ^γ)\n η = -ρ * s / (γ - 1)\n return η\nend\n\nfunction numerical_volume_conservative_flux_first_order!(\n ::EntropyConservative,\n ::DryAtmosModel,\n F::Grad,\n state_1::Vars,\n aux_1::Vars,\n state_2::Vars,\n aux_2::Vars,\n)\n FT = eltype(F)\n ρ_1, ρu_1, ρe_1 = state_1.ρ, state_1.ρu, state_1.ρe\n ρ_2, ρu_2, ρe_2 = state_2.ρ, state_2.ρu, state_2.ρe\n Φ_1, Φ_2 = aux_1.Φ, aux_2.Φ\n u_1 = ρu_1 / ρ_1\n u_2 = ρu_2 / ρ_2\n p_1 = pressure(ρ_1, ρu_1, ρe_1, Φ_1)\n p_2 = pressure(ρ_2, ρu_2, ρe_2, Φ_2)\n b_1 = ρ_1 / 2p_1\n b_2 = ρ_2 / 2p_2\n\n ρ_avg = ave(ρ_1, ρ_2)\n u_avg = ave(u_1, u_2)\n b_avg = ave(b_1, b_2)\n Φ_avg = ave(Φ_1, Φ_2)\n\n usq_avg = ave(dot(u_1, u_1), dot(u_2, u_2))\n\n ρ_log = logave(ρ_1, ρ_2)\n b_log = logave(b_1, b_2)\n α = b_avg * ρ_log / 2b_1\n\n γ = FT(gamma(param_set))\n\n Fρ = u_avg * ρ_log\n Fρu = u_avg * Fρ' + ρ_avg / 2b_avg * I\n if total_energy\n Fρe =\n (1 / (2 * (γ - 1) * b_log) - usq_avg / 2 + Φ_avg) * Fρ + Fρu * u_avg\n else\n Fρe = (1 / (2 * (γ - 1) * b_log) - usq_avg / 2) * Fρ + Fρu * u_avg\n end\n\n F.ρ += Fρ\n F.ρu += Fρu\n F.ρe += Fρe\nend\n\nfunction numerical_volume_fluctuation_flux_first_order!(\n ::NumericalFluxFirstOrder,\n ::DryAtmosModel,\n D::Grad,\n state_1::Vars,\n aux_1::Vars,\n state_2::Vars,\n aux_2::Vars,\n)\n if fluctuation_gravity\n FT = eltype(D)\n ρ_1, ρu_1, ρe_1 = state_1.ρ, state_1.ρu, state_1.ρe\n ρ_2, ρu_2, ρe_2 = state_2.ρ, state_2.ρu, state_2.ρe\n Φ_1, Φ_2 = aux_1.Φ, aux_2.Φ\n p_1 = pressure(ρ_1, ρu_1, ρe_1, Φ_1)\n p_2 = pressure(ρ_2, ρu_2, ρe_2, Φ_2)\n b_1 = ρ_1 / 2p_1\n b_2 = ρ_2 / 2p_2\n\n ρ_log = logave(ρ_1, ρ_2)\n b_avg = ave(b_1, b_2)\n α = b_avg * ρ_log / 2b_1\n\n D.ρu -= α * (Φ_1 - Φ_2) * I\n end\nend\n\nstruct CentralVolumeFlux <: NumericalFluxFirstOrder end\nfunction numerical_volume_conservative_flux_first_order!(\n ::CentralVolumeFlux,\n m::DryAtmosModel,\n F::Grad,\n state_1::Vars,\n aux_1::Vars,\n state_2::Vars,\n aux_2::Vars,\n)\n FT = eltype(F)\n F_1 = similar(F)\n flux_first_order!(m, F_1, state_1, aux_1, FT(0), EveryDirection())\n\n F_2 = similar(F)\n flux_first_order!(m, F_2, state_2, aux_2, FT(0), EveryDirection())\n\n parent(F) .= (parent(F_1) .+ parent(F_2)) ./ 2\nend\n\nstruct KGVolumeFlux <: NumericalFluxFirstOrder end\nfunction numerical_volume_conservative_flux_first_order!(\n ::KGVolumeFlux,\n m::DryAtmosModel,\n F::Grad,\n state_1::Vars,\n aux_1::Vars,\n state_2::Vars,\n aux_2::Vars,\n)\n Φ_1 = aux_1.Φ\n ρ_1 = state_1.ρ\n ρu_1 = state_1.ρu\n ρe_1 = state_1.ρe\n u_1 = ρu_1 / ρ_1\n e_1 = ρe_1 / ρ_1\n p_1 = pressure(ρ_1, ρu_1, ρe_1, Φ_1)\n\n Φ_2 = aux_2.Φ\n ρ_2 = state_2.ρ\n ρu_2 = state_2.ρu\n ρe_2 = state_2.ρe\n u_2 = ρu_2 / ρ_2\n e_2 = ρe_2 / ρ_2\n p_2 = pressure(ρ_2, ρu_2, ρe_2, Φ_2)\n\n ρ_avg = ave(ρ_1, ρ_2)\n u_avg = ave(u_1, u_2)\n e_avg = ave(e_1, e_2)\n p_avg = ave(p_1, p_2)\n\n F.ρ = ρ_avg * u_avg\n F.ρu = p_avg * I + ρ_avg * u_avg .* u_avg'\n F.ρe = ρ_avg * u_avg * e_avg + p_avg * u_avg\nend\n\n\nstruct Coriolis end\nfunction source!(\n m::DryAtmosModel,\n ::Coriolis,\n source,\n state_prognostic,\n state_auxiliary,\n)\n FT = eltype(state_prognostic)\n _Omega::FT = Omega(param_set)\n # note: this assumes a SphericalOrientation\n source.ρu -= SVector(0, 0, 2 * _Omega) × state_prognostic.ρu\nend\n\nfunction source!(m::DryAtmosModel, source, state_prognostic, state_auxiliary)\n ntuple(Val(length(m.sources))) do s\n Base.@_inline_meta\n source!(m, m.sources[s], source, state_prognostic, state_auxiliary)\n end\nend\n\n\nstruct EntropyConservativeWithPenalty <: NumericalFluxFirstOrder end\nfunction numerical_flux_first_order!(\n numerical_flux::EntropyConservativeWithPenalty,\n balance_law::BalanceLaw,\n fluxᵀn::Vars{S},\n normal_vector::SVector,\n state_prognostic⁻::Vars{S},\n state_auxiliary⁻::Vars{A},\n state_prognostic⁺::Vars{S},\n state_auxiliary⁺::Vars{A},\n t,\n direction,\n) where {S, A}\n\n FT = eltype(fluxᵀn)\n\n numerical_flux_first_order!(\n EntropyConservative(),\n balance_law,\n fluxᵀn,\n normal_vector,\n state_prognostic⁻,\n state_auxiliary⁻,\n state_prognostic⁺,\n state_auxiliary⁺,\n t,\n direction,\n )\n fluxᵀn = parent(fluxᵀn)\n\n wavespeed⁻ = wavespeed(\n balance_law,\n normal_vector,\n state_prognostic⁻,\n state_auxiliary⁻,\n t,\n direction,\n )\n wavespeed⁺ = wavespeed(\n balance_law,\n normal_vector,\n state_prognostic⁺,\n state_auxiliary⁺,\n t,\n direction,\n )\n max_wavespeed = max.(wavespeed⁻, wavespeed⁺)\n penalty =\n max_wavespeed .* (parent(state_prognostic⁻) - parent(state_prognostic⁺))\n\n fluxᵀn .+= penalty / 2\nend\n\nBase.@kwdef struct MatrixFlux{FT} <: NumericalFluxFirstOrder\n Mcut::FT = 0\n low_mach::Bool = false\n kinetic_energy_preserving::Bool = false\nend\n\nfunction numerical_flux_first_order!(\n numerical_flux::MatrixFlux,\n balance_law::BalanceLaw,\n fluxᵀn::Vars{S},\n normal_vector::SVector,\n state_prognostic⁻::Vars{S},\n state_auxiliary⁻::Vars{A},\n state_prognostic⁺::Vars{S},\n state_auxiliary⁺::Vars{A},\n t,\n direction,\n) where {S, A}\n\n FT = eltype(fluxᵀn)\n numerical_flux_first_order!(\n EntropyConservative(),\n balance_law,\n fluxᵀn,\n normal_vector,\n state_prognostic⁻,\n state_auxiliary⁻,\n state_prognostic⁺,\n state_auxiliary⁺,\n t,\n direction,\n )\n fluxᵀn = parent(fluxᵀn)\n\n γ = FT(gamma(param_set))\n\n low_mach = numerical_flux.low_mach\n Mcut = numerical_flux.Mcut\n kinetic_energy_preserving = numerical_flux.kinetic_energy_preserving\n\n ω = FT(π) / 3\n δ = FT(π) / 5\n random_unit_vector = SVector(sin(ω) * cos(δ), cos(ω) * cos(δ), sin(δ))\n # tangent space basis\n τ1 = random_unit_vector × normal_vector\n τ2 = τ1 × normal_vector\n\n ρ⁻ = state_prognostic⁻.ρ\n ρu⁻ = state_prognostic⁻.ρu\n ρe⁻ = state_prognostic⁻.ρe\n\n Φ⁻ = state_auxiliary⁻.Φ\n u⁻ = ρu⁻ / ρ⁻\n p⁻ = pressure(ρ⁻, ρu⁻, ρe⁻, Φ⁻)\n β⁻ = ρ⁻ / 2p⁻\n\n Φ⁺ = state_auxiliary⁺.Φ\n ρ⁺ = state_prognostic⁺.ρ\n ρu⁺ = state_prognostic⁺.ρu\n ρe⁺ = state_prognostic⁺.ρe\n\n u⁺ = ρu⁺ / ρ⁺\n p⁺ = pressure(ρ⁺, ρu⁺, ρe⁺, Φ⁺)\n β⁺ = ρ⁺ / 2p⁺\n\n ρ_log = logave(ρ⁻, ρ⁺)\n β_log = logave(β⁻, β⁺)\n if total_energy\n Φ_avg = ave(Φ⁻, Φ⁺)\n else\n Φ_avg = 0\n end\n u_avg = ave.(u⁻, u⁺)\n p_avg = ave(ρ⁻, ρ⁺) / 2ave(β⁻, β⁺)\n u²_bar = 2 * sum(u_avg .^ 2) - sum(ave(u⁻ .^ 2, u⁺ .^ 2))\n\n h_bar = γ / (2 * β_log * (γ - 1)) + u²_bar / 2 + Φ_avg\n c_bar = sqrt(γ * p_avg / ρ_log)\n\n umc = u_avg - c_bar * normal_vector\n upc = u_avg + c_bar * normal_vector\n u_avgᵀn = u_avg' * normal_vector\n R = hcat(\n SVector(1, umc[1], umc[2], umc[3], h_bar - c_bar * u_avgᵀn),\n SVector(1, u_avg[1], u_avg[2], u_avg[3], u²_bar / 2 + Φ_avg),\n SVector(0, τ1[1], τ1[2], τ1[3], τ1' * u_avg),\n SVector(0, τ2[1], τ2[2], τ2[3], τ2' * u_avg),\n SVector(1, upc[1], upc[2], upc[3], h_bar + c_bar * u_avgᵀn),\n )\n\n if low_mach\n M = abs(u_avg' * normal_vector) / c_bar\n c_bar *= max(min(M, FT(1)), Mcut)\n end\n\n if kinetic_energy_preserving\n λl = abs(u_avgᵀn) + c_bar\n λr = λl\n else\n λl = abs(u_avgᵀn - c_bar)\n λr = abs(u_avgᵀn + c_bar)\n end\n\n Λ = SDiagonal(λl, abs(u_avgᵀn), abs(u_avgᵀn), abs(u_avgᵀn), λr)\n #Ξ = sqrt(abs((p⁺ - p⁻) / (p⁺ + p⁻)))\n #Λ = Ξ * abs(u_avgᵀn + c_bar) * I + (1 - Ξ) * ΛM\n\n T = SDiagonal(ρ_log / 2γ, ρ_log * (γ - 1) / γ, p_avg, p_avg, ρ_log / 2γ)\n\n entropy⁻ = similar(parent(state_prognostic⁻), Size(6))\n state_to_entropy_variables!(\n balance_law,\n Vars{vars_state(balance_law, Entropy(), FT)}(entropy⁻),\n state_prognostic⁻,\n state_auxiliary⁻,\n )\n\n entropy⁺ = similar(parent(state_prognostic⁺), Size(6))\n state_to_entropy_variables!(\n balance_law,\n Vars{vars_state(balance_law, Entropy(), FT)}(entropy⁺),\n state_prognostic⁺,\n state_auxiliary⁺,\n )\n\n Δentropy = parent(entropy⁺) - parent(entropy⁻)\n\n fluxᵀn .-= R * Λ * T * R' * Δentropy[SOneTo(5)] / 2\nend\n\nfunction vertical_unit_vector(m::DryAtmosModel, aux::Vars)\n FT = eltype(aux)\n aux.∇Φ / FT(grav(param_set))\nend\n\nnorm_u(state::Vars, k̂::AbstractVector, ::VerticalDirection) =\n abs(dot(state.ρu, k̂)) / state.ρ\nnorm_u(state::Vars, k̂::AbstractVector, ::HorizontalDirection) =\n norm((state.ρu .- dot(state.ρu, k̂) * k̂) / state.ρ)\nnorm_u(state::Vars, k̂::AbstractVector, ::Direction) = norm(state.ρu / state.ρ)\n\nfunction advective_courant(\n m::DryAtmosModel,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n t,\n direction,\n)\n k̂ = vertical_unit_vector(m, aux)\n normu = norm_u(state, k̂, direction)\n return Δt * normu / Δx\nend\n\nfunction nondiffusive_courant(\n m::DryAtmosModel,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n t,\n direction,\n)\n ρ = state.ρ\n ρu = state.ρu\n ρe = state.ρe\n Φ = aux.Φ\n p = pressure(ρ, ρu, ρe, Φ)\n\n k̂ = vertical_unit_vector(m, aux)\n normu = norm_u(state, k̂, direction)\n ss = soundspeed(ρ, p)\n return Δt * (normu + ss) / Δx\nend\n\nfunction drag_source!(m::DryAtmosModel, args...)\n drag_source!(m, m.drag_source, args...)\nend\nstruct NoDrag end\ndrag_source!(m::DryAtmosModel, ::NoDrag, args...) = nothing\n\nstruct Gravity end\nfunction source!(m::DryAtmosModel, ::Gravity, source, state, aux)\n ∇Φ = aux.∇Φ\n if !fluctuation_gravity\n source.ρu -= state.ρ * ∇Φ\n end\n if !total_energy\n source.ρe -= state.ρu' * ∇Φ\n end\nend\n\n# Numerical Flux \n#=\nnumerical_flux_first_order!(::Nothing, _...) = nothing\nnumerical_boundary_flux_first_order!(::Nothing, _...) = nothing\nnumerical_flux_second_order!(::Nothing, _...) = nothing\nnumerical_boundary_flux_second_order!(::Nothing, _...) = nothing\n=#\nnumerical_flux_first_order!(::Nothing, ::DryAtmosModel, _...) = nothing\nnumerical_flux_second_order!(::Nothing, ::DryAtmosModel, _...) = nothing\nnumerical_boundary_flux_second_order!(::Nothing, a, ::DryAtmosModel, _...) =\n nothing\n\ninclude(\"linear.jl\")\n\n# Throwing this in for convenience\nfunction cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\n function f(sR, ξ, η)\n X, Y = tan(π * ξ / 4), tan(π * η / 4)\n x1 = sR / sqrt(X^2 + Y^2 + 1)\n x2, x3 = X * x1, Y * x1\n x1, x2, x3\n end\n\n fdim = argmax(abs.((a, b, c)))\n if fdim == 1 && a < 0\n # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(-R, b / a, c / a)\n elseif fdim == 2 && b < 0\n # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(-R, a / b, c / b)\n elseif fdim == 1 && a > 0\n # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(R, b / a, c / a)\n elseif fdim == 2 && b > 0\n # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(R, a / b, c / b)\n elseif fdim == 3 && c > 0\n # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(R, b / c, a / c)\n elseif fdim == 3 && c < 0\n # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(-R, b / c, a / c)\n else\n error(\"invalid case for cubedshellwarp: $a, $b, $c\")\n end\n\n return x1, x2, x3\nend\n" }, { "path": "test/Numerics/ESDGMethods/DryAtmos/baroclinic_wave.jl", "content": "using ClimateMachine\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Mesh.Topologies: StackedCubedSphereTopology, grid1d\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.Atmos: AtmosFilterPerturbations\nusing ClimateMachine.DGMethods: ESDGModel, init_ode_state, courant\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.VTK: writevtk, writepvtu\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing Thermodynamics: soundspeed_air\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.VariableTemplates: flattenednames\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: R_d, cv_d, Omega, planet_radius, MSLP\nimport CLIMAParameters\n\nusing MPI, Logging, StaticArrays, LinearAlgebra, Printf, Dates, Test\nusing CUDA\n\nconst output_vtk = false\n\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n#const X = 20\nconst X = 1\nCLIMAParameters.Planet.planet_radius(::EarthParameterSet) = 6.371e6 / X\nCLIMAParameters.Planet.Omega(::EarthParameterSet) = 7.2921159e-5 * X\n\n# No this isn't great but w/e \n\ninclude(\"DryAtmos.jl\")\n\nfunction sphr_to_cart_vec(vec, lat, lon)\n FT = eltype(vec)\n slat, clat = sin(lat), cos(lat)\n slon, clon = sin(lon), cos(lon)\n u = MVector{3, FT}(\n -slon * vec[1] - slat * clon * vec[2] + clat * clon * vec[3],\n clon * vec[1] - slat * slon * vec[2] + clat * slon * vec[3],\n clat * vec[2] + slat * vec[3],\n )\n return u\nend\n\nstruct BaroclinicWave <: AbstractDryAtmosProblem end\n\nfunction init_state_prognostic!(\n bl::DryAtmosModel,\n ::BaroclinicWave,\n state,\n aux,\n localgeo,\n t,\n)\n coords = localgeo.coord\n FT = eltype(state)\n\n # parameters\n _grav::FT = grav(param_set)\n _R_d::FT = R_d(param_set)\n _cv_d::FT = cv_d(param_set)\n _Ω::FT = Omega(param_set)\n _a::FT = planet_radius(param_set)\n _p_0::FT = MSLP(param_set)\n\n k::FT = 3\n T_E::FT = 310\n T_P::FT = 240\n T_0::FT = 0.5 * (T_E + T_P)\n Γ::FT = 0.005\n A::FT = 1 / Γ\n B::FT = (T_0 - T_P) / T_0 / T_P\n C::FT = 0.5 * (k + 2) * (T_E - T_P) / T_E / T_P\n b::FT = 2\n H::FT = _R_d * T_0 / _grav\n z_t::FT = 15e3\n λ_c::FT = π / 9\n φ_c::FT = 2 * π / 9\n d_0::FT = _a / 6\n V_p::FT = 1\n M_v::FT = 0.608\n p_w::FT = 34e3 ## Pressure width parameter for specific humidity\n η_crit::FT = 10 * _p_0 / p_w ## Critical pressure coordinate\n q_0::FT = 0 ## Maximum specific humidity (default: 0.018)\n q_t::FT = 1e-12 ## Specific humidity above artificial tropopause\n φ_w::FT = 2π / 9 ## Specific humidity latitude wind parameter\n\n # grid\n λ = @inbounds atan(coords[2], coords[1])\n φ = @inbounds asin(coords[3] / norm(coords, 2))\n z = norm(coords) - _a\n\n r::FT = z + _a\n γ::FT = 1 # set to 0 for shallow-atmosphere case and to 1 for deep atmosphere case\n\n # convenience functions for temperature and pressure\n τ_z_1::FT = exp(Γ * z / T_0)\n τ_z_2::FT = 1 - 2 * (z / b / H)^2\n τ_z_3::FT = exp(-(z / b / H)^2)\n τ_1::FT = 1 / T_0 * τ_z_1 + B * τ_z_2 * τ_z_3\n τ_2::FT = C * τ_z_2 * τ_z_3\n τ_int_1::FT = A * (τ_z_1 - 1) + B * z * τ_z_3\n τ_int_2::FT = C * z * τ_z_3\n I_T::FT =\n (cos(φ) * (1 + γ * z / _a))^k -\n k / (k + 2) * (cos(φ) * (1 + γ * z / _a))^(k + 2)\n\n # base state virtual temperature, pressure, specific humidity, density\n T_v::FT = (τ_1 - τ_2 * I_T)^(-1)\n p::FT = _p_0 * exp(-_grav / _R_d * (τ_int_1 - τ_int_2 * I_T))\n\n # base state velocity\n U::FT =\n _grav * k / _a *\n τ_int_2 *\n T_v *\n (\n (cos(φ) * (1 + γ * z / _a))^(k - 1) -\n (cos(φ) * (1 + γ * z / _a))^(k + 1)\n )\n u_ref::FT =\n -_Ω * (_a + γ * z) * cos(φ) +\n sqrt((_Ω * (_a + γ * z) * cos(φ))^2 + (_a + γ * z) * cos(φ) * U)\n v_ref::FT = 0\n w_ref::FT = 0\n\n # velocity perturbations\n F_z::FT = 1 - 3 * (z / z_t)^2 + 2 * (z / z_t)^3\n if z > z_t\n F_z = FT(0)\n end\n d::FT = _a * acos(sin(φ) * sin(φ_c) + cos(φ) * cos(φ_c) * cos(λ - λ_c))\n c3::FT = cos(π * d / 2 / d_0)^3\n s1::FT = sin(π * d / 2 / d_0)\n if 0 < d < d_0 && d != FT(_a * π)\n u′::FT =\n -16 * V_p / 3 / sqrt(3) *\n F_z *\n c3 *\n s1 *\n (-sin(φ_c) * cos(φ) + cos(φ_c) * sin(φ) * cos(λ - λ_c)) /\n sin(d / _a)\n v′::FT =\n 16 * V_p / 3 / sqrt(3) * F_z * c3 * s1 * cos(φ_c) * sin(λ - λ_c) /\n sin(d / _a)\n else\n u′ = FT(0)\n v′ = FT(0)\n end\n w′::FT = 0\n u_sphere = SVector{3, FT}(u_ref + u′, v_ref + v′, w_ref + w′)\n #u_sphere = SVector{3, FT}(u_ref, v_ref, w_ref)\n u_cart = sphr_to_cart_vec(u_sphere, φ, λ)\n\n ## temperature & density\n T::FT = T_v\n ρ::FT = p / (_R_d * T)\n ## potential & kinetic energy\n e_pot = aux.Φ\n e_kin::FT = 0.5 * u_cart' * u_cart\n e_int = _cv_d * T\n\n ## Assign state variables\n state.ρ = ρ\n state.ρu = ρ * u_cart\n if total_energy\n state.ρe = ρ * (e_int + e_kin + e_pot)\n else\n state.ρe = ρ * (e_int + e_kin)\n end\n\n nothing\nend\n\n\nfunction main()\n ClimateMachine.init(parse_clargs = true)\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 3\n numelem_horz = 8\n numelem_vert = 5\n\n timeend = 10 * 24 * 3600\n outputtime = 24 * 3600\n\n FT = Float64\n result = run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n )\nend\n\nfunction run(\n mpicomm,\n polynomialorder,\n numelem_horz,\n numelem_vert,\n timeend,\n outputtime,\n ArrayType,\n FT,\n)\n _planet_radius::FT = planet_radius(param_set)\n domain_height = FT(30e3)\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + domain_height),\n nelem = numelem_vert,\n )\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n meshwarp = cubedshellwarp,\n )\n\n T_profile =\n DecayingTemperatureProfile{FT}(param_set, FT(290), FT(220), FT(8e3))\n\n\n if total_energy\n sources = (Coriolis(),)\n else\n sources = (Coriolis(), Gravity())\n end\n problem = BaroclinicWave()\n model = DryAtmosModel{FT}(\n SphericalOrientation(),\n problem,\n ref_state = DryReferenceState(T_profile),\n sources = sources,\n )\n\n esdg = ESDGModel(\n model,\n grid,\n #volume_numerical_flux_first_order = CentralVolumeFlux(),\n #volume_numerical_flux_first_order = EntropyConservative(),\n volume_numerical_flux_first_order = KGVolumeFlux(),\n #surface_numerical_flux_first_order = MatrixFlux(),\n surface_numerical_flux_first_order = RusanovNumericalFlux(),\n )\n\n linearmodel = DryAtmosAcousticGravityLinearModel(model)\n lineardg = DGModel(\n linearmodel,\n grid,\n RusanovNumericalFlux(),\n #CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = esdg.state_auxiliary,\n )\n\n # determine the time step\n element_size = (domain_height / numelem_vert)\n acoustic_speed = soundspeed_air(param_set, FT(330))\n\n dx = min_node_distance(grid)\n cfl = 3\n dt = cfl * dx / acoustic_speed\n\n Q = init_ode_state(esdg, FT(0))\n\n #odesolver = LSRK144NiegemannDiehlBusch(esdg, Q; dt = dt, t0 = 0)\n\n linearsolver = ManyColumnLU()\n odesolver = ARK2GiraldoKellyConstantinescu(\n esdg,\n lineardg,\n LinearBackwardEulerSolver(linearsolver; isadjustable = false),\n Q;\n dt = dt,\n t0 = 0,\n split_explicit_implicit = false,\n )\n\n eng0 = norm(Q)\n @info @sprintf \"\"\"Starting\n ArrayType = %s\n FT = %s\n polynomialorder = %d\n numelem_horz = %d\n numelem_vert = %d\n dt = %.16e\n norm(Q₀) = %.16e\n \"\"\" \"$ArrayType\" \"$FT\" polynomialorder numelem_horz numelem_vert dt eng0\n\n # Set up the information callback\n starttime = Ref(now())\n cbinfo = EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q)\n runtime = Dates.format(\n convert(DateTime, now() - starttime[]),\n dateformat\"HH:MM:SS\",\n )\n @info @sprintf \"\"\"Update\n simtime = %.16e\n runtime = %s\n norm(Q) = %.16e\n \"\"\" gettime(odesolver) runtime energy\n end\n end\n cbcfl = EveryXSimulationSteps(100) do\n simtime = gettime(odesolver)\n\n @views begin\n ρ = Array(Q.data[:, 1, :])\n ρu = Array(Q.data[:, 2, :])\n ρv = Array(Q.data[:, 3, :])\n ρw = Array(Q.data[:, 4, :])\n end\n\n u = ρu ./ ρ\n v = ρv ./ ρ\n w = ρw ./ ρ\n\n ue = extrema(u)\n ve = extrema(v)\n we = extrema(w)\n\n @info @sprintf \"\"\"CFL\n simtime = %.16e\n u = (%.4e, %.4e)\n v = (%.4e, %.4e)\n w = (%.4e, %.4e)\n \"\"\" simtime ue... ve... we...\n end\n callbacks = (cbinfo, cbcfl)\n\n\n #filterorder = 32\n #filter = ExponentialFilter(grid, 0, filterorder)\n #cbfilter = EveryXSimulationSteps(1) do\n # Filters.apply!(\n # Q,\n # #AtmosFilterPerturbations(model),\n # :,\n # grid,\n # filter,\n # # state_auxiliary = esdg.state_auxiliary,\n # )\n # nothing\n #end\n #callbacks = (callbacks..., cbfilter)\n\n if output_vtk\n # create vtk dir\n vtkdir =\n \"vtk_esdg_total_KG_ncg_hires_baroclinic\" *\n \"_poly$(polynomialorder)_horz$(numelem_horz)_vert$(numelem_vert)\" *\n \"_$(ArrayType)_$(FT)\"\n mkpath(vtkdir)\n\n vtkstep = 0\n # output initial step\n do_output(mpicomm, vtkdir, vtkstep, esdg, Q, model)\n\n # setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n do_output(mpicomm, vtkdir, vtkstep, esdg, Q, model)\n end\n callbacks = (callbacks..., cbvtk)\n end\n\n ## Create a callback to report state statistics for main MPIStateArrays\n ## every ntFreq timesteps.\n nt_freq = floor(Int, 1 // 10 * timeend / dt)\n cbsc =\n ClimateMachine.StateCheck.sccreate([(Q, \"state\")], nt_freq; prec = 12)\n callbacks = (callbacks..., cbsc)\n\n solve!(\n Q,\n odesolver;\n timeend = timeend,\n adjustfinalstep = false,\n callbacks = callbacks,\n )\n\n # final statistics\n engf = norm(Q)\n @info @sprintf \"\"\"Finished\n norm(Q) = %.16e\n norm(Q) / norm(Q₀) = %.16e\n norm(Q) - norm(Q₀) = %.16e\n \"\"\" engf engf / eng0 engf - eng0\n engf\n\n ## Check results against reference if present\n ClimateMachine.StateCheck.scprintref(cbsc)\n #! format: off\n refDat = (\n [\n [ \"state\", \"ρ\", 1.26085994139264381125e-02, 1.51502814805562069367e+00, 3.46330071392505767225e-01, 3.42221017037761976454e-01 ],\n [ \"state\", \"ρu[1]\", -1.26747868050267172180e+02, 1.22735648852872344605e+02, -6.56484249582622303443e-02, 1.07588365672914481053e+01 ],\n [ \"state\", \"ρu[2]\", -1.44635478251794808102e+02, 1.11383888659731638882e+02, -1.53109264073002888581e-03, 1.06376480052955262323e+01 ],\n [ \"state\", \"ρu[3]\", -1.50479775987282266669e+02, 1.62284843398145170568e+02, 5.84596035289077428643e-02, 9.73728925076320095400e+00 ],\n [ \"state\", \"ρe\", 1.73490214503025617887e+03, 2.70352534694924892392e+05, 6.42401799036320589948e+04, 7.10054852130306971958e+04 ],\n ],\n [\n [ \"state\", \"ρ\", 12, 12, 12, 12 ],\n [ \"state\", \"ρu[1]\", 12, 12, 12, 12 ],\n [ \"state\", \"ρu[2]\", 12, 12, 12, 12 ],\n [ \"state\", \"ρu[3]\", 12, 12, 12, 12 ],\n [ \"state\", \"ρe\", 12, 12, 12, 12 ],\n ],\n )\n #! format: on\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(cbsc, refDat)\n end\nend\n\n\nfunction do_output(\n mpicomm,\n vtkdir,\n vtkstep,\n dg,\n Q,\n model,\n testname = \"baroclinicwave\",\n)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(\n filename,\n Q,\n dg,\n statenames,\n dg.state_auxiliary,\n auxnames;\n number_sample_points = 10,\n )\n\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(pvtuprefix, prefixes, (statenames..., auxnames...), eltype(Q))\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend\n\n@testset \"$(@__FILE__)\" begin\n tic = Base.time()\n\n main()\n\n toc = Base.time()\n time = toc - tic\n println(time)\nend\n" }, { "path": "test/Numerics/ESDGMethods/DryAtmos/linear.jl", "content": "using ClimateMachine.DGMethods: DGModel\nusing ClimateMachine.MPIStateArrays: MPIStateArray\nusing ClimateMachine.DGMethods.NumericalFluxes: NumericalFluxSecondOrder\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.Mesh.Grids: Direction\n\nimport ClimateMachine.BalanceLaws:\n flux_second_order!,\n indefinite_stack_integral!,\n reverse_indefinite_stack_integral!,\n integral_load_auxiliary_state!,\n integral_set_auxiliary_state!,\n reverse_integral_load_auxiliary_state!,\n reverse_integral_set_auxiliary_state!\n\n@inline function linearized_pressure(ρ, ρe, Φ)\n FT = eltype(ρ)\n γ = FT(gamma(param_set))\n if total_energy\n (γ - 1) * (ρe - ρ * Φ)\n else\n (γ - 1) * ρe\n end\nend\n\nabstract type DryAtmosLinearModel <: BalanceLaw end\n\nfunction vars_state(lm::DryAtmosLinearModel, ::Prognostic, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n end\nend\nvars_state(lm::DryAtmosLinearModel, st::Auxiliary, FT) =\n vars_state(lm.atmos, st, FT)\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n lm::DryAtmosLinearModel,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n return false\nend\nfunction flux_second_order!(\n lm::DryAtmosLinearModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n)\n nothing\nend\nintegral_load_auxiliary_state!(\n lm::DryAtmosLinearModel,\n integ::Vars,\n state::Vars,\n aux::Vars,\n) = nothing\nintegral_set_auxiliary_state!(lm::DryAtmosLinearModel, aux::Vars, integ::Vars) =\n nothing\nreverse_integral_load_auxiliary_state!(\n lm::DryAtmosLinearModel,\n integ::Vars,\n state::Vars,\n aux::Vars,\n) = nothing\nreverse_integral_set_auxiliary_state!(\n lm::DryAtmosLinearModel,\n aux::Vars,\n integ::Vars,\n) = nothing\nflux_second_order!(\n lm::DryAtmosLinearModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n t::Real,\n) = nothing\nfunction wavespeed(\n lm::DryAtmosLinearModel,\n nM,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n ref = aux.ref_state\n return soundspeed(ref.ρ, ref.p)\nend\n\nboundary_conditions(lm::DryAtmosLinearModel) = (1, 2)\nfunction boundary_state!(\n nf::NumericalFluxFirstOrder,\n bc,\n lm::DryAtmosLinearModel,\n args...,\n)\n boundary_state!(nf, bc, lm.atmos, args...)\nend\nfunction boundary_state!(\n nf::NumericalFluxSecondOrder,\n bc,\n lm::DryAtmosLinearModel,\n args...,\n)\n nothing\nend\ninit_state_auxiliary!(lm::DryAtmosLinearModel, aux::Vars, geom::LocalGeometry) =\n nothing\ninit_state_prognostic!(\n lm::DryAtmosLinearModel,\n state::Vars,\n aux::Vars,\n coords,\n t,\n) = nothing\n\nstruct DryAtmosAcousticGravityLinearModel{M} <: DryAtmosLinearModel\n atmos::M\n function DryAtmosAcousticGravityLinearModel(atmos::M) where {M}\n if atmos.ref_state === NoReferenceState()\n error(\"DryAtmosAcousticGravityLinearModel needs a model with a reference state\")\n end\n new{M}(atmos)\n end\nend\nfunction flux_first_order!(\n lm::DryAtmosAcousticGravityLinearModel,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n direction,\n)\n FT = eltype(state)\n ref = aux.ref_state\n\n flux.ρ = state.ρu\n pL = linearized_pressure(state.ρ, state.ρe, aux.Φ)\n flux.ρu += pL * I\n flux.ρe = ((ref.ρe + ref.p) / ref.ρ) * state.ρu\n nothing\nend\nfunction source!(\n lm::DryAtmosAcousticGravityLinearModel,\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n t::Real,\n ::NTuple{1, Dir},\n) where {Dir <: Direction}\n if Dir === VerticalDirection || Dir === EveryDirection\n ∇Φ = aux.∇Φ\n source.ρu -= state.ρ * ∇Φ\n if !total_energy\n source.ρe -= state.ρu' * ∇Φ\n end\n end\n nothing\nend\n" }, { "path": "test/Numerics/ESDGMethods/DryAtmos/run_tests.jl", "content": "using Test\nusing ClimateMachine\nimport ClimateMachine.BalanceLaws\nimport ClimateMachine.BalanceLaws: boundary_state!\nusing ClimateMachine.DGMethods: ESDGModel, init_ode_state\nusing ClimateMachine.DGMethods.NumericalFluxes:\n numerical_volume_flux_first_order!\nusing ClimateMachine.Mesh.Topologies: BrickTopology\nusing ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid\nusing ClimateMachine.VariableTemplates: varsindex\nusing StaticArrays: MArray, @SVector\nusing KernelAbstractions: wait\nusing Random\nusing DoubleFloats\nusing MPI\nusing ClimateMachine.MPIStateArrays\n\nRandom.seed!(7)\n\ninclude(\"DryAtmos.jl\")\n\nboundary_state!(::Nothing, _...) = nothing\n\nstruct TestProblem <: AbstractDryAtmosProblem end\n\n# Random initialization function\nfunction init_state_prognostic!(\n ::DryAtmosModel,\n ::TestProblem,\n state_prognostic,\n state_auxiliary,\n _...,\n) where {dim}\n FT = eltype(state_prognostic)\n ρ = state_prognostic.ρ = rand(FT) + 1\n ρu = state_prognostic.ρu = 2 * (@SVector rand(FT, 3)) .- 1\n p = rand(FT) + 1\n Φ = state_auxiliary.Φ\n state_prognostic.ρe = totalenergy(ρ, ρu, p, Φ)\nend\n\nfunction check_operators(FT, dim, mpicomm, N, ArrayType)\n # Create a warped mesh so the metrics are not constant\n Ne = (8, 9, 10)\n brickrange = (\n range(FT(-1); length = Ne[1] + 1, stop = 1),\n range(FT(-1); length = Ne[2] + 1, stop = 1),\n range(FT(-1); length = Ne[3] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n ntuple(k -> brickrange[k], dim);\n periodicity = ntuple(k -> true, dim),\n )\n warpfun =\n (x1, x2, x3) -> begin\n α = (4 / π) * (1 - x1^2) * (1 - x2^2) * (1 - x3^2)\n # Rotate by α with x1 and x2\n x1, x2 = cos(α) * x1 - sin(α) * x2, sin(α) * x1 + cos(α) * x2\n # Rotate by α with x1 and x3\n if dim == 3\n x1, x3 = cos(α) * x1 - sin(α) * x3, sin(α) * x1 + cos(α) * x3\n end\n return (x1, x2, x3)\n end\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n # Orientation does not matter since we will be setting the geopotential to a\n # random field\n model = DryAtmosModel{dim}(FlatOrientation(), TestProblem())\n\n ##################################################################\n # check that the volume terms lead to only surface contributions #\n ##################################################################\n # Create the ES model\n esdg = ESDGModel(\n model,\n grid;\n volume_numerical_flux_first_order = EntropyConservative(),\n surface_numerical_flux_first_order = nothing,\n )\n\n # Make the Geopotential random\n esdg.state_auxiliary .= ArrayType(2rand(FT, size(esdg.state_auxiliary)))\n start_exchange = MPIStateArrays.begin_ghost_exchange!(esdg.state_auxiliary)\n end_exchange = MPIStateArrays.end_ghost_exchange!(\n esdg.state_auxiliary,\n dependencies = start_exchange,\n )\n wait(end_exchange)\n\n # Create a random state\n state_prognostic = init_ode_state(esdg; init_on_cpu = true)\n\n # Storage for the tendency\n volume_tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(volume_tendency, state_prognostic, nothing, 0)\n\n # Check that the volume terms only lead to surface integrals of\n # ∑_{j} n_j ψ_j\n # where Ψ_j = β^T f_j - ζ_j = ρu_j\n Np, K = (N + 1)^dim, length(esdg.grid.topology.realelems)\n\n # Get the mass matrix on the host\n _M = ClimateMachine.Grids._M\n M = Array(grid.vgeo[:, _M:_M, 1:K])\n\n # Get the state, tendency, and aux on the host\n Q = Array(state_prognostic.data[:, :, 1:K])\n dQ = Array(volume_tendency.data[:, :, 1:K])\n A = Array(esdg.state_auxiliary.data[:, :, 1:K])\n\n # Compute the entropy variables\n β = similar(Q, Np, number_states(model, Entropy()), K)\n @views for e in 1:K\n for i in 1:Np\n state_to_entropy_variables!(\n model,\n β[i, :, e],\n Q[i, :, e],\n A[i, :, e],\n )\n end\n end\n\n # Get the unit normals and surface mass matrix\n sgeo = Array(grid.sgeo)\n n1 = sgeo[ClimateMachine.Grids._n1, :, :, 1:K]\n n2 = sgeo[ClimateMachine.Grids._n2, :, :, 1:K]\n n3 = sgeo[ClimateMachine.Grids._n3, :, :, 1:K]\n sM = sgeo[ClimateMachine.Grids._sM, :, :, 1:K]\n\n # Get the Ψs\n fmask = Array(grid.vmap⁻[:, :, 1])\n _ρu = varsindex(vars_state(model, Prognostic(), FT), :ρu)\n Ψ1 = Q[fmask, _ρu[1], 1:K]\n Ψ2 = Q[fmask, _ρu[2], 1:K]\n Ψ3 = Q[fmask, _ρu[3], 1:K]\n\n # Compute the surface integral:\n # ∫_Ωf ∑_j n_j * Ψ_j\n surface = sum(sM .* (n1 .* Ψ1 + n2 .* Ψ2 + n3 .* Ψ3), dims = (1, 2))[:]\n\n # Compute the volume integral:\n # -∫_Ω ∑_j β^T (dq/dt)\n # (tendency is -dq / dt)\n num_state = number_states(model, Prognostic())\n volume = sum(β[:, 1:num_state, :] .* M .* dQ, dims = (1, 2))[:]\n\n @test all(isapprox.(\n surface,\n volume;\n atol = 10eps(FT),\n rtol = sqrt(eps(FT)),\n ))\n\n ###########################################\n # check that the volume and surface match #\n ###########################################\n esdg = ESDGModel(\n model,\n grid;\n state_auxiliary = esdg.state_auxiliary,\n volume_numerical_flux_first_order = nothing,\n surface_numerical_flux_first_order = EntropyConservative(),\n )\n\n surface_tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(surface_tendency, state_prognostic, nothing, 0)\n\n # Surface integral should be equal and opposite to the volume integral\n dQ = Array(surface_tendency.data[:, :, 1:K])\n volume_integral = MPI.Allreduce(sum(volume), +, mpicomm)\n surface_integral =\n MPI.Allreduce(sum(β[:, 1:num_state, :] .* M .* dQ), +, mpicomm)\n @test volume_integral ≈ -surface_integral\n\n ########################################################\n # check that the full tendency is entropy conservative #\n ########################################################\n esdg = ESDGModel(\n model,\n grid;\n state_auxiliary = esdg.state_auxiliary,\n volume_numerical_flux_first_order = EntropyConservative(),\n surface_numerical_flux_first_order = EntropyConservative(),\n )\n\n tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(tendency, state_prognostic, nothing, 0)\n\n # Check for entropy conservation\n dQ = Array(tendency.data[:, :, 1:K])\n integral = MPI.Allreduce(sum(β[:, 1:num_state, :] .* M .* dQ), +, mpicomm)\n @test isapprox(integral, 0, atol = sqrt(eps(sum(volume))))\nend\n\nlet\n model = DryAtmosModel{3}(FlatOrientation(), TestProblem())\n num_state = number_states(model, Prognostic())\n num_aux = number_states(model, Auxiliary())\n num_entropy = number_states(model, Entropy())\n\n @testset \"state to entropy variable transforms\" begin\n for FT in (Float32, Float64)\n state_in =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_in = rand(FT, num_aux)\n state_out = similar(state_in)\n aux_out = similar(aux_in)\n entropy = similar(state_in, num_entropy)\n\n state_to_entropy_variables!(model, entropy, state_in, aux_in)\n entropy_variables_to_state!(model, state_out, aux_out, entropy)\n\n @test all(state_in .≈ state_out)\n #@test all(aux_in .≈ aux_out)\n end\n end\n\n @testset \"test numerical flux for Tadmor shuffle\" begin\n for FT in (Float32, Float64)\n # Create some random states\n state_1 =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_1 = 0 * rand(FT, num_aux)\n\n state_2 =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_2 = 0 * rand(FT, num_aux)\n\n # Get the entropy variables for the two states\n entropy_1 = similar(state_1, num_entropy)\n state_to_entropy_variables!(model, entropy_1, state_1, aux_1)\n\n entropy_2 = similar(state_1, num_entropy)\n state_to_entropy_variables!(model, entropy_2, state_2, aux_2)\n\n # Get the values of Ψ_j = β^T f_j - ζ_j = ρu_j where β is the\n # entropy variables, f_j is the conservative flux, and ζ_j is the\n # entropy flux. For conservation laws this is the entropy potential.\n Ψ_1 = Vars{vars_state(model, Prognostic(), FT)}(state_1).ρu\n Ψ_2 = Vars{vars_state(model, Prognostic(), FT)}(state_2).ρu\n\n # Evaluate the flux with both orders of the two states\n H_12 = fill!(MArray{Tuple{3, num_state}, FT}(undef), -zero(FT))\n numerical_volume_flux_first_order!(\n EntropyConservative(),\n model,\n H_12,\n state_1,\n aux_1,\n state_2,\n aux_2,\n )\n\n H_21 = fill!(MArray{Tuple{3, num_state}, FT}(undef), -zero(FT))\n numerical_volume_flux_first_order!(\n EntropyConservative(),\n model,\n H_21,\n state_2,\n aux_2,\n state_1,\n aux_1,\n )\n\n # Check that we satisfy the Tadmor shuffle\n @test all(\n H_12 * entropy_1[1:num_state] - H_21 * entropy_2[1:num_state] .≈\n Ψ_1 - Ψ_2,\n )\n end\n end\n\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n polynomialorder = 4\n test_types = (Float32, Float64)\n for FT in test_types\n for dim in 2:3\n @testset \"check ESDGMethods relations for dim = $dim and FT = $FT\" begin\n check_operators(FT, dim, mpicomm, polynomialorder, ArrayType)\n end\n end\n end\nend\n" }, { "path": "test/Numerics/ESDGMethods/DryAtmos/run_tests_mpo.jl", "content": "using Test\nusing ClimateMachine\nimport ClimateMachine.BalanceLaws\nimport ClimateMachine.BalanceLaws: boundary_state!\nusing ClimateMachine.DGMethods: ESDGModel, init_ode_state\nusing ClimateMachine.DGMethods.NumericalFluxes:\n numerical_volume_flux_first_order!\nusing ClimateMachine.Mesh.Topologies: BrickTopology\nusing ClimateMachine.Mesh.Grids: DiscontinuousSpectralElementGrid\nusing ClimateMachine.VariableTemplates: varsindex\nusing StaticArrays: MArray, @SVector\nusing KernelAbstractions: wait\nusing Random\nusing DoubleFloats\nusing MPI\nusing ClimateMachine.MPIStateArrays\n\nRandom.seed!(7)\n\ninclude(\"DryAtmos.jl\")\n\nboundary_state!(::Nothing, _...) = nothing\n\nstruct TestProblem <: AbstractDryAtmosProblem end\n\n# Random initialization function\nfunction init_state_prognostic!(\n ::DryAtmosModel,\n ::TestProblem,\n state_prognostic,\n state_auxiliary,\n _...,\n) where {dim}\n FT = eltype(state_prognostic)\n ρ = state_prognostic.ρ = rand(FT) + 1\n ρu = state_prognostic.ρu = 2 * (@SVector rand(FT, 3)) .- 1\n p = rand(FT) + 1\n Φ = state_auxiliary.Φ\n state_prognostic.ρe = totalenergy(ρ, ρu, p, Φ)\nend\n\nfunction check_operators(FT, dim, mpicomm, N, ArrayType)\n # Create a warped mesh so the metrics are not constant\n Ne = (8, 9, 10)\n brickrange = (\n range(FT(-1); length = Ne[1] + 1, stop = 1),\n range(FT(-1); length = Ne[2] + 1, stop = 1),\n range(FT(-1); length = Ne[3] + 1, stop = 1),\n )\n topl = BrickTopology(\n mpicomm,\n ntuple(k -> brickrange[k], dim);\n periodicity = ntuple(k -> true, dim),\n )\n warpfun =\n (x1, x2, x3) -> begin\n α = (4 / π) * (1 - x1^2) * (1 - x2^2) * (1 - x3^2)\n # Rotate by α with x1 and x2\n x1, x2 = cos(α) * x1 - sin(α) * x2, sin(α) * x1 + cos(α) * x2\n # Rotate by α with x1 and x3\n if dim == 3\n x1, x3 = cos(α) * x1 - sin(α) * x3, sin(α) * x1 + cos(α) * x3\n end\n return (x1, x2, x3)\n end\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n # Orientation does not matter since we will be setting the geopotential to a\n # random field\n model = DryAtmosModel{dim}(FlatOrientation(), TestProblem())\n\n ##################################################################\n # check that the volume terms lead to only surface contributions #\n ##################################################################\n # Create the ES model\n esdg = ESDGModel(\n model,\n grid;\n volume_numerical_flux_first_order = EntropyConservative(),\n surface_numerical_flux_first_order = nothing,\n )\n\n # Make the Geopotential random\n esdg.state_auxiliary .= ArrayType(2rand(FT, size(esdg.state_auxiliary)))\n start_exchange = MPIStateArrays.begin_ghost_exchange!(esdg.state_auxiliary)\n end_exchange = MPIStateArrays.end_ghost_exchange!(\n esdg.state_auxiliary,\n dependencies = start_exchange,\n )\n wait(end_exchange)\n\n # Create a random state\n state_prognostic = init_ode_state(esdg; init_on_cpu = true)\n\n # Storage for the tendency\n volume_tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(volume_tendency, state_prognostic, nothing, 0)\n\n # Check that the volume terms only lead to surface integrals of\n # ∑_{j} n_j ψ_j\n # where Ψ_j = β^T f_j - ζ_j = ρu_j\n Np, K = prod(N .+ 1), length(esdg.grid.topology.realelems)\n\n # Get the mass matrix on the host\n _M = ClimateMachine.Grids._M\n M = Array(grid.vgeo[:, _M:_M, 1:K])\n\n # Get the state, tendency, and aux on the host\n Q = Array(state_prognostic.data[:, :, 1:K])\n dQ = Array(volume_tendency.data[:, :, 1:K])\n A = Array(esdg.state_auxiliary.data[:, :, 1:K])\n\n # Compute the entropy variables\n β = similar(Q, Np, number_states(model, Entropy()), K)\n @views for e in 1:K\n for i in 1:Np\n state_to_entropy_variables!(\n model,\n β[i, :, e],\n Q[i, :, e],\n A[i, :, e],\n )\n end\n end\n\n # Get the unit normals and surface mass matrix\n sgeo = Array(grid.sgeo)\n n1 = sgeo[ClimateMachine.Grids._n1, :, :, 1:K]\n n2 = sgeo[ClimateMachine.Grids._n2, :, :, 1:K]\n n3 = sgeo[ClimateMachine.Grids._n3, :, :, 1:K]\n sM = sgeo[ClimateMachine.Grids._sM, :, :, 1:K]\n\n num_state = number_states(model, Prognostic())\n volume = sum(β[:, 1:num_state, :] .* M .* dQ, dims = (1, 2))[:]\n\n ###########################################\n # check that the volume and surface match #\n ###########################################\n esdg = ESDGModel(\n model,\n grid;\n state_auxiliary = esdg.state_auxiliary,\n volume_numerical_flux_first_order = nothing,\n surface_numerical_flux_first_order = EntropyConservative(),\n )\n\n surface_tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(surface_tendency, state_prognostic, nothing, 0)\n\n # Surface integral should be equal and opposite to the volume integral\n dQ = Array(surface_tendency.data[:, :, 1:K])\n volume_integral = MPI.Allreduce(sum(volume), +, mpicomm)\n surface_integral =\n MPI.Allreduce(sum(β[:, 1:num_state, :] .* M .* dQ), +, mpicomm)\n @test volume_integral ≈ -surface_integral\n\n ########################################################\n # check that the full tendency is entropy conservative #\n ########################################################\n esdg = ESDGModel(\n model,\n grid;\n state_auxiliary = esdg.state_auxiliary,\n volume_numerical_flux_first_order = EntropyConservative(),\n surface_numerical_flux_first_order = EntropyConservative(),\n )\n\n tendency = similar(state_prognostic)\n\n # Compute the tendency function\n esdg(tendency, state_prognostic, nothing, 0)\n\n # Check for entropy conservation\n dQ = Array(tendency.data[:, :, 1:K])\n integral = MPI.Allreduce(sum(β[:, 1:num_state, :] .* M .* dQ), +, mpicomm)\n @test isapprox(integral, 0, atol = sqrt(eps(sum(volume))))\nend\n\nlet\n model = DryAtmosModel{3}(FlatOrientation(), TestProblem())\n num_state = number_states(model, Prognostic())\n num_aux = number_states(model, Auxiliary())\n num_entropy = number_states(model, Entropy())\n\n @testset \"state to entropy variable transforms\" begin\n for FT in (Float32, Float64)\n state_in =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_in = rand(FT, num_aux)\n state_out = similar(state_in)\n aux_out = similar(aux_in)\n entropy = similar(state_in, num_entropy)\n\n state_to_entropy_variables!(model, entropy, state_in, aux_in)\n entropy_variables_to_state!(model, state_out, aux_out, entropy)\n\n @test all(state_in .≈ state_out)\n #@test all(aux_in .≈ aux_out)\n end\n end\n\n @testset \"test numerical flux for Tadmor shuffle\" begin\n for FT in (Float32, Float64)\n # Create some random states\n state_1 =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_1 = 0 * rand(FT, num_aux)\n\n state_2 =\n [1, 2, 2, 2, 1] .* rand(FT, num_state) + [3, -1, -1, -1, 100]\n aux_2 = 0 * rand(FT, num_aux)\n\n # Get the entropy variables for the two states\n entropy_1 = similar(state_1, num_entropy)\n state_to_entropy_variables!(model, entropy_1, state_1, aux_1)\n\n entropy_2 = similar(state_1, num_entropy)\n state_to_entropy_variables!(model, entropy_2, state_2, aux_2)\n\n # Get the values of Ψ_j = β^T f_j - ζ_j = ρu_j where β is the\n # entropy variables, f_j is the conservative flux, and ζ_j is the\n # entropy flux. For conservation laws this is the entropy potential.\n Ψ_1 = Vars{vars_state(model, Prognostic(), FT)}(state_1).ρu\n Ψ_2 = Vars{vars_state(model, Prognostic(), FT)}(state_2).ρu\n\n # Evaluate the flux with both orders of the two states\n H_12 = fill!(MArray{Tuple{3, num_state}, FT}(undef), -zero(FT))\n numerical_volume_flux_first_order!(\n EntropyConservative(),\n model,\n H_12,\n state_1,\n aux_1,\n state_2,\n aux_2,\n )\n\n H_21 = fill!(MArray{Tuple{3, num_state}, FT}(undef), -zero(FT))\n numerical_volume_flux_first_order!(\n EntropyConservative(),\n model,\n H_21,\n state_2,\n aux_2,\n state_1,\n aux_1,\n )\n\n # Check that we satisfy the Tadmor shuffle\n @test all(\n H_12 * entropy_1[1:num_state] - H_21 * entropy_2[1:num_state] .≈\n Ψ_1 - Ψ_2,\n )\n end\n end\n\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n polynomialorders = ((2, 2, 1), (1, 1, 2))\n test_types = (Float32, Float64)\n for FT in test_types\n for dim in (3,)\n for polynomialorder in polynomialorders\n dim = 3\n @testset \"check ESDGMethods relations for dim = $dim, FT = $FT, and polynomial order $polynomialorder\" begin\n check_operators(\n FT,\n dim,\n mpicomm,\n polynomialorder,\n ArrayType,\n )\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/ESDGMethods/diagnostics.jl", "content": "using KernelAbstractions\nusing ClimateMachine.MPIStateArrays: array_device, weightedsum\nusing KernelAbstractions.Extras: @unroll\n\nfunction entropy_integral(dg, entropy, state_prognostic)\n balance_law = dg.balance_law\n state_auxiliary = dg.state_auxiliary\n device = array_device(state_prognostic)\n grid = dg.grid\n topology = grid.topology\n Np = dofs_per_element(grid)\n dim = dimensionality(grid)\n # XXX: Needs updating for multiple polynomial orders\n N = polynomialorders(grid)\n # Currently only support single polynomial order\n @assert all(N[1] .== N)\n N = N[1]\n\n realelems = topology.realelems\n\n event = Event(device)\n event = esdg_compute_entropy!(device, min(Np, 1024))(\n balance_law,\n Val(dim),\n Val(N),\n entropy.data,\n state_prognostic.data,\n state_auxiliary.data,\n realelems,\n ndrange = Np * length(realelems),\n dependencies = event,\n )\n wait(event)\n\n weightedsum(entropy)\nend\n\n@kernel function esdg_compute_entropy!(\n balance_law::BalanceLaw,\n ::Val{dim},\n ::Val{N},\n entropy,\n state_prognostic,\n state_auxiliary,\n elems,\n) where {dim, N}\n\n FT = eltype(state_prognostic)\n num_state_prognostic = number_states(balance_law, Prognostic())\n num_state_auxiliary = number_states(balance_law, Auxiliary())\n\n Nq = N + 1\n\n Nqk = dim == 2 ? 1 : Nq\n\n Np = Nq * Nq * Nqk\n\n local_state_prognostic = MArray{Tuple{num_state_prognostic}, FT}(undef)\n local_state_auxiliary = MArray{Tuple{num_state_auxiliary}, FT}(undef)\n\n I = @index(Global, Linear)\n eI = (I - 1) ÷ Np + 1\n n = (I - 1) % Np + 1\n\n @inbounds begin\n e = elems[eI]\n @unroll for s in 1:num_state_prognostic\n local_state_prognostic[s] = state_prognostic[n, s, e]\n end\n\n @unroll for s in 1:num_state_auxiliary\n local_state_auxiliary[s] = state_auxiliary[n, s, e]\n end\n\n entropy[n, 1, e] = state_to_entropy(\n balance_law,\n Vars{vars_state(balance_law, Prognostic(), FT)}(\n local_state_prognostic,\n ),\n Vars{vars_state(balance_law, Auxiliary(), FT)}(\n local_state_auxiliary,\n ),\n )\n end\nend\n\nfunction entropy_product(dg, entropy, state_prognostic, tendency)\n balance_law = dg.balance_law\n state_auxiliary = dg.state_auxiliary\n device = array_device(state_prognostic)\n grid = dg.grid\n topology = grid.topology\n Np = dofs_per_element(grid)\n dim = dimensionality(grid)\n # XXX: Needs updating for multiple polynomial orders\n N = polynomialorders(grid)\n # Currently only support single polynomial order\n @assert all(N[1] .== N)\n N = N[1]\n\n realelems = topology.realelems\n\n event = Event(device)\n event = esdg_compute_entropy_product!(device, min(Np, 1024))(\n balance_law,\n Val(dim),\n Val(N),\n entropy.data,\n state_prognostic.data,\n tendency.data,\n state_auxiliary.data,\n realelems,\n ndrange = Np * length(realelems),\n dependencies = event,\n )\n wait(event)\n\n weightedsum(entropy)\nend\n\n@kernel function esdg_compute_entropy_product!(\n balance_law::BalanceLaw,\n ::Val{dim},\n ::Val{N},\n entropy,\n state_prognostic,\n tendency,\n state_auxiliary,\n elems,\n) where {dim, N}\n\n FT = eltype(state_prognostic)\n num_state_prognostic = number_states(balance_law, Prognostic())\n num_state_entropy = number_states(balance_law, Entropy())\n num_state_auxiliary = number_states(balance_law, Auxiliary())\n\n Nq = N + 1\n\n Nqk = dim == 2 ? 1 : Nq\n\n Np = Nq * Nq * Nqk\n\n local_state_entropy = MArray{Tuple{num_state_entropy}, FT}(undef)\n local_state_prognostic = MArray{Tuple{num_state_prognostic}, FT}(undef)\n local_tendency = MArray{Tuple{num_state_prognostic}, FT}(undef)\n local_state_auxiliary = MArray{Tuple{num_state_auxiliary}, FT}(undef)\n\n I = @index(Global, Linear)\n eI = (I - 1) ÷ Np + 1\n n = (I - 1) % Np + 1\n\n @inbounds begin\n e = elems[eI]\n\n @unroll for s in 1:num_state_prognostic\n local_state_prognostic[s] = state_prognostic[n, s, e]\n end\n\n @unroll for s in 1:num_state_prognostic\n local_tendency[s] = tendency[n, s, e]\n end\n\n @unroll for s in 1:num_state_auxiliary\n local_state_auxiliary[s] = state_auxiliary[n, s, e]\n end\n\n state_to_entropy_variables!(\n balance_law,\n Vars{vars_state(balance_law, Entropy(), FT)}(local_state_entropy),\n Vars{vars_state(balance_law, Prognostic(), FT)}(\n local_state_prognostic,\n ),\n Vars{vars_state(balance_law, Auxiliary(), FT)}(\n local_state_auxiliary,\n ),\n )\n\n local_product = -zero(FT)\n # not that tendency related to the last entropy variable is assumed zero\n @unroll for s in 1:num_state_prognostic\n local_product += local_state_entropy[s] * local_tendency[s]\n end\n entropy[n, 1, e] = local_product\n end\nend\n" }, { "path": "test/Numerics/Mesh/BrickMesh.jl", "content": "using ClimateMachine.Mesh.BrickMesh\nusing ClimateMachine.Mesh.Grids: mappings, commmapping\nusing Test\nusing MPI\n\nMPI.Initialized() || MPI.Init()\n\n@testset \"Linear Parition\" begin\n @test BrickMesh.linearpartition(1, 1, 1) == 1:1\n @test BrickMesh.linearpartition(20, 1, 1) == 1:20\n @test BrickMesh.linearpartition(10, 1, 2) == 1:5\n @test BrickMesh.linearpartition(10, 2, 2) == 6:10\nend\n\n@testset \"Hilbert Code\" begin\n @test BrickMesh.hilbertcode([0, 0], bits = 1) == [0, 0]\n @test BrickMesh.hilbertcode([0, 1], bits = 1) == [0, 1]\n @test BrickMesh.hilbertcode([1, 1], bits = 1) == [1, 0]\n @test BrickMesh.hilbertcode([1, 0], bits = 1) == [1, 1]\n @test BrickMesh.hilbertcode([0, 0], bits = 2) == [0, 0]\n @test BrickMesh.hilbertcode([1, 0], bits = 2) == [0, 1]\n @test BrickMesh.hilbertcode([1, 1], bits = 2) == [0, 2]\n @test BrickMesh.hilbertcode([0, 1], bits = 2) == [0, 3]\n @test BrickMesh.hilbertcode([0, 2], bits = 2) == [1, 0]\n @test BrickMesh.hilbertcode([0, 3], bits = 2) == [1, 1]\n @test BrickMesh.hilbertcode([1, 3], bits = 2) == [1, 2]\n @test BrickMesh.hilbertcode([1, 2], bits = 2) == [1, 3]\n @test BrickMesh.hilbertcode([2, 2], bits = 2) == [2, 0]\n @test BrickMesh.hilbertcode([2, 3], bits = 2) == [2, 1]\n @test BrickMesh.hilbertcode([3, 3], bits = 2) == [2, 2]\n @test BrickMesh.hilbertcode([3, 2], bits = 2) == [2, 3]\n @test BrickMesh.hilbertcode([3, 1], bits = 2) == [3, 0]\n @test BrickMesh.hilbertcode([2, 1], bits = 2) == [3, 1]\n @test BrickMesh.hilbertcode([2, 0], bits = 2) == [3, 2]\n @test BrickMesh.hilbertcode([3, 0], bits = 2) == [3, 3]\n\n @test BrickMesh.hilbertcode(UInt64.([14, 3, 4])) ==\n UInt64.([0x0, 0x0, 0xe25])\nend\n\n@testset \"Mesh to Hilbert Code\" begin\n let\n etc = Array{Float64}(undef, 2, 4, 6)\n etc[:, :, 1] = [2.0 3.0 2.0 3.0; 4.0 4.0 5.0 5.0]\n etc[:, :, 2] = [3.0 4.0 3.0 4.0; 4.0 4.0 5.0 5.0]\n etc[:, :, 3] = [4.0 5.0 4.0 5.0; 4.0 4.0 5.0 5.0]\n etc[:, :, 4] = [2.0 3.0 2.0 3.0; 5.0 5.0 6.0 6.0]\n etc[:, :, 5] = [3.0 4.0 3.0 4.0; 5.0 5.0 6.0 6.0]\n etc[:, :, 6] = [4.0 5.0 4.0 5.0; 5.0 5.0 6.0 6.0]\n\n code_exect = UInt64[\n 0x0000000000000000 0x1555555555555555 0xffffffffffffffff 0x5555555555555555 0x6aaaaaaaaaaaaaaa 0xaaaaaaaaaaaaaaaa\n 0x0000000000000000 0x5555555555555555 0xffffffffffffffff 0x5555555555555555 0xaaaaaaaaaaaaaaaa 0xaaaaaaaaaaaaaaaa\n ]\n\n code = centroidtocode(MPI.COMM_SELF, etc)\n\n @test code == code_exect\n end\n\n let\n nelem = 1\n d = 2\n\n etc = Array{Float64}(undef, d, d^2, nelem)\n etc[:, :, 1] = [2.0 3.0 2.0 3.0; 4.0 4.0 5.0 5.0]\n code = centroidtocode(MPI.COMM_SELF, etc)\n\n @test code == zeros(eltype(code), d, nelem)\n end\nend\n\n@testset \"Vertex Ordering\" begin\n @test ((1,), 1) == BrickMesh.vertsortandorder(1)\n\n @test ((1, 2), 1) == BrickMesh.vertsortandorder(1, 2)\n @test ((1, 2), 2) == BrickMesh.vertsortandorder(2, 1)\n\n @test ((1, 2, 3), 1) == BrickMesh.vertsortandorder(1, 2, 3)\n @test ((1, 2, 3), 2) == BrickMesh.vertsortandorder(3, 1, 2)\n @test ((1, 2, 3), 3) == BrickMesh.vertsortandorder(2, 3, 1)\n @test ((1, 2, 3), 4) == BrickMesh.vertsortandorder(2, 1, 3)\n @test ((1, 2, 3), 5) == BrickMesh.vertsortandorder(3, 2, 1)\n @test ((1, 2, 3), 6) == BrickMesh.vertsortandorder(1, 3, 2)\n\n @test_throws ErrorException BrickMesh.vertsortandorder(2, 1, 1)\n\n @test ((1, 2, 3, 4), 1) == BrickMesh.vertsortandorder(1, 2, 3, 4)\n @test ((1, 2, 3, 4), 2) == BrickMesh.vertsortandorder(1, 3, 2, 4)\n @test ((1, 2, 3, 4), 3) == BrickMesh.vertsortandorder(2, 1, 3, 4)\n @test ((1, 2, 3, 4), 4) == BrickMesh.vertsortandorder(2, 4, 1, 3)\n @test ((1, 2, 3, 4), 5) == BrickMesh.vertsortandorder(3, 1, 4, 2)\n @test ((1, 2, 3, 4), 6) == BrickMesh.vertsortandorder(3, 4, 1, 2)\n @test ((1, 2, 3, 4), 7) == BrickMesh.vertsortandorder(4, 2, 3, 1)\n @test ((1, 2, 3, 4), 8) == BrickMesh.vertsortandorder(4, 3, 2, 1)\n\n @test_throws ErrorException BrickMesh.vertsortandorder(1, 3, 3, 1)\nend\n\n@testset \"Mesh\" begin\n let\n (etv, etc, etb, fc) = brickmesh((4:7,), (false,))\n etv_expect = [\n 1 2 3\n 2 3 4\n ]\n etb_expect = [\n 1 0 0\n 0 0 1\n ]\n fc_expect = Array{Int64, 1}[]\n\n @test etv == etv_expect\n @test etb == etb_expect\n @test fc == fc_expect\n @test etc[:, :, 1] == [4 5]\n @test etc[:, :, 2] == [5 6]\n @test etc[:, :, 3] == [6 7]\n end\n\n let\n (etv, etc, etb, fc) = brickmesh((4:7,), (true,))\n etv_expect = [\n 1 2 3\n 2 3 4\n ]\n etb_expect = [\n 0 0 0\n 0 0 0\n ]\n fc_expect = Array{Int64, 1}[[3, 2, 1]]\n\n @test etv == etv_expect\n @test etb == etb_expect\n @test fc == fc_expect\n @test etc[:, :, 1] == [4 5]\n @test etc[:, :, 2] == [5 6]\n @test etc[:, :, 3] == [6 7]\n end\n\n let\n (etv, etc, etb, fc) = brickmesh((2:5, 4:6), (false, true))\n\n etv_expect = [\n 1 2 5 6\n 2 3 6 7\n 3 4 7 8\n 5 6 9 10\n 6 7 10 11\n 7 8 11 12\n ]'\n etb_expect = [\n 1 0 0 1 0 0\n 0 0 1 0 0 1\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n ]\n fc_expect = Array{Int64, 1}[[4, 4, 1, 2], [5, 4, 2, 3], [6, 4, 3, 4]]\n\n @test etv == etv_expect\n @test etb == etb_expect\n @test fc == fc_expect\n @test etc[:, :, 1] == [\n 2 3 2 3\n 4 4 5 5\n ]\n @test etc[:, :, 5] == [\n 3 4 3 4\n 5 5 6 6\n ]\n end\n\n let\n (etv, etc, etb, fc) =\n brickmesh((-1:2:1, -1:2:1, -1:1:1), (true, true, true))\n etv_expect = [\n 1 5\n 2 6\n 3 7\n 4 8\n 5 9\n 6 10\n 7 11\n 8 12\n ]\n etb_expect = zeros(Int64, 6, 2)\n\n fc_expect = Array{Int64, 1}[\n [1, 2, 1, 3, 5, 7],\n [1, 4, 1, 2, 5, 6],\n [2, 2, 5, 7, 9, 11],\n [2, 4, 5, 6, 9, 10],\n [2, 6, 1, 2, 3, 4],\n ]\n\n @test etv == etv_expect\n @test etb == etb_expect\n @test fc == fc_expect\n\n @test etc[:, :, 1] == [\n -1 1 -1 1 -1 1 -1 1\n -1 -1 1 1 -1 -1 1 1\n -1 -1 -1 -1 0 0 0 0\n ]\n\n @test etc[:, :, 2] == [\n -1 1 -1 1 -1 1 -1 1\n -1 -1 1 1 -1 -1 1 1\n 0 0 0 0 1 1 1 1\n ]\n end\n\n let\n (etv, etc, etb, fc) = brickmesh(\n (-1:1, -1:1, -1:1),\n (false, false, false),\n boundary = ((11, 12), (13, 14), (15, 16)),\n )\n\n @test etb == [\n 11 0 11 0 11 0 11 0\n 0 12 0 12 0 12 0 12\n 13 13 0 0 13 13 0 0\n 0 0 14 14 0 0 14 14\n 15 15 15 15 0 0 0 0\n 0 0 0 0 16 16 16 16\n ]\n end\n\n let\n x = (1:1000,)\n p = (false,)\n b = ((1, 2),)\n\n (etv, etc, etb, fc) = brickmesh(x, p, boundary = b)\n\n n = 50\n (etv_parts, etc_parts, etb_parts, fc_parts) =\n brickmesh(x, p, boundary = b, part = 1, numparts = n)\n for j in 2:n\n (etv_j, etc_j, etb_j, fc_j) =\n brickmesh(x, p, boundary = b, part = j, numparts = n)\n etv_parts = cat(etv_parts, etv_j; dims = 2)\n etc_parts = cat(etc_parts, etc_j; dims = 3)\n etb_parts = cat(etb_parts, etb_j; dims = 2)\n end\n\n @test etv == etv_parts\n @test etc == etc_parts\n @test etb == etb_parts\n end\n\n\n let\n x = (-1:2:10, -1:1:1, -4:1:1)\n p = (true, false, true)\n b = ((1, 2), (3, 4), (5, 6))\n\n (etv, etc, etb, fc) = brickmesh(x, p, boundary = b)\n\n n = 50\n (etv_parts, etc_parts, etb_parts, fc_parts) =\n brickmesh(x, p, boundary = b, part = 1, numparts = n)\n for j in 2:n\n (etv_j, etc_j, etb_j, fc_j) =\n brickmesh(x, p, boundary = b, part = j, numparts = n)\n etv_parts = cat(etv_parts, etv_j; dims = 2)\n etc_parts = cat(etc_parts, etc_j; dims = 3)\n etb_parts = cat(etb_parts, etb_j; dims = 2)\n end\n\n @test etv == etv_parts\n @test etc == etc_parts\n @test etb == etb_parts\n end\nend\n\n@testset \"Connect\" begin\n let\n comm = MPI.COMM_SELF\n\n mesh = connectmesh(\n comm,\n partition(comm, brickmesh((0:10,), (true,))...)[1:4]...,\n )\n\n nelem = 10\n\n @test mesh[:elemtocoord][:, :, 1] == [0 1]\n @test mesh[:elemtocoord][:, :, 2] == [1 2]\n @test mesh[:elemtocoord][:, :, 3] == [2 3]\n @test mesh[:elemtocoord][:, :, 4] == [3 4]\n @test mesh[:elemtocoord][:, :, 5] == [4 5]\n @test mesh[:elemtocoord][:, :, 6] == [5 6]\n @test mesh[:elemtocoord][:, :, 7] == [6 7]\n @test mesh[:elemtocoord][:, :, 8] == [7 8]\n @test mesh[:elemtocoord][:, :, 9] == [8 9]\n @test mesh[:elemtocoord][:, :, 10] == [9 10]\n\n @test mesh[:elemtoelem] == [\n 10 1 2 3 4 5 6 7 8 9\n 2 3 4 5 6 7 8 9 10 1\n ]\n\n @test mesh[:elemtoface] == [\n 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1\n ]\n\n @test mesh[:elemtoordr] == ones(Int, size(mesh[:elemtoordr]))\n @test mesh[:elemtobndy] == zeros(Int, size(mesh[:elemtoordr]))\n\n @test mesh[:elems] == 1:nelem\n @test mesh[:realelems] == 1:nelem\n @test mesh[:ghostelems] == nelem .+ (1:0)\n\n @test length(mesh[:sendelems]) == 0\n\n @test mesh[:nabrtorank] == Int[]\n @test mesh[:nabrtorecv] == UnitRange{Int}[]\n @test mesh[:nabrtosend] == UnitRange{Int}[]\n end\n\n let\n comm = MPI.COMM_SELF\n mesh = connectmesh(\n comm,\n partition(comm, brickmesh((0:4, 5:9), (false, true))...)[1:4]...,\n )\n\n nelem = 16\n\n @test mesh[:elemtocoord][:, :, 1] == [0 1 0 1; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 2] == [1 2 1 2; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 3] == [1 2 1 2; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 4] == [0 1 0 1; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 5] == [0 1 0 1; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 6] == [0 1 0 1; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 7] == [1 2 1 2; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 8] == [1 2 1 2; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 9] == [2 3 2 3; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 10] == [2 3 2 3; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 11] == [3 4 3 4; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 12] == [3 4 3 4; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 13] == [3 4 3 4; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 14] == [2 3 2 3; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 15] == [2 3 2 3; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 16] == [3 4 3 4; 5 5 6 6]\n\n @test mesh[:elemtoelem] == [\n 1 1 4 4 5 6 6 5 8 7 10 9 14 3 2 15\n 2 15 14 3 8 7 10 9 12 11 11 12 13 13 16 16\n 6 7 2 1 4 5 8 3 14 9 12 13 16 15 10 11\n 4 3 8 5 6 1 2 7 10 15 16 11 12 9 14 13\n ]\n\n @test mesh[:elemtoface] == [\n 1 2 2 1 1 1 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n @test mesh[:elemtoordr] == ones(Int, size(mesh[:elemtoordr]))\n\n @test mesh[:elems] == 1:nelem\n @test mesh[:realelems] == 1:nelem\n @test mesh[:ghostelems] == nelem .+ (1:0)\n\n @test length(mesh[:sendelems]) == 0\n\n @test mesh[:nabrtorank] == Int[]\n @test mesh[:nabrtorecv] == UnitRange{Int}[]\n @test mesh[:nabrtosend] == UnitRange{Int}[]\n end\n\n let\n comm = MPI.COMM_SELF\n mesh = connectmeshfull(\n comm,\n partition(comm, brickmesh((0:4, 5:9), (false, true))...)[1:4]...,\n )\n\n nelem = 16\n\n @test mesh[:elemtocoord][:, :, 1] == [0 1 0 1; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 2] == [1 2 1 2; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 3] == [1 2 1 2; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 4] == [0 1 0 1; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 5] == [0 1 0 1; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 6] == [0 1 0 1; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 7] == [1 2 1 2; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 8] == [1 2 1 2; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 9] == [2 3 2 3; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 10] == [2 3 2 3; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 11] == [3 4 3 4; 8 8 9 9]\n @test mesh[:elemtocoord][:, :, 12] == [3 4 3 4; 7 7 8 8]\n @test mesh[:elemtocoord][:, :, 13] == [3 4 3 4; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 14] == [2 3 2 3; 6 6 7 7]\n @test mesh[:elemtocoord][:, :, 15] == [2 3 2 3; 5 5 6 6]\n @test mesh[:elemtocoord][:, :, 16] == [3 4 3 4; 5 5 6 6]\n\n @test mesh[:elemtoelem] == [\n 1 1 4 4 5 6 6 5 8 7 10 9 14 3 2 15\n 2 15 14 3 8 7 10 9 12 11 11 12 13 13 16 16\n 6 7 2 1 4 5 8 3 14 9 12 13 16 15 10 11\n 4 3 8 5 6 1 2 7 10 15 16 11 12 9 14 13\n ]\n\n @test mesh[:elemtoface] == [\n 1 2 2 1 1 1 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n @test mesh[:elemtoordr] == ones(Int, size(mesh[:elemtoordr]))\n\n @test mesh[:elems] == 1:nelem\n @test mesh[:realelems] == 1:nelem\n @test mesh[:ghostelems] == nelem .+ (1:0)\n\n @test length(mesh[:sendelems]) == 0\n\n @test mesh[:nabrtorank] == Int[]\n @test mesh[:nabrtorecv] == UnitRange{Int}[]\n @test mesh[:nabrtosend] == UnitRange{Int}[]\n end\nend\n\n@testset \"Mappings\" begin\n let\n comm = MPI.COMM_SELF\n x = (0:4,)\n mesh =\n connectmesh(comm, partition(comm, brickmesh(x, (true,))...)[1:4]...)\n\n N = 3\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Nfp = (N + 1)^(d - 1)\n\n vmap⁻, vmap⁺ = mappings(\n ntuple(j -> N, d),\n mesh[:elemtoelem],\n mesh[:elemtoface],\n mesh[:elemtoordr],\n )\n\n @test vmap⁻ == reshape([1, 4, 5, 8, 9, 12, 13, 16], Nfp, nface, nelem)\n @test vmap⁺ == reshape([16, 5, 4, 9, 8, 13, 12, 1], Nfp, nface, nelem)\n end\n\n # Single polynomial order\n let\n comm = MPI.COMM_SELF\n x = (-1:1, 0:1)\n p = (false, true)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = 2\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Nfp = (N + 1)^(d - 1)\n\n vmap⁻, vmap⁺ = mappings(\n ntuple(j -> N, d),\n mesh[:elemtoelem],\n mesh[:elemtoface],\n mesh[:elemtoordr],\n )\n\n #! format: off\n @test vmap⁻ == reshape(\n [\n 1, 4, 7, # f=1 e=1\n 3, 6, 9, # f=2 e=1\n 1, 2, 3, # f=3 e=1\n 7, 8, 9, # f=4 e=1\n 10, 13, 16, # f=1 e=2\n 12, 15, 18, # f=2 e=2\n 10, 11, 12, # f=3 e=2\n 16, 17, 18, # f=4 e=2\n ],\n Nfp,\n nface,\n nelem,\n )\n\n @test vmap⁺ == reshape(\n [\n 1, 4, 7, # f=1 e=1\n 10, 13, 16, # f=1 e=2\n 7, 8, 9, # f=4 e=1\n 1, 2, 3, # f=3 e=1\n 3, 6, 9, # f=2 e=1\n 12, 15, 18, # f=2 e=2\n 16, 17, 18, # f=4 e=2\n 10, 11, 12, # f=3 e=2\n ],\n Nfp,\n nface,\n nelem,\n )\n #! format: on\n end\n\n # Two polynomial orders\n let\n comm = MPI.COMM_SELF\n x = (-1:1, 0:1)\n p = (false, true)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = (2, 3)\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Nfp = div.(prod(N .+ 1), ((N .+ 1) .^ (d - 1)))\n\n vmap⁻, vmap⁺ =\n mappings(N, mesh[:elemtoelem], mesh[:elemtoface], mesh[:elemtoordr])\n\n #! format: off\n @test vmap⁻ == reshape(\n [\n 1, 4, 7, 10, # f=1 e=1\n 3, 6, 9, 12, # f=2 e=1\n 1, 2, 3, 0, # f=3 e=1\n 10, 11, 12, 0, # f=4 e=1\n 13, 16, 19, 22, # f=1 e=2\n 15, 18, 21, 24, # f=2 e=2\n 13, 14, 15, 0, # f=3 e=2\n 22, 23, 24, 0, # f=4 e=2\n ],\n maximum(Nfp),\n nface,\n nelem,\n )\n\n @test vmap⁺ == reshape(\n [\n 1, 4, 7, 10, # f=1 e=1\n 13, 16, 19, 22, # f=2 e=1\n 10, 11, 12, 0, # f=3 e=1\n 1, 2, 3, 0, # f=4 e=1\n 3, 6, 9, 12, # f=1 e=2\n 15, 18, 21, 24, # f=2 e=2\n 22, 23, 24, 0, # f=3 e=2\n 13, 14, 15, 0, # f=4 e=2\n ],\n maximum(Nfp),\n nface,\n nelem,\n )\n #! format: on\n end\n\n # Single polynomial order\n let\n comm = MPI.COMM_SELF\n x = (0:1, 0:1, -1:1)\n p = (false, true, false)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = 2\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Np = (N + 1)^d\n Nfp = (N + 1)^(d - 1)\n\n vmap⁻, vmap⁺ = mappings(\n ntuple(j -> N, d),\n mesh[:elemtoelem],\n mesh[:elemtoface],\n mesh[:elemtoordr],\n )\n\n #! format: off\n fmask = [\n 1 3 1 7 1 19\n 4 6 2 8 2 20\n 7 9 3 9 3 21\n 10 12 10 16 4 22\n 13 15 11 17 5 23\n 16 18 12 18 6 24\n 19 21 19 25 7 25\n 22 24 20 26 8 26\n 25 27 21 27 9 27\n ]\n #! format: on\n\n @test vmap⁻ == reshape([fmask[:]; fmask[:] .+ Np], Nfp, nface, nelem)\n\n @test vmap⁺ == reshape(\n [\n fmask[:, 1]\n fmask[:, 2]\n fmask[:, 4]\n fmask[:, 3]\n fmask[:, 5]\n fmask[:, 5] .+ Np\n fmask[:, 1] .+ Np\n fmask[:, 2] .+ Np\n fmask[:, 4] .+ Np\n fmask[:, 3] .+ Np\n fmask[:, 6]\n fmask[:, 6] .+ Np\n ],\n Nfp,\n nface,\n nelem,\n )\n end\n\n # Multiple polynomial orders\n let\n comm = MPI.COMM_SELF\n x = (0:1, 0:1, -1:1)\n p = (false, true, false)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = (1, 2, 3)\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Np = prod(N .+ 1)\n Nfp = div.(Np, N .+ 1)\n\n vmap⁻, vmap⁺ =\n mappings(N, mesh[:elemtoelem], mesh[:elemtoface], mesh[:elemtoordr])\n\n #! format: off\n fmask =\n (\n [ 1 3 5 7 9 11 13 15 17 19 21 23],\n [ 2 4 6 8 10 12 14 16 18 20 22 24],\n [ 1 2 7 8 13 14 19 20 0 0 0 0],\n [ 5 6 11 12 17 18 23 24 0 0 0 0],\n [ 1 2 3 4 5 6 0 0 0 0 0 0],\n [19 20 21 22 23 24 0 0 0 0 0 0],\n )\n #! format: on\n\n @test vmap⁻ == reshape(\n [\n fmask[1]\n fmask[2]\n fmask[3]\n fmask[4]\n fmask[5]\n fmask[6]\n fmask[1] .+ Np .* (fmask[1] .> 0)\n fmask[2] .+ Np .* (fmask[2] .> 0)\n fmask[3] .+ Np .* (fmask[3] .> 0)\n fmask[4] .+ Np .* (fmask[4] .> 0)\n fmask[5] .+ Np .* (fmask[5] .> 0)\n fmask[6] .+ Np .* (fmask[6] .> 0)\n ]',\n maximum(Nfp),\n nface,\n nelem,\n )\n\n @test vmap⁺ == reshape(\n [\n vmap⁻[:, 1, 1]\n vmap⁻[:, 2, 1]\n vmap⁻[:, 4, 1]\n vmap⁻[:, 3, 1]\n vmap⁻[:, 5, 1]\n vmap⁻[:, 5, 2]\n vmap⁻[:, 1, 2]\n vmap⁻[:, 2, 2]\n vmap⁻[:, 4, 2]\n vmap⁻[:, 3, 2]\n vmap⁻[:, 6, 1]\n vmap⁻[:, 6, 2]\n ],\n maximum(Nfp),\n nface,\n nelem,\n )\n end\n\n # Test polynomial order 0\n let\n comm = MPI.COMM_SELF\n x = (0:1, -1:1)\n p = (false, true)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = (5, 0)\n Nq = N .+ 1\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n vmap⁻, vmap⁺ =\n mappings(N, mesh[:elemtoelem], mesh[:elemtoface], mesh[:elemtoordr])\n\n #! format: off\n p = reshape(1:Np, Nq)\n fmask = ntuple(j -> zeros(Int, maximum(Nfp)), 2d)\n fmask[1][1:Nfp[1]] .= p[1, :][:]\n fmask[2][1:Nfp[1]] .= p[end, :][:]\n fmask[3][1:Nfp[2]] .= p[:, 1][:]\n fmask[4][1:Nfp[2]] .= p[:, end][:]\n #! format: on\n\n @test vmap⁻ == reshape(\n [\n fmask[1]\n fmask[2]\n fmask[3]\n fmask[4]\n fmask[1] .+ Np .* (fmask[1] .> 0)\n fmask[2] .+ Np .* (fmask[2] .> 0)\n fmask[3] .+ Np .* (fmask[3] .> 0)\n fmask[4] .+ Np .* (fmask[4] .> 0)\n ]',\n maximum(Nfp),\n nface,\n nelem,\n )\n\n @test vmap⁺ == reshape(\n [\n vmap⁻[:, 1, 1]\n vmap⁻[:, 2, 1]\n vmap⁻[:, 4, 2]\n vmap⁻[:, 3, 2]\n vmap⁻[:, 1, 2]\n vmap⁻[:, 2, 2]\n vmap⁻[:, 4, 1]\n vmap⁻[:, 3, 1]\n ],\n maximum(Nfp),\n nface,\n nelem,\n )\n end\n\n # Test polynomial order 0\n let\n comm = MPI.COMM_SELF\n x = (0:1, 0:1, -1:1)\n p = (false, true, false)\n mesh = connectmesh(comm, partition(comm, brickmesh(x, p)...)[1:4]...)\n\n N = (7, 2, 0)\n Nq = N .+ 1\n d = length(x)\n nelem = prod(length.(x) .- 1)\n nface = 2d\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n vmap⁻, vmap⁺ =\n mappings(N, mesh[:elemtoelem], mesh[:elemtoface], mesh[:elemtoordr])\n\n #! format: off\n p = reshape(1:Np, Nq)\n fmask = ntuple(j -> zeros(Int, maximum(Nfp)), 2d)\n fmask[1][1:Nfp[1]] .= p[1, :, :][:]\n fmask[2][1:Nfp[1]] .= p[end, :, :][:]\n fmask[3][1:Nfp[2]] .= p[:, 1, :][:]\n fmask[4][1:Nfp[2]] .= p[:, end, :][:]\n fmask[5][1:Nfp[3]] .= p[:, :, 1][:]\n fmask[6][1:Nfp[3]] .= p[:, :, end][:]\n #! format: on\n\n @test vmap⁻ == reshape(\n [\n fmask[1]\n fmask[2]\n fmask[3]\n fmask[4]\n fmask[5]\n fmask[6]\n fmask[1] .+ Np .* (fmask[1] .> 0)\n fmask[2] .+ Np .* (fmask[2] .> 0)\n fmask[3] .+ Np .* (fmask[3] .> 0)\n fmask[4] .+ Np .* (fmask[4] .> 0)\n fmask[5] .+ Np .* (fmask[5] .> 0)\n fmask[6] .+ Np .* (fmask[6] .> 0)\n ]',\n maximum(Nfp),\n nface,\n nelem,\n )\n\n @test vmap⁺ == reshape(\n [\n vmap⁻[:, 1, 1]\n vmap⁻[:, 2, 1]\n vmap⁻[:, 4, 1]\n vmap⁻[:, 3, 1]\n vmap⁻[:, 5, 1]\n vmap⁻[:, 5, 2]\n vmap⁻[:, 1, 2]\n vmap⁻[:, 2, 2]\n vmap⁻[:, 4, 2]\n vmap⁻[:, 3, 2]\n vmap⁻[:, 6, 1]\n vmap⁻[:, 6, 2]\n ],\n maximum(Nfp),\n nface,\n nelem,\n )\n end\nend\n\n@testset \"Get Partition\" begin\n let\n Nelem = 150\n (so, ss, rs) = BrickMesh.getpartition(MPI.COMM_SELF, Nelem:-1:1)\n @test so == Nelem:-1:1\n @test ss == [1, Nelem + 1]\n @test rs == [1, Nelem + 1]\n end\n\n let\n Nelem = 111\n code = [ones(1, Nelem); collect(Nelem:-1:1)']\n (so, ss, rs) = BrickMesh.getpartition(MPI.COMM_SELF, Nelem:-1:1)\n @test so == Nelem:-1:1\n @test ss == [1, Nelem + 1]\n @test rs == [1, Nelem + 1]\n end\nend\n\n@testset \"Partition\" begin\n (etv, etc, etb, fc) =\n brickmesh((-1:2:1, -1:2:1, -2:1:2), (true, true, true))\n (netv, netc, netb, nfc) = partition(MPI.COMM_SELF, etv, etc, etb, fc)[1:4]\n @test etv == netv\n @test etc == netc\n @test etb == netb\n @test fc == nfc\nend\n\n@testset \"Comm Mappings\" begin\n let\n N = 1\n d = 2\n nface = 2d\n\n commelems = [1, 2, 5]\n commfaces = BitArray(undef, nface, length(commelems))\n commfaces .= false\n nabrtocomm = [1:2, 3:3]\n\n vmapC, nabrtovmapC =\n commmapping(ntuple(j -> N, d), commelems, commfaces, nabrtocomm)\n\n @test vmapC == Int[]\n @test nabrtovmapC == UnitRange{Int64}[1:0, 1:0]\n end\n\n let\n N = 1\n d = 2\n nface = 2d\n\n commelems = [1, 2, 5]\n commfaces = BitArray([\n false false false\n false true false\n false true false\n false false true\n ])\n nabrtocomm = [1:2, 3:3]\n\n vmapC, nabrtovmapC =\n commmapping(ntuple(j -> N, d), commelems, commfaces, nabrtocomm)\n\n @test vmapC == [5, 6, 8, 19, 20]\n @test nabrtovmapC == UnitRange{Int64}[1:3, 4:5]\n end\n\n # 2D, single polynomial order\n let\n N = 2\n d = 2\n nface = 2d\n\n commelems = [2, 4, 5]\n commfaces = BitArray([\n true true false\n false false false\n false true true\n false false true\n ])\n nabrtocomm = [1:1, 2:3]\n\n vmapC, nabrtovmapC =\n commmapping(ntuple(j -> N, d), commelems, commfaces, nabrtocomm)\n\n @test vmapC == [10, 13, 16, 28, 29, 30, 31, 34, 37, 38, 39, 43, 44, 45]\n @test nabrtovmapC == UnitRange{Int64}[1:3, 4:14]\n end\n\n # 2D, multiple polynomial orders\n let\n N = (2, 3)\n Np = prod(N .+ 1)\n d = length(N)\n nface = 2d\n\n commelems = [2, 4, 5]\n #! format: off\n commfaces = BitArray([\n true true false false # faces of commelems[1] to send\n false false false true # faces of commelems[2] to send\n true false false true # faces of commelems[3] to send\n ]')\n #! format: on\n nabrtocomm = [1:1, 2:3]\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n @test vmapC == [\n (commelems[1] - 1) * Np + 1\n (commelems[1] - 1) * Np + 3\n (commelems[1] - 1) * Np + 4\n (commelems[1] - 1) * Np + 6\n (commelems[1] - 1) * Np + 7\n (commelems[1] - 1) * Np + 9\n (commelems[1] - 1) * Np + 10\n (commelems[1] - 1) * Np + 12\n (commelems[2] - 1) * Np + 10\n (commelems[2] - 1) * Np + 11\n (commelems[2] - 1) * Np + 12\n (commelems[3] - 1) * Np + 1\n (commelems[3] - 1) * Np + 4\n (commelems[3] - 1) * Np + 7\n (commelems[3] - 1) * Np + 10\n (commelems[3] - 1) * Np + 11\n (commelems[3] - 1) * Np + 12\n ]\n\n @test nabrtovmapC == UnitRange{Int64}[1:8, 9:17]\n end\n\n # 3D, single polynomial order\n let\n N = 2\n d = 3\n nface = 2d\n\n commelems = [3, 4, 7, 9]\n commfaces = BitArray([\n true true true false\n false true false false\n false true false false\n false true false false\n false true true true\n false true false true\n ])\n nabrtocomm = [1:1, 2:4]\n\n vmapC, nabrtovmapC =\n commmapping(ntuple(j -> N, d), commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC == [\n 55, 58, 61, 64, 67, 70, 73, 76, 79, 82, 83, 84, 85, 86, 87, 88, 89,\n 90, 91, 92, 93, 94, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105,\n 106, 107, 108, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172,\n 175, 178, 181, 184, 187, 217, 218, 219, 220, 221, 222, 223, 224,\n 225, 235, 236, 237, 238, 239, 240, 241, 242, 243,\n ]\n #! format: on\n @test nabrtovmapC == UnitRange{Int64}[1:9, 10:68]\n end\n\n # 3D, multiple polynomial order\n let\n N = (2, 3, 4)\n Nq = N .+ 1\n Np = prod(Nq)\n d = length(N)\n nface = 2d\n\n commelems = [3, 4, 7, 9]\n commfaces = BitArray([\n true true true false\n false true false false\n false true false false\n false true false false\n false true true true\n false true false true\n ])\n\n p = reshape(1:Np, Nq)\n fmask = (\n p[1, :, :][:], # Face 1\n p[Nq[1], :, :][:], # Face 2\n p[:, 1, :][:], # Face 3\n p[:, Nq[2], :][:], # Face 4\n p[:, :, 1][:], # Face 5\n p[:, :, Nq[3]][:], # Face 6\n )\n\n nabrtocomm = [1:1, 2:4]\n\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC ==\n [\n sort(unique(vcat(fmask[commfaces[:, 1]]...))) .+ (commelems[1] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 2]]...))) .+ (commelems[2] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 3]]...))) .+ (commelems[3] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 4]]...))) .+ (commelems[4] - 1) *Np\n ]\n #! format: on\n\n @test nabrtovmapC == UnitRange{Int64}[1:20, 21:126]\n end\n\n # Test mappings with polyorder 0 in one dimension\n let\n N = (0, 1, 2)\n Nq = N .+ 1\n Np = prod(Nq)\n d = length(N)\n nface = 2d\n\n commelems = [3, 4, 7, 9]\n commfaces = BitArray([\n true true true false\n false true false false\n false true false false\n false true false false\n false true true true\n false true false true\n ])\n\n p = reshape(1:Np, Nq)\n fmask = (\n p[1, :, :][:], # Face 1\n p[Nq[1], :, :][:], # Face 2\n p[:, 1, :][:], # Face 3\n p[:, Nq[2], :][:], # Face 4\n p[:, :, 1][:], # Face 5\n p[:, :, Nq[3]][:], # Face 6\n )\n\n nabrtocomm = [1:1, 2:4]\n\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC ==\n [\n sort(unique(vcat(fmask[commfaces[:, 1]]...))) .+ (commelems[1] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 2]]...))) .+ (commelems[2] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 3]]...))) .+ (commelems[3] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 4]]...))) .+ (commelems[4] - 1) *Np\n ]\n #! format: on\n\n @test nabrtovmapC == UnitRange{Int64}[1:6, 7:22]\n end\n\n let\n N = (3, 2, 0)\n Nq = N .+ 1\n Np = prod(Nq)\n d = length(N)\n nface = 2d\n\n commelems = [3, 4, 7, 9, 19]\n commfaces = BitArray([\n true true true false true\n false true false false false\n false true false false true\n false true false false false\n false true true true false\n false true false true false\n ])\n\n p = reshape(1:Np, Nq)\n fmask = (\n p[1, :, :][:], # Face 1\n p[Nq[1], :, :][:], # Face 2\n p[:, 1, :][:], # Face 3\n p[:, Nq[2], :][:], # Face 4\n p[:, :, 1][:], # Face 5\n p[:, :, Nq[3]][:], # Face 6\n )\n\n nabrtocomm = [1:1, 2:5]\n\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC ==\n [\n sort(unique(vcat(fmask[commfaces[:, 1]]...))) .+ (commelems[1] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 2]]...))) .+ (commelems[2] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 3]]...))) .+ (commelems[3] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 4]]...))) .+ (commelems[4] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 5]]...))) .+ (commelems[5] - 1) *Np\n ]\n #! format: on\n\n @test nabrtovmapC == UnitRange{Int64}[1:3, 4:45]\n end\n\n let\n N = (0, 2)\n Nq = N .+ 1\n Np = prod(Nq)\n d = length(N)\n nface = 2d\n\n commelems = [3, 4, 7, 9]\n commfaces = BitArray([\n true true true false\n false true false true\n false true false false\n false true false true\n ])\n\n p = reshape(1:Np, Nq)\n fmask = (\n p[1, :][:], # Face 1\n p[Nq[1], :][:], # Face 2\n p[:, 1][:], # Face 3\n p[:, Nq[2]][:], # Face 4\n )\n\n nabrtocomm = [1:1, 2:4]\n\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC ==\n [\n sort(unique(vcat(fmask[commfaces[:, 1]]...))) .+ (commelems[1] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 2]]...))) .+ (commelems[2] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 3]]...))) .+ (commelems[3] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 4]]...))) .+ (commelems[4] - 1) *Np\n ]\n #! format: on\n\n @test nabrtovmapC == UnitRange{Int64}[1:3, 4:12]\n end\n\n let\n N = (3, 0)\n Nq = N .+ 1\n Np = prod(Nq)\n d = length(N)\n nface = 2d\n\n commelems = [3, 4, 7, 9, 19]\n commfaces = BitArray([\n true true true false true\n false true false false true\n false true false true false\n false true false true false\n ])\n\n p = reshape(1:Np, Nq)\n fmask = (\n p[1, :][:], # Face 1\n p[Nq[1], :][:], # Face 2\n p[:, 1][:], # Face 3\n p[:, Nq[2]][:], # Face 4\n )\n\n nabrtocomm = [1:1, 2:5]\n\n vmapC, nabrtovmapC = commmapping(N, commelems, commfaces, nabrtocomm)\n\n #! format: off\n @test vmapC ==\n [\n sort(unique(vcat(fmask[commfaces[:, 1]]...))) .+ (commelems[1] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 2]]...))) .+ (commelems[2] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 3]]...))) .+ (commelems[3] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 4]]...))) .+ (commelems[4] - 1) *Np\n sort(unique(vcat(fmask[commfaces[:, 5]]...))) .+ (commelems[5] - 1) *Np\n ]\n #! format: on\n\n @test nabrtovmapC == UnitRange{Int64}[1:1, 2:12]\n end\n\nend\n" }, { "path": "test/Numerics/Mesh/DSS.jl", "content": "using Test\nusing MPI\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.Mesh.DSS\nusing KernelAbstractions\n\nDA = ClimateMachine.array_type()\n\nfunction test_dss()\n FT = Float64\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 1\n\n N = (4, 4, 5)\n brickrange = (0:2, 5:6, 0:1)\n periodicity = (false, false, false)\n nvars = 1\n\n topl = StackedBrickTopology(\n comm,\n brickrange,\n periodicity = periodicity,\n boundary = ((1, 2), (3, 4), (5, 6)),\n connectivity = :full,\n )\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = DA,\n polynomialorder = N,\n )\n\n Nq = N .+ 1\n Np = prod(Nq)\n Q = MPIStateArray{FT}(\n comm,\n DA,\n Np,\n nvars,\n length(topl.elems),\n realelems = topl.realelems,\n ghostelems = topl.ghostelems,\n vmaprecv = grid.vmaprecv,\n vmapsend = grid.vmapsend,\n nabrtorank = topl.nabrtorank,\n nabrtovmaprecv = grid.nabrtovmaprecv,\n nabrtovmapsend = grid.nabrtovmapsend,\n )\n realelems = Q.realelems\n ghostelems = Q.ghostelems\n\n ldof = DA{FT, 1}(reshape(1:Np, Np))\n\n for ivar in 1:nvars\n for ielem in realelems\n Q.data[:, ivar, ielem] .= ldof .+ FT((ielem - 1) * Np)\n end\n Q.data[:, ivar, ghostelems] .= FT(0)\n end\n\n pre_dss = Array(Q.data)\n\n dss!(Q, grid)\n #---------Tests-------------------------\n nodes = reshape(1:Np, Nq)\n interior = nodes[2:(Nq[1] - 1), 2:(Nq[2] - 1), 2:(Nq[3] - 1)][:]\n vertmap = Array(grid.vertmap)\n edgemap = Array(grid.edgemap)\n facemap = Array(grid.facemap)\n post_dss = Array(Q.data)\n\n ne1r = 1:max(Nq[1] - 2, 0)\n ne2r = 1:max(Nq[2] - 2, 0)\n ne3r = 1:max(Nq[3] - 2, 0)\n\n nf1r = 1:max((Nq[2] - 2) * (Nq[3] - 2), 0)\n nf2r = 1:max((Nq[1] - 2) * (Nq[3] - 2), 0)\n nf3r = 1:max((Nq[1] - 2) * (Nq[2] - 2), 0)\n\n compare(el1, i1, el2, i2, efmap, rng, post, pre) =\n post[efmap[rng, i1, 1], 1, el1] == pre[efmap[rng, i2, 1], 1, el2]\n\n compare(el1, i1, el2, i2, el3, i3, efmap, rng, post, pre) =\n post[efmap[rng, i1, 1], 1, el1] ==\n pre[efmap[rng, i2, 1], 1, el2] .+ pre[efmap[rng, i3, 1], 1, el3]\n\n el1, el2 = 1, 2\n # Element # 1 --------------------------------------------------------------------------\n # interior dof should not be affected\n @test pre_dss[interior, 1, 1] == post_dss[interior, 1, 1]\n # vertex check\n @test post_dss[vertmap, 1, 1] == [\n pre_dss[vertmap[1], 1, 1],\n pre_dss[vertmap[2], 1, 1] + pre_dss[vertmap[1], 1, 2],\n pre_dss[vertmap[3], 1, 1],\n pre_dss[vertmap[4], 1, 1] + pre_dss[vertmap[3], 1, 2],\n pre_dss[vertmap[5], 1, 1],\n pre_dss[vertmap[6], 1, 1] + pre_dss[vertmap[5], 1, 2],\n pre_dss[vertmap[7], 1, 1],\n pre_dss[vertmap[8], 1, 1] + pre_dss[vertmap[7], 1, 2],\n ]\n # edge check\n @test compare(el1, 1, el1, 1, edgemap, ne1r, post_dss, pre_dss) # edge 1 (unaffected)\n @test compare(el1, 2, el1, 2, edgemap, ne1r, post_dss, pre_dss) # edge 2 (unaffected)\n @test compare(el1, 3, el1, 3, edgemap, ne1r, post_dss, pre_dss) # edge 3 (unaffected)\n @test compare(el1, 4, el1, 4, edgemap, ne1r, post_dss, pre_dss) # edge 4 (unaffected)\n\n @test compare(el1, 5, el1, 5, edgemap, ne2r, post_dss, pre_dss) # edge 5 (unaffected)\n @test compare(el1, 6, el1, 6, el2, 5, edgemap, ne2r, post_dss, pre_dss) # edge 6 (shared edge)\n @test compare(el1, 7, el1, 7, edgemap, ne2r, post_dss, pre_dss) # edge 7 (unaffected)\n @test compare(el1, 8, el1, 8, el2, 7, edgemap, ne2r, post_dss, pre_dss) # edge 8 (shared edge)\n\n @test compare(el1, 9, el1, 9, edgemap, ne3r, post_dss, pre_dss) # edge 9 (unaffected)\n @test compare(el1, 10, el1, 10, el2, 9, edgemap, ne3r, post_dss, pre_dss) # edge 10 (shared edge)\n @test compare(el1, 11, el1, 11, edgemap, ne3r, post_dss, pre_dss) # edge 11 (unaffected)\n @test compare(el1, 12, el1, 12, el2, 11, edgemap, ne3r, post_dss, pre_dss) # edge 12 (shared edge)\n\n # face check\n @test compare(el1, 1, el1, 1, facemap, nf1r, post_dss, pre_dss) # face 1 (unaffected)\n @test compare(el1, 2, el1, 2, el2, 1, facemap, nf1r, post_dss, pre_dss) # face 2 (shared face)\n\n @test compare(el1, 3, el1, 3, facemap, nf2r, post_dss, pre_dss) # face 3 (unaffected)\n @test compare(el1, 4, el1, 4, facemap, nf2r, post_dss, pre_dss) # face 4 (unaffected)\n\n @test compare(el1, 5, el1, 5, facemap, nf3r, post_dss, pre_dss) # face 5 (unaffected)\n @test compare(el1, 6, el1, 6, facemap, nf3r, post_dss, pre_dss) # face 6 (unaffected)\n\n # Element # 2 --------------------------------------------------------------------------\n # interior dof should not be affected\n @test pre_dss[interior, 1, 2] == post_dss[interior, 1, 2]\n # vertex check\n @test post_dss[vertmap, 1, 2] == [\n pre_dss[vertmap[1], 1, 2] + pre_dss[vertmap[2], 1, 1],\n pre_dss[vertmap[2], 1, 2],\n pre_dss[vertmap[3], 1, 2] + pre_dss[vertmap[4], 1, 1],\n pre_dss[vertmap[4], 1, 2],\n pre_dss[vertmap[5], 1, 2] + pre_dss[vertmap[6], 1, 1],\n pre_dss[vertmap[6], 1, 2],\n pre_dss[vertmap[7], 1, 2] + pre_dss[vertmap[8], 1, 1],\n pre_dss[vertmap[8], 1, 2],\n ]\n # edge check\n @test compare(el2, 1, el2, 1, edgemap, ne1r, post_dss, pre_dss) # edge 1 (unaffected)\n @test compare(el2, 2, el2, 2, edgemap, ne1r, post_dss, pre_dss) # edge 2 (unaffected)\n @test compare(el2, 3, el2, 3, edgemap, ne1r, post_dss, pre_dss) # edge 3 (unaffected)\n @test compare(el2, 4, el2, 4, edgemap, ne1r, post_dss, pre_dss) # edge 4 (unaffected)\n\n @test compare(el2, 5, el2, 5, el1, 6, edgemap, ne2r, post_dss, pre_dss) # edge 5 (shared edge)\n @test compare(el2, 6, el2, 6, edgemap, ne2r, post_dss, pre_dss) # edge 6 (unaffected)\n @test compare(el2, 7, el2, 7, el1, 8, edgemap, ne2r, post_dss, pre_dss) # edge 7 (shared edge)\n @test compare(el2, 8, el2, 8, edgemap, ne2r, post_dss, pre_dss) # edge 8 (unaffected)\n\n @test compare(el2, 9, el2, 9, el1, 10, edgemap, ne3r, post_dss, pre_dss) # edge 9 (shared edge)\n @test compare(el2, 10, el2, 10, edgemap, ne3r, post_dss, pre_dss) # edge 10 (unaffected)\n @test compare(el2, 11, el2, 11, el1, 12, edgemap, ne3r, post_dss, pre_dss) # edge 11 (shared edge)\n @test compare(el2, 12, el2, 12, edgemap, ne3r, post_dss, pre_dss) # edge 12 (unaffected)\n\n # face check\n @test compare(el2, 1, el2, 1, el1, 2, facemap, nf1r, post_dss, pre_dss) # face 1 (shared face)\n @test compare(el2, 2, el2, 2, facemap, nf1r, post_dss, pre_dss) # face 2 (unaffected)\n\n @test compare(el2, 3, el2, 3, facemap, nf2r, post_dss, pre_dss) # face 3 (unaffected)\n @test compare(el2, 4, el2, 4, facemap, nf2r, post_dss, pre_dss) # face 4 (unaffected)\n\n @test compare(el2, 5, el2, 5, facemap, nf3r, post_dss, pre_dss) # face 5 (unaffected)\n @test compare(el2, 6, el2, 6, facemap, nf3r, post_dss, pre_dss) # face 6 (unaffected)\n\nend\n\ntest_dss()\n" }, { "path": "test/Numerics/Mesh/DSS_mpi.jl", "content": "using Test\nusing MPI\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.Mesh.DSS\nusing KernelAbstractions\n\nDA = ClimateMachine.array_type()\n\nfunction test_dss_stacked_3d()\n FT = Float64\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 3\n\n N = (4, 4, 5)\n brickrange = (0:4, 5:9, 0:3)\n periodicity = (false, false, false)\n nvars = 3\n\n topl = StackedBrickTopology(\n comm,\n brickrange,\n periodicity = periodicity,\n boundary = ((1, 2), (3, 4), (5, 6)),\n connectivity = :full,\n )\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = DA,\n polynomialorder = N,\n )\n\n Nq = N .+ 1\n Np = prod(Nq)\n Q = MPIStateArray{FT}(\n comm,\n DA,\n Np,\n nvars,\n length(topl.elems),\n realelems = topl.realelems,\n ghostelems = topl.ghostelems,\n vmaprecv = grid.vmaprecv,\n vmapsend = grid.vmapsend,\n nabrtorank = topl.nabrtorank,\n nabrtovmaprecv = grid.nabrtovmaprecv,\n nabrtovmapsend = grid.nabrtovmapsend,\n )\n realelems = Q.realelems\n ghostelems = Q.ghostelems\n\n Q.data[:, :, realelems] .= FT(1)\n Q.data[:, :, ghostelems] .= FT(0)\n\n dss!(Q, grid)\n #---------Tests-------------------------\n nodes = reshape(1:Np, Nq)\n interior = nodes[2:(Nq[1] - 1), 2:(Nq[2] - 1), 2:(Nq[3] - 1)][:]\n vertmap = Array(grid.vertmap)\n edgemap = Array(grid.edgemap)\n facemap = Array(grid.facemap)\n data = Array(Q.data)\n compare(data, ivar, iel, efmap) =\n unique(data[setdiff(efmap, [-1]), ivar, lel])\n lel = 1 # local element number for each process\n\n if crank == 0\n for ivar in 1:nvars\n # local element #1, global elememt # 1\n # interior dof should not be affected\n idof = unique(data[interior, ivar, lel])\n @test idof == [FT(1)]\n # vertex dof check\n @test data[vertmap, ivar, lel] == FT.([1, 2, 2, 4, 2, 4, 4, 8])\n # edge dof check\n @test compare(data, ivar, lel, edgemap[:, 1, 1]) == [FT(1)] # edge # 1\n @test compare(data, ivar, lel, edgemap[:, 2, 1]) == [FT(2)] # edge # 2\n @test compare(data, ivar, lel, edgemap[:, 3, 1]) == [FT(2)] # edge # 3\n @test compare(data, ivar, lel, edgemap[:, 4, 1]) == [FT(4)] # edge # 4\n @test compare(data, ivar, lel, edgemap[:, 5, 1]) == [FT(1)] # edge # 5\n @test compare(data, ivar, lel, edgemap[:, 6, 1]) == [FT(2)] # edge # 6\n @test compare(data, ivar, lel, edgemap[:, 7, 1]) == [FT(2)] # edge # 7\n @test compare(data, ivar, lel, edgemap[:, 8, 1]) == [FT(4)] # edge # 8\n @test compare(data, ivar, lel, edgemap[:, 9, 1]) == [FT(1)] # edge # 9\n @test compare(data, ivar, lel, edgemap[:, 10, 1]) == [FT(2)] # edge # 10\n @test compare(data, ivar, lel, edgemap[:, 11, 1]) == [FT(2)] # edge # 11\n @test compare(data, ivar, lel, edgemap[:, 12, 1]) == [FT(4)] # edge # 12\n\n # face dof check\n @test compare(data, ivar, lel, facemap[:, 1, 1]) == [FT(1)] # face # 1\n @test compare(data, ivar, lel, facemap[:, 2, 1]) == [FT(2)] # face # 2\n @test compare(data, ivar, lel, facemap[:, 3, 1]) == [FT(1)] # face # 3\n @test compare(data, ivar, lel, facemap[:, 4, 1]) == [FT(2)] # face # 4\n @test compare(data, ivar, lel, facemap[:, 5, 1]) == [FT(1)] # face # 5\n @test compare(data, ivar, lel, facemap[:, 6, 1]) == [FT(2)] # face # 6\n end\n\n elseif crank == 1\n for ivar in 1:nvars\n # local element #1, global elememt # 6 (in 2D)\n # interior dof should not be affected\n idof = unique(data[interior, ivar, lel])\n @test idof == [FT(1)]\n # vertex dof check\n @test data[vertmap, ivar, lel] == FT.([2, 4, 1, 2, 4, 8, 2, 4])\n # edge dof check\n @test compare(data, ivar, lel, edgemap[:, 1, 1]) == [FT(2)] # edge # 1\n @test compare(data, ivar, lel, edgemap[:, 2, 1]) == [FT(1)] # edge # 2\n @test compare(data, ivar, lel, edgemap[:, 3, 1]) == [FT(4)] # edge # 3\n @test compare(data, ivar, lel, edgemap[:, 4, 1]) == [FT(2)] # edge # 4\n @test compare(data, ivar, lel, edgemap[:, 5, 1]) == [FT(1)] # edge # 5\n @test compare(data, ivar, lel, edgemap[:, 6, 1]) == [FT(2)] # edge # 6\n @test compare(data, ivar, lel, edgemap[:, 7, 1]) == [FT(2)] # edge # 7\n @test compare(data, ivar, lel, edgemap[:, 8, 1]) == [FT(4)] # edge # 8\n @test compare(data, ivar, lel, edgemap[:, 9, 1]) == [FT(2)] # edge # 9\n @test compare(data, ivar, lel, edgemap[:, 10, 1]) == [FT(4)] # edge # 10\n @test compare(data, ivar, lel, edgemap[:, 11, 1]) == [FT(1)] # edge # 11\n @test compare(data, ivar, lel, edgemap[:, 12, 1]) == [FT(2)] # edge # 12\n\n\n # face dof check\n @test compare(data, ivar, lel, facemap[:, 1, 1]) == [FT(1)] # face # 1\n @test compare(data, ivar, lel, facemap[:, 2, 1]) == [FT(2)] # face # 2\n @test compare(data, ivar, lel, facemap[:, 3, 1]) == [FT(2)] # face # 3\n @test compare(data, ivar, lel, facemap[:, 4, 1]) == [FT(1)] # face # 4\n @test compare(data, ivar, lel, facemap[:, 5, 1]) == [FT(1)] # face # 5\n @test compare(data, ivar, lel, facemap[:, 6, 1]) == [FT(2)] # face # 6\n end\n else # crank == 2\n for ivar in 1:nvars\n # local element #1, global elememt # 11 (in 2D)\n # interior dof should not be affected\n idof = unique(data[interior, ivar, lel])\n @test idof == [FT(1)]\n # vertex dof check\n @test data[vertmap, ivar, lel] == FT.([4, 2, 2, 1, 8, 4, 4, 2])\n # edge dof check\n @test compare(data, ivar, lel, edgemap[:, 1, 1]) == [FT(2)] # edge # 1\n @test compare(data, ivar, lel, edgemap[:, 2, 1]) == [FT(1)] # edge # 2\n @test compare(data, ivar, lel, edgemap[:, 3, 1]) == [FT(4)] # edge # 3\n @test compare(data, ivar, lel, edgemap[:, 4, 1]) == [FT(2)] # edge # 4\n @test compare(data, ivar, lel, edgemap[:, 5, 1]) == [FT(2)] # edge # 5\n @test compare(data, ivar, lel, edgemap[:, 6, 1]) == [FT(1)] # edge # 6\n @test compare(data, ivar, lel, edgemap[:, 7, 1]) == [FT(4)] # edge # 7\n @test compare(data, ivar, lel, edgemap[:, 8, 1]) == [FT(2)] # edge # 8\n @test compare(data, ivar, lel, edgemap[:, 9, 1]) == [FT(4)] # edge # 9\n @test compare(data, ivar, lel, edgemap[:, 10, 1]) == [FT(2)] # edge # 10\n @test compare(data, ivar, lel, edgemap[:, 11, 1]) == [FT(2)] # edge # 11\n @test compare(data, ivar, lel, edgemap[:, 12, 1]) == [FT(1)] # edge # 12\n # face dof check\n @test compare(data, ivar, lel, facemap[:, 1, 1]) == [FT(2)] # face # 1\n @test compare(data, ivar, lel, facemap[:, 2, 1]) == [FT(1)] # face # 2\n @test compare(data, ivar, lel, facemap[:, 3, 1]) == [FT(2)] # face # 3\n @test compare(data, ivar, lel, facemap[:, 4, 1]) == [FT(1)] # face # 4\n @test compare(data, ivar, lel, facemap[:, 5, 1]) == [FT(1)] # face # 5\n @test compare(data, ivar, lel, facemap[:, 6, 1]) == [FT(2)] # face # 6\n end\n end\nend\n\ntest_dss_stacked_3d()\n" }, { "path": "test/Numerics/Mesh/Elements.jl", "content": "using ClimateMachine.Mesh.Elements\nusing GaussQuadrature\nusing LinearAlgebra\nusing Test\n\n@testset \"GaussQuadrature\" begin\n for T in (Float32, Float64, BigFloat)\n let\n x, w = GaussQuadrature.legendre(T, 1)\n @test iszero(x)\n @test w ≈ [2 * one(T)]\n end\n\n let\n endpt = GaussQuadrature.left\n x, w = GaussQuadrature.legendre(T, 1, endpt)\n @test x ≈ [-one(T)]\n @test w ≈ [2 * one(T)]\n end\n\n let\n endpt = GaussQuadrature.right\n x, w = GaussQuadrature.legendre(T, 1, endpt)\n @test x ≈ [one(T)]\n @test w ≈ [2 * one(T)]\n end\n\n let\n endpt = GaussQuadrature.left\n x, w = GaussQuadrature.legendre(T, 2, endpt)\n @test x ≈ [-one(T); T(1 // 3)]\n @test w ≈ [T(1 // 2); T(3 // 2)]\n end\n\n let\n endpt = GaussQuadrature.right\n x, w = GaussQuadrature.legendre(T, 2, endpt)\n @test x ≈ [T(-1 // 3); one(T)]\n @test w ≈ [T(3 // 2); T(1 // 2)]\n end\n end\n\n let\n err = ErrorException(\"Must have at least two points for both ends.\")\n endpt = GaussQuadrature.both\n @test_throws err GaussQuadrature.legendre(1, endpt)\n end\n\n let\n T = Float64\n n = 100\n endpt = GaussQuadrature.both\n\n a, b = GaussQuadrature.legendre_coefs(T, n)\n\n err = ErrorException(\n \"No convergence after 1 iterations \" * \"(try increasing maxits)\",\n )\n\n @test_throws err GaussQuadrature.custom_gauss_rule(\n -one(T),\n one(T),\n a,\n b,\n endpt,\n 1,\n )\n end\nend\n\n@testset \"Operators\" begin\n P5(r::AbstractVector{T}) where {T} =\n T(1) / T(8) * (T(15) * r - T(70) * r .^ 3 + T(63) * r .^ 5)\n\n P6(r::AbstractVector{T}) where {T} =\n T(1) / T(16) *\n (-T(5) .+ T(105) * r .^ 2 - T(315) * r .^ 4 + T(231) * r .^ 6)\n DP6(r::AbstractVector{T}) where {T} =\n T(1) / T(16) *\n (T(2 * 105) * r - T(4 * 315) * r .^ 3 + T(6 * 231) * r .^ 5)\n\n IPN(::Type{T}, N) where {T} = T(2) / T(2 * N + 1)\n\n N = 6\n for test_type in (Float32, Float64, BigFloat)\n r, w = Elements.lglpoints(test_type, N)\n D = Elements.spectralderivative(r)\n x = LinRange{test_type}(-1, 1, 101)\n I = Elements.interpolationmatrix(r, x)\n\n @test sum(P5(r) .^ 2 .* w) ≈ IPN(test_type, 5)\n @test D * P6(r) ≈ DP6(r)\n @test I * P6(r) ≈ P6(x)\n end\n\n for test_type in (Float32, Float64, BigFloat)\n r, w = Elements.glpoints(test_type, N)\n D = Elements.spectralderivative(r)\n\n @test sum(P5(r) .^ 2 .* w) ≈ IPN(test_type, 5)\n @test sum(P6(r) .^ 2 .* w) ≈ IPN(test_type, 6)\n @test D * P6(r) ≈ DP6(r)\n end\nend\n\n@testset \"Jacobip\" begin\n for T in (Float32, Float64, BigFloat)\n let\n α, β, N = T(0), T(0), 3 # α, β (for Legendre polynomials) & polynomial order\n x, wt = Elements.lglpoints(T, N + 1) # lgl points for polynomial order N\n V = Elements.jacobip(α, β, N, x)\n # compare with orthonormalized exact solution\n # https://en.wikipedia.org/wiki/Legendre_polynomials\n V_exact = similar(V)\n V_exact[:, 1] .= 1\n V_exact[:, 2] .= x\n V_exact[:, 3] .= (3 * x .^ 2 .- 1) / 2\n V_exact[:, 4] .= (5 * x .^ 3 .- 3 * x) / 2\n scale = 1 ./ sqrt.(diag(V_exact' * Diagonal(wt) * V_exact))\n V_exact = V_exact * Diagonal(scale)\n @test V ≈ V_exact\n end\n end\nend\n" }, { "path": "test/Numerics/Mesh/Geometry.jl", "content": "using Test, MPI\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Geometry\nusing StaticArrays\n\nMPI.Initialized() || MPI.Init()\n\n@testset \"LocalGeometry\" begin\n FT = Float64\n ArrayType = Array\n\n xmin = 0\n ymin = 0\n zmin = 0\n xmax = 2000\n ymax = 400\n zmax = 2000\n\n Ne = (20, 2, 20)\n\n polynomialorder = 4\n\n brickrange = (\n range(FT(xmin); length = Ne[1] + 1, stop = xmax),\n range(FT(ymin); length = Ne[2] + 1, stop = ymax),\n range(FT(zmin); length = Ne[3] + 1, stop = zmax),\n )\n topl = StackedBrickTopology(\n MPI.COMM_SELF,\n brickrange,\n periodicity = (true, true, false),\n )\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = polynomialorder,\n )\n\n S = (\n (xmax - xmin) / (polynomialorder * Ne[1]),\n (ymax - ymin) / (polynomialorder * Ne[2]),\n (zmax - zmin) / (polynomialorder * Ne[3]),\n )\n Savg = cbrt(prod(S))\n Shoriavg = (S[1] + S[2]) / 2\n M = SDiagonal(S .^ -2)\n\n N = (polynomialorder, polynomialorder)\n Np = (polynomialorder + 1)^3\n for e in 1:size(grid.vgeo, 3)\n for n in 1:size(grid.vgeo, 1)\n g = LocalGeometry{Np, N}(grid.vgeo, n, e)\n @test lengthscale(g) ≈ Savg\n @test lengthscale_horizontal(g) ≈ Shoriavg\n end\n end\nend\n" }, { "path": "test/Numerics/Mesh/Grids.jl", "content": "using ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Topologies: BrickTopology\nusing Test\nusing MPI\n\nMPI.Initialized() || MPI.Init()\n\n@testset \"2-D Mass matrix\" begin\n topology = BrickTopology(MPI.COMM_WORLD, ntuple(_ -> (-1, 1), 2);)\n\n @testset for N in ((2, 2), (2, 3), (3, 2))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = Float64,\n DeviceArray = Array,\n polynomialorder = N,\n )\n\n @views begin\n ω = grid.ω\n M = grid.vgeo[:, Grids._M, 1]\n MH = grid.vgeo[:, Grids._MH, 1]\n sM = grid.sgeo[Grids._sM, :, :, 1]\n\n M = reshape(M, Nq[1], Nq[2])\n @test M ≈ [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]\n\n MH = reshape(MH, Nq[1], Nq[2])\n @test MH ≈ [ω[1][i] for i in 1:Nq[1], j in 1:Nq[2]]\n\n sM12 = sM[1:Nfp[1], 1:2]\n @test sM12[:, 1] ≈ [ω[2][j] for j in 1:Nq[2]]\n @test sM12[:, 2] ≈ [ω[2][j] for j in 1:Nq[2]]\n\n sM34 = sM[1:Nfp[2], 3:4]\n @test sM34[:, 1] ≈ [ω[1][j] for j in 1:Nq[1]]\n @test sM34[:, 2] ≈ [ω[1][j] for j in 1:Nq[1]]\n end\n end\nend\n\n@testset \"3-D Mass matrix\" begin\n topology = BrickTopology(MPI.COMM_WORLD, ntuple(_ -> (-1, 1), 3);)\n\n @testset for N in ((2, 2, 2), (2, 3, 4), (4, 3, 2), (2, 4, 3))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = Float64,\n DeviceArray = Array,\n polynomialorder = N,\n )\n\n @views begin\n ω = grid.ω\n M = grid.vgeo[:, Grids._M, 1]\n MH = grid.vgeo[:, Grids._MH, 1]\n sM = grid.sgeo[Grids._sM, :, :, 1]\n\n M = reshape(M, Nq[1], Nq[2], Nq[3])\n @test M ≈ [\n ω[1][i] * ω[2][j] * ω[3][k]\n for i in 1:Nq[1], j in 1:Nq[2], k in 1:Nq[3]\n ]\n\n MH = reshape(MH, Nq[1], Nq[2], Nq[3])\n @test MH ≈ [\n ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2], k in 1:Nq[3]\n ]\n\n sM12 = reshape(sM[1:Nfp[1], 1:2], Nq[2], Nq[3], 2)\n @test sM12[:, :, 1] ≈\n [ω[2][j] * ω[3][k] for j in 1:Nq[2], k in 1:Nq[3]]\n @test sM12[:, :, 2] ≈\n [ω[2][j] * ω[3][k] for j in 1:Nq[2], k in 1:Nq[3]]\n\n sM34 = reshape(sM[1:Nfp[2], 3:4], Nq[1], Nq[3], 2)\n @test sM34[:, :, 1] ≈\n [ω[1][i] * ω[3][k] for i in 1:Nq[1], k in 1:Nq[3]]\n @test sM34[:, :, 2] ≈\n [ω[1][i] * ω[3][k] for i in 1:Nq[1], k in 1:Nq[3]]\n\n sM56 = reshape(sM[1:Nfp[3], 5:6], Nq[1], Nq[2], 2)\n @test sM56[:, :, 1] ≈\n [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]\n @test sM56[:, :, 2] ≈\n [ω[1][i] * ω[2][j] for i in 1:Nq[1], j in 1:Nq[2]]\n end\n end\nend\n" }, { "path": "test/Numerics/Mesh/Metrics.jl", "content": "using ClimateMachine.Mesh.Elements\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.GeometricFactors\nusing ClimateMachine.Mesh.Metrics\nusing LinearAlgebra: I\nusing Test\nusing Random: MersenneTwister\n\nconst _ξ1x1, _ξ2x1, _ξ3x1 = Grids._ξ1x1, Grids._ξ2x1, Grids._ξ3x1\nconst _ξ1x2, _ξ2x2, _ξ3x2 = Grids._ξ1x2, Grids._ξ2x2, Grids._ξ3x2\nconst _ξ1x3, _ξ2x3, _ξ3x3 = Grids._ξ1x3, Grids._ξ2x3, Grids._ξ3x3\nconst _M, _MI = Grids._M, Grids._MI\nconst _x1, _x2, _x3 = Grids._x1, Grids._x2, Grids._x3\nconst _JcV = Grids._JcV\nconst _nvgeo = Grids._nvgeo\n\nconst _n1, _n2, _n3 = Grids._n1, Grids._n2, Grids._n3\nconst _sM, _vMI = Grids._sM, Grids._vMI\nconst _nsgeo = Grids._nsgeo\n\n@testset \"1-D Metric terms\" begin\n for FT in (Float32, Float64)\n #{{{\n let\n N = (4,)\n Nq = N .+ 1\n Np = prod(Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.lglpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n dim = 1\n e2c = Array{FT, 3}(undef, 1, 2, 2)\n e2c[:, :, 1] = [-1 0]\n e2c[:, :, 2] = [0 10]\n nelem = size(e2c, 3)\n\n (vgeo, sgeo, _) =\n Grids.computegeometry(e2c, D, ξ, ω, (x...) -> identity(x))\n\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n @test vgeo[:, _x1, 1] ≈ (ξ[1] .- 1) / 2\n @test vgeo[:, _x1, 2] ≈ 5 * (ξ[1] .+ 1)\n\n @test vgeo[:, _M, 1] ≈ ω[1] .* ones(FT, Nq) / 2\n @test vgeo[:, _M, 2] ≈ 5 * ω[1] .* ones(FT, Nq)\n @test vgeo[:, _ξ1x1, 1] ≈ 2 * ones(FT, Nq)\n @test vgeo[:, _ξ1x1, 2] ≈ ones(FT, Nq) / 5\n @test sgeo.n1[1, 1, :] ≈ -ones(FT, nelem)\n @test sgeo.n1[1, 2, :] ≈ ones(FT, nelem)\n @test sgeo.sωJ[1, 1, :] ≈ ones(FT, nelem)\n @test sgeo.sωJ[1, 2, :] ≈ ones(FT, nelem)\n end\n #}}}\n end\n\n # N = 0 test\n for FT in (Float32, Float64)\n #{{{\n let\n N = (0,)\n Nq = N .+ 1\n Np = prod(Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.glpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n dim = 1\n e2c = Array{FT, 3}(undef, 1, 2, 2)\n e2c[:, :, 1] = [-1 0]\n e2c[:, :, 2] = [0 10]\n nelem = size(e2c, 3)\n\n (vgeo, sgeo, _) =\n Grids.computegeometry(e2c, D, ξ, ω, (x...) -> identity(x))\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n @test vgeo[1, _x1, 1] ≈ sum(e2c[:, :, 1]) / 2\n @test vgeo[1, _x1, 2] ≈ sum(e2c[:, :, 2]) / 2\n\n @test vgeo[:, _M, 1] ≈ ω[1] .* ones(FT, Nq) / 2\n @test vgeo[:, _M, 2] ≈ 5 * ω[1] .* ones(FT, Nq)\n @test vgeo[:, _ξ1x1, 1] ≈ 2 * ones(FT, Nq)\n @test vgeo[:, _ξ1x1, 2] ≈ ones(FT, Nq) / 5\n\n @test sgeo.n1[1, 1, :] ≈ -ones(FT, nelem)\n @test sgeo.n1[1, 2, :] ≈ ones(FT, nelem)\n @test sgeo.sωJ[1, 1, :] ≈ ones(FT, nelem)\n @test sgeo.sωJ[1, 2, :] ≈ ones(FT, nelem)\n end\n #}}}\n end\nend\n\n@testset \"2-D Metric terms\" begin\n for FT in (Float32, Float64), N in ((4, 4), (4, 6), (6, 4))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.lglpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n # linear and rotation test\n #{{{\n let\n e2c = Array{FT, 3}(undef, 2, 4, 4)\n e2c[:, :, 1] = [\n 0 2 0 2\n 0 0 2 2\n ]\n e2c[:, :, 2] = [\n 2 2 0 0\n 0 2 0 2\n ]\n e2c[:, :, 3] = [\n 2 0 2 0\n 2 2 0 0\n ]\n e2c[:, :, 4] = [\n 0 0 2 2\n 2 0 2 0\n ]\n nelem = size(e2c, 3)\n\n x_exact = Array{FT, 3}(undef, Nq..., nelem)\n x_exact[:, :, 1] .= 1 .+ ξ[1]\n x_exact[:, :, 2] .= 1 .- ξ[2]'\n x_exact[:, :, 3] .= 1 .- ξ[1]\n x_exact[:, :, 4] .= 1 .+ ξ[2]'\n\n y_exact = Array{FT, 3}(undef, Nq..., nelem)\n y_exact[:, :, 1] .= 1 .+ ξ[2]'\n y_exact[:, :, 2] .= 1 .+ ξ[1]\n y_exact[:, :, 3] .= 1 .- ξ[2]'\n y_exact[:, :, 4] .= 1 .- ξ[1]\n\n M_exact =\n ones(FT, Nq..., nelem) .* reshape(kron(reverse(ω)...), Nq..., 1)\n\n ξ1x1_exact = zeros(FT, Nq..., nelem)\n ξ1x1_exact[:, :, 1] .= 1\n ξ1x1_exact[:, :, 3] .= -1\n\n ξ1x2_exact = zeros(FT, Nq..., nelem)\n ξ1x2_exact[:, :, 2] .= 1\n ξ1x2_exact[:, :, 4] .= -1\n\n ξ2x1_exact = zeros(FT, Nq..., nelem)\n ξ2x1_exact[:, :, 2] .= -1\n ξ2x1_exact[:, :, 4] .= 1\n\n ξ2x2_exact = zeros(FT, Nq..., nelem)\n ξ2x2_exact[:, :, 1] .= 1\n ξ2x2_exact[:, :, 3] .= -1\n\n sM_exact = fill(FT(NaN), maximum(Nfp), nface, nelem)\n sM_exact[1:Nfp[1], 1, :] .= 1 .* ω[2]\n sM_exact[1:Nfp[1], 2, :] .= 1 .* ω[2]\n sM_exact[1:Nfp[2], 3, :] .= 1 .* ω[1]\n sM_exact[1:Nfp[2], 4, :] .= 1 .* ω[1]\n\n nx_exact = fill(FT(NaN), maximum(Nfp), nface, nelem)\n nx_exact[1:Nfp[1], 1:2, :] .= 0\n nx_exact[1:Nfp[2], 3:4, :] .= 0\n\n nx_exact[1:Nfp[1], 1, 1] .= -1\n nx_exact[1:Nfp[1], 2, 1] .= 1\n nx_exact[1:Nfp[2], 3, 2] .= 1\n nx_exact[1:Nfp[2], 4, 2] .= -1\n nx_exact[1:Nfp[1], 1, 3] .= 1\n nx_exact[1:Nfp[1], 2, 3] .= -1\n nx_exact[1:Nfp[2], 3, 4] .= -1\n nx_exact[1:Nfp[2], 4, 4] .= 1\n\n ny_exact = fill(FT(NaN), maximum(Nfp), nface, nelem)\n ny_exact[1:Nfp[1], 1:2, :] .= 0\n ny_exact[1:Nfp[2], 3:4, :] .= 0\n\n ny_exact[1:Nfp[2], 3, 1] .= -1\n ny_exact[1:Nfp[2], 4, 1] .= 1\n ny_exact[1:Nfp[1], 1, 2] .= -1\n ny_exact[1:Nfp[1], 2, 2] .= 1\n ny_exact[1:Nfp[2], 3, 3] .= 1\n ny_exact[1:Nfp[2], 4, 3] .= -1\n ny_exact[1:Nfp[1], 1, 4] .= 1\n ny_exact[1:Nfp[1], 2, 4] .= -1\n\n (vgeo, sgeo, _) =\n Grids.computegeometry(e2c, D, ξ, ω, (x...) -> identity(x))\n\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n @test (@view vgeo[:, :, _x1, :]) ≈ x_exact\n @test (@view vgeo[:, :, _x2, :]) ≈ y_exact\n @test (@view vgeo[:, :, _M, :]) ≈ M_exact\n @test (@view vgeo[:, :, _ξ1x1, :]) ≈ ξ1x1_exact\n @test (@view vgeo[:, :, _ξ1x2, :]) ≈ ξ1x2_exact\n @test (@view vgeo[:, :, _ξ2x1, :]) ≈ ξ2x1_exact\n @test (@view vgeo[:, :, _ξ2x2, :]) ≈ ξ2x2_exact\n msk = isfinite.(sM_exact)\n @test sgeo.sωJ[:, :, :][msk] ≈ sM_exact[msk]\n @test sgeo.n1[:, :, :][msk] ≈ nx_exact[msk]\n @test sgeo.n2[:, :, :][msk] ≈ ny_exact[msk]\n\n nothing\n end\n #}}}\n\n # Polynomial 2-D test\n #{{{\n let\n f(ξ1, ξ2) = (\n 9 .* ξ1 - (1 .+ ξ1) .* ξ2 .^ 2 +\n (ξ1 .- 1) .^ 2 .* (1 .- ξ2 .^ 2 .+ ξ2 .^ 3),\n 10 .* ξ2 .+ ξ1 .^ 4 .* (1 .- ξ2) .+ ξ1 .^ 2 .* ξ2 .* (1 .+ ξ2),\n )\n fx1ξ1(ξ1, ξ2) =\n 7 .+ ξ2 .^ 2 .- 2 .* ξ2 .^ 3 .+\n 2 .* ξ1 .* (1 .- ξ2 .^ 2 .+ ξ2 .^ 3)\n fx1ξ2(ξ1, ξ2) =\n -2 .* (1 .+ ξ1) .* ξ2 .+\n (-1 .+ ξ1) .^ 2 .* ξ2 .* (-2 .+ 3 .* ξ2)\n fx2ξ1(ξ1, ξ2) =\n -4 .* ξ1 .^ 3 .* (-1 .+ ξ2) .+ 2 .* ξ1 .* ξ2 .* (1 .+ ξ2)\n fx2ξ2(ξ1, ξ2) = 10 .- ξ1 .^ 4 .+ ξ1 .^ 2 .* (1 .+ 2 .* ξ2)\n\n e2c = Array{FT, 3}(undef, 2, 4, 1)\n e2c[:, :, 1] = [-1 1 -1 1; -1 -1 1 1]\n nelem = size(e2c, 3)\n\n # Create the metrics\n (x1ξ1, x1ξ2, x2ξ1, x2ξ2) = let\n (vgeo, _) = Grids.computegeometry(\n e2c,\n D,\n ξ,\n ω,\n (x...) -> identity(x),\n )\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n ξ1, ξ2 = vgeo[:, :, _x1, :], vgeo[:, :, _x2, :]\n (fx1ξ1(ξ1, ξ2), fx1ξ2(ξ1, ξ2), fx2ξ1(ξ1, ξ2), fx2ξ2(ξ1, ξ2))\n end\n J = (x1ξ1 .* x2ξ2 - x1ξ2 .* x2ξ1)\n M = J .* reshape(kron(reverse(ω)...), Nq..., 1)\n\n meshwarp(ξ1, ξ2, _) = (f(ξ1, ξ2)..., 0)\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n x1 = @view vgeo[:, :, _x1, :]\n x2 = @view vgeo[:, :, _x2, :]\n\n @test M ≈ (@view vgeo[:, :, _M, :])\n @test (@view vgeo[:, :, _ξ1x1, :]) ≈ x2ξ2 ./ J\n @test (@view vgeo[:, :, _ξ2x1, :]) ≈ -x2ξ1 ./ J\n @test (@view vgeo[:, :, _ξ1x2, :]) ≈ -x1ξ2 ./ J\n @test (@view vgeo[:, :, _ξ2x2, :]) ≈ x1ξ1 ./ J\n\n # check the normals?\n sM = @view sgeo.sωJ[:, :, :]\n n1 = @view sgeo.n1[:, :, :]\n n2 = @view sgeo.n2[:, :, :]\n @test all(hypot.(n1[1:Nfp[1], 1:2, :], n2[1:Nfp[1], 1:2, :]) .≈ 1)\n @test all(hypot.(n1[1:Nfp[2], 3:4, :], n2[1:Nfp[2], 3:4, :]) .≈ 1)\n @test sM[1:Nfp[1], 1, :] .* n1[1:Nfp[1], 1, :] ≈\n -x2ξ2[1, :, :] .* ω[2]\n @test sM[1:Nfp[1], 1, :] .* n2[1:Nfp[1], 1, :] ≈\n x1ξ2[1, :, :] .* ω[2]\n @test sM[1:Nfp[1], 2, :] .* n1[1:Nfp[1], 2, :] ≈\n x2ξ2[Nq[1], :, :] .* ω[2]\n @test sM[1:Nfp[1], 2, :] .* n2[1:Nfp[1], 2, :] ≈\n -x1ξ2[Nq[1], :, :] .* ω[2]\n @test sM[1:Nfp[2], 3, :] .* n1[1:Nfp[2], 3, :] ≈\n x2ξ1[:, 1, :] .* ω[1]\n @test sM[1:Nfp[2], 3, :] .* n2[1:Nfp[2], 3, :] ≈\n -x1ξ1[:, 1, :] .* ω[1]\n @test sM[1:Nfp[2], 4, :] .* n1[1:Nfp[2], 4, :] ≈\n -x2ξ1[:, Nq[2], :] .* ω[1]\n @test sM[1:Nfp[2], 4, :] .* n2[1:Nfp[2], 4, :] ≈\n x1ξ1[:, Nq[2], :] .* ω[1]\n end\n #}}}\n\n # Constant preserving test\n #{{{\n let\n rng = MersenneTwister(777)\n f(ξ1, ξ2) = (\n ξ1 + (ξ1 * ξ2 * rand(rng) + rand(rng)) / 10,\n ξ2 + (ξ1 * ξ2 * rand(rng) + rand(rng)) / 10,\n )\n\n e2c = Array{FT, 3}(undef, 2, 4, 1)\n e2c[:, :, 1] = [-1 1 -1 1; -1 -1 1 1]\n nelem = size(e2c, 3)\n\n meshwarp(ξ1, ξ2, _) = (f(ξ1, ξ2)..., 0)\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n x1 = @view vgeo[:, :, _x1, :]\n x2 = @view vgeo[:, :, _x2, :]\n\n (Cx1, Cx2) = (zeros(FT, Nq...), zeros(FT, Nq...))\n\n J = vgeo[:, :, _M, :] ./ reshape(kron(reverse(ω)...), Nq..., 1)\n ξ1x1 = @view vgeo[:, :, _ξ1x1, :]\n ξ1x2 = @view vgeo[:, :, _ξ1x2, :]\n ξ2x1 = @view vgeo[:, :, _ξ2x1, :]\n ξ2x2 = @view vgeo[:, :, _ξ2x2, :]\n\n e = 1\n for n in 1:Nq[2]\n Cx1[:, n] += D[1] * (J[:, n, e] .* ξ1x1[:, n, e])\n Cx2[:, n] += D[1] * (J[:, n, e] .* ξ1x2[:, n, e])\n end\n for n in 1:Nq[1]\n Cx1[n, :] += D[2] * (J[n, :, e] .* ξ2x1[n, :, e])\n Cx2[n, :] += D[2] * (J[n, :, e] .* ξ2x2[n, :, e])\n end\n @test maximum(abs.(Cx1)) ≤ 100 * eps(FT)\n @test maximum(abs.(Cx2)) ≤ 100 * eps(FT)\n end\n #}}}\n end\n\n #N = 0 test\n #{{{\n let\n for FT in (Float32, Float64)\n N = (2, 0)\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n Nq[j] == 1 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n fx1(ξ1, ξ2) = ξ1 + (1 + ξ1)^2 * ξ2 / 10\n fx1ξ1(ξ1, ξ2) = 1 + 2 * (1 + ξ1) * ξ2 / 10\n fx1ξ2(ξ1, ξ2) = (1 + ξ1)^2 / 10\n fx2(ξ1, ξ2) = ξ2 - (1 + ξ1)^2\n fx2ξ1(ξ1, ξ2) = -2 * (1 + ξ1)\n fx2ξ2(ξ1, ξ2) = 1\n\n e2c = Array{FT, 3}(undef, 2, 4, 1)\n e2c[:, :, 1] = [-1 1 -1 1; -1 -1 1 1]\n nelem = size(e2c, 3)\n\n # Create the metrics\n (x1, x2, x1ξ1, x1ξ2, x2ξ1, x2ξ2) = let\n vgeo = VolumeGeometry(FT, Nq, nelem)\n Metrics.creategrid!(vgeo, e2c, ξ)\n (\n fx1.(vgeo.x1, vgeo.x2),\n fx2.(vgeo.x1, vgeo.x2),\n fx1ξ1.(vgeo.x1, vgeo.x2),\n fx1ξ2.(vgeo.x1, vgeo.x2),\n fx2ξ1.(vgeo.x1, vgeo.x2),\n fx2ξ2.(vgeo.x1, vgeo.x2),\n )\n end\n J = (x1ξ1 .* x2ξ2 - x1ξ2 .* x2ξ1)\n\n M = J .* reshape(kron(reverse(ω)...), Nq..., 1)\n\n meshwarp(ξ1, ξ2, _) = (fx1(ξ1, ξ2), fx2(ξ1, ξ2), 0)\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n @test x1 ≈ vgeo[:, :, _x1, :]\n @test x2 ≈ vgeo[:, :, _x2, :]\n\n @test M ≈ vgeo[:, :, _M, :]\n @test (@view vgeo[:, :, _ξ2x1, :]) .* J ≈ -x2ξ1\n @test (@view vgeo[:, :, _ξ2x2, :]) .* J ≈ x1ξ1\n @test (@view vgeo[:, :, _ξ1x1, :]) .* J ≈ x2ξ2\n @test (@view vgeo[:, :, _ξ1x2, :]) .* J ≈ -x1ξ2\n\n # check the normals?\n sM = @view sgeo.sωJ[:, :, :]\n n1 = @view sgeo.n1[:, :, :]\n n2 = @view sgeo.n2[:, :, :]\n @test all(hypot.(n1[1:Nfp[1], 1:2, :], n2[1:Nfp[1], 1:2, :]) .≈ 1)\n @test all(hypot.(n1[1:Nfp[2], 3:4, :], n2[1:Nfp[2], 3:4, :]) .≈ 1)\n @test sM[1:Nfp[1], 1, :] .* n1[1:Nfp[1], 1, :] ≈\n -x2ξ2[1, :, :] .* ω[2]\n @test sM[1:Nfp[1], 1, :] .* n2[1:Nfp[1], 1, :] ≈\n x1ξ2[1, :, :] .* ω[2]\n @test sM[1:Nfp[1], 2, :] .* n1[1:Nfp[1], 2, :] ≈\n x2ξ2[Nq[1], :, :] .* ω[2]\n @test sM[1:Nfp[1], 2, :] .* n2[1:Nfp[1], 2, :] ≈\n -x1ξ2[Nq[1], :, :] .* ω[2]\n\n # for these faces we need the N = 1 metrics\n (x1ξ1, x2ξ1) = let\n @assert Nq[2] == 1 && Nq[1] != 1\n Nq_N1 = max.(2, Nq)\n vgeo_N1 = VolumeGeometry(FT, Nq_N1, nelem)\n Metrics.creategrid!(\n vgeo_N1,\n e2c,\n (ξ[1], Elements.lglpoints(FT, 1)[1]),\n )\n x1 = reshape(vgeo_N1.x1, (Nq_N1..., nelem))\n x2 = reshape(vgeo_N1.x2, (Nq_N1..., nelem))\n (fx1ξ1.(x1, x2), fx2ξ1.(x1, x2))\n end\n\n @test sM[1:Nfp[2], 3, :] .* n1[1:Nfp[2], 3, :] ≈\n x2ξ1[:, 1, :] .* ω[1]\n @test sM[1:Nfp[2], 3, :] .* n2[1:Nfp[2], 3, :] ≈\n -x1ξ1[:, 1, :] .* ω[1]\n @test sM[1:Nfp[2], 4, :] .* n1[1:Nfp[2], 4, :] ≈\n -x2ξ1[:, 2, :] .* ω[1]\n @test sM[1:Nfp[2], 4, :] .* n2[1:Nfp[2], 4, :] ≈\n x1ξ1[:, 2, :] .* ω[1]\n end\n end\n #}}}\n\n # Constant preserving test for N = 0\n #{{{\n let\n for FT in (Float32, Float64), N in ((4, 0), (0, 4))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n Nq[j] == 1 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n rng = MersenneTwister(777)\n fx1(ξ1, ξ2) = ξ1 + (ξ1 * ξ2 * rand(rng) + rand(rng)) / 10\n fx2(ξ1, ξ2) = ξ2 + (ξ1 * ξ2 * rand(rng) + rand(rng)) / 10\n\n e2c = Array{FT, 3}(undef, 2, 4, 1)\n e2c[:, :, 1] = [-1 1 -1 1; -1 -1 1 1]\n nelem = size(e2c, 3)\n\n meshwarp(ξ1, ξ2, _) = (fx1(ξ1, ξ2), fx2(ξ1, ξ2), 0)\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n\n vgeo = reshape(\n vgeo.array,\n prod(Nq),\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n M = vgeo[:, _M, :]\n ξ1x1 = vgeo[:, _ξ1x1, :]\n ξ2x1 = vgeo[:, _ξ2x1, :]\n ξ1x2 = vgeo[:, _ξ1x2, :]\n ξ2x2 = vgeo[:, _ξ2x2, :]\n\n I1 = Matrix(I, Nq[1], Nq[1])\n I2 = Matrix(I, Nq[2], Nq[2])\n D1 = kron(I2, D[1])\n D2 = kron(D[2], I1)\n\n # Face interpolation operators\n L = (\n kron(I2, I1[1, :]'),\n kron(I2, I1[Nq[1], :]'),\n kron(I2[1, :]', I1),\n kron(I2[Nq[2], :]', I1),\n )\n sM = ntuple(f -> sgeo.sωJ[1:Nfp[cld(f, 2)], f, :], nface)\n n1 = ntuple(f -> sgeo.n1[1:Nfp[cld(f, 2)], f, :], nface)\n n2 = ntuple(f -> sgeo.n2[1:Nfp[cld(f, 2)], f, :], nface)\n\n # If constant preserving then:\n # \\sum_{j} = D' * M * ξjxk = \\sum_{f} L_f' * sM_f * n1_f\n @test D1' * (M .* ξ1x1) + D2' * (M .* ξ2x1) ≈\n mapreduce((L, sM, n1) -> L' * (sM .* n1), +, L, sM, n1)\n @test D1' * (M .* ξ1x2) + D2' * (M .* ξ2x2) ≈\n mapreduce((L, sM, n2) -> L' * (sM .* n2), +, L, sM, n2)\n end\n end\n #}}}\nend\n\n@testset \"3-D Metric terms\" begin\n # linear test\n #{{{\n for FT in (Float32, Float64), N in ((2, 2, 2), (2, 3, 4), (4, 3, 2))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.lglpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n e2c = Array{FT, 3}(undef, dim, 8, 2)\n e2c[:, :, 1] = [\n 0 2 0 2 0 2 0 2\n 0 0 2 2 0 0 2 2\n 0 0 0 0 2 2 2 2\n ]\n e2c[:, :, 2] = [\n 2 2 0 0 2 2 0 0\n 0 2 0 2 0 2 0 2\n 0 0 0 0 2 2 2 2\n ]\n\n nelem = size(e2c, 3)\n\n x_exact = Array{FT, 4}(undef, Nq..., nelem)\n x_exact[:, :, :, 1] .= 1 .+ ξ[1]\n x_exact[:, :, :, 2] .= 1 .- ξ[2]'\n\n ξ1x1_exact = zeros(Int, Nq..., nelem)\n ξ1x1_exact[:, :, :, 1] .= 1\n\n ξ1x2_exact = zeros(Int, Nq..., nelem)\n ξ1x2_exact[:, :, :, 2] .= 1\n\n ξ2x1_exact = zeros(Int, Nq..., nelem)\n ξ2x1_exact[:, :, :, 2] .= -1\n\n ξ2x2_exact = zeros(Int, Nq..., nelem)\n ξ2x2_exact[:, :, :, 1] .= 1\n\n ξ3x3_exact = ones(Int, Nq..., nelem)\n\n y_exact = Array{FT, 4}(undef, Nq..., nelem)\n y_exact[:, :, :, 1] .= 1 .+ ξ[2]'\n y_exact[:, :, :, 2] .= 1 .+ ξ[1]\n\n z_exact = Array{FT, 4}(undef, Nq..., nelem)\n z_exact[:, :, :, :] .= reshape(1 .+ ξ[3], 1, 1, Nq[3])\n\n M_exact =\n ones(Int, Nq..., nelem) .* reshape(kron(reverse(ω)...), Nq..., 1)\n\n sJ_exact = ones(Int, maximum(Nfp), nface, nelem)\n\n nx_exact = zeros(Int, maximum(Nfp), nface, nelem)\n nx_exact[:, 1, 1] .= -1\n nx_exact[:, 2, 1] .= 1\n nx_exact[:, 3, 2] .= 1\n nx_exact[:, 4, 2] .= -1\n\n ny_exact = zeros(Int, maximum(Nfp), nface, nelem)\n ny_exact[:, 3, 1] .= -1\n ny_exact[:, 4, 1] .= 1\n ny_exact[:, 1, 2] .= -1\n ny_exact[:, 2, 2] .= 1\n\n nz_exact = zeros(Int, maximum(Nfp), nface, nelem)\n nz_exact[:, 5, 1:2] .= -1\n nz_exact[:, 6, 1:2] .= 1\n\n (vgeo, sgeo, _) =\n Grids.computegeometry(e2c, D, ξ, ω, (x...) -> identity(x))\n\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n @test (@view vgeo[:, :, :, _x1, :]) ≈ x_exact\n @test (@view vgeo[:, :, :, _x2, :]) ≈ y_exact\n @test (@view vgeo[:, :, :, _x3, :]) ≈ z_exact\n @test (@view vgeo[:, :, :, _M, :]) ≈ M_exact\n @test (@view vgeo[:, :, :, _ξ1x1, :]) ≈ ξ1x1_exact\n @test (@view vgeo[:, :, :, _ξ1x2, :]) ≈ ξ1x2_exact\n @test maximum(abs.(@view vgeo[:, :, :, _ξ1x3, :])) ≤ 100 * eps(FT)\n @test (@view vgeo[:, :, :, _ξ2x1, :]) ≈ ξ2x1_exact\n @test (@view vgeo[:, :, :, _ξ2x2, :]) ≈ ξ2x2_exact\n @test maximum(abs.(@view vgeo[:, :, :, _ξ2x3, :])) ≤ 100 * eps(FT)\n @test maximum(abs.(@view vgeo[:, :, :, _ξ3x1, :])) ≤ 100 * eps(FT)\n @test maximum(abs.(@view vgeo[:, :, :, _ξ3x2, :])) ≤ 100 * eps(FT)\n @test (@view vgeo[:, :, :, _ξ3x3, :]) ≈ ξ3x3_exact\n for d in 1:dim\n for f in (2d - 1):(2d)\n ωf = ntuple(j -> ω[mod1(d + j, dim)], dim - 1)\n if !(dim == 3 && d == 2)\n ωf = reverse(ωf)\n end\n Mf = kron(1, ωf...)\n @test isapprox(\n (@view sgeo.sωJ[1:Nfp[d], f, :]),\n sJ_exact[1:Nfp[d], f, :] .* Mf,\n atol = √eps(FT),\n rtol = √eps(FT),\n )\n @test isapprox(\n (@view sgeo.n1[1:Nfp[d], f, :]),\n nx_exact[1:Nfp[d], f, :];\n atol = √eps(FT),\n rtol = √eps(FT),\n )\n @test isapprox(\n (@view sgeo.n2[1:Nfp[d], f, :]),\n ny_exact[1:Nfp[d], f, :];\n atol = √eps(FT),\n rtol = √eps(FT),\n )\n @test isapprox(\n (@view sgeo.n3[1:Nfp[d], f, :]),\n nz_exact[1:Nfp[d], f, :];\n atol = √eps(FT),\n rtol = √eps(FT),\n )\n end\n end\n end\n #}}}\n\n # Polynomial 3-D test\n #{{{\n for FT in (Float32, Float64), N in ((9, 9, 9), (9, 9, 10), (10, 9, 11))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.lglpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n f(ξ1, ξ2, ξ3) = @.( (\n ξ2 + ξ1 * ξ3 - (ξ1^2 * ξ2^2 * ξ3^2) / 4,\n ξ3 - ((ξ1 * ξ2 * ξ3 + 1) / 2)^3 + 1,\n ξ1 + 2 * ((ξ1 + 1) / 2)^6 * ((ξ2 + 1) / 2)^6 * ((ξ3 + 1) / 2)^6,\n ))\n\n fx1ξ1(ξ1, ξ2, ξ3) = @.(ξ3 - (ξ1 * ξ2^2 * ξ3^2) / 2)\n fx1ξ2(ξ1, ξ2, ξ3) = @.(1 - (ξ1^2 * ξ2 * ξ3^2) / 2)\n fx1ξ3(ξ1, ξ2, ξ3) = @.(ξ1 - (ξ1^2 * ξ2^2 * ξ3) / 2)\n fx2ξ1(ξ1, ξ2, ξ3) = @.(-(3 * ξ2 * ξ3 * ((ξ1 * ξ2 * ξ3 + 1) / 2)^2) / 2)\n fx2ξ2(ξ1, ξ2, ξ3) = @.(-(3 * ξ1 * ξ3 * ((ξ1 * ξ2 * ξ3 + 1) / 2)^2) / 2)\n fx2ξ3(ξ1, ξ2, ξ3) =\n @.(1 - (3 * ξ1 * ξ2 * ((ξ1 * ξ2 * ξ3 + 1) / 2)^2) / 2)\n fx3ξ1(ξ1, ξ2, ξ3) =\n @.(6 * ((ξ1 + 1) / 2)^5 * ((ξ2 + 1) / 2)^6 * ((ξ3 + 1) / 2)^6 + 1)\n fx3ξ2(ξ1, ξ2, ξ3) =\n @.(6 * ((ξ1 + 1) / 2)^6 * ((ξ2 + 1) / 2)^5 * ((ξ3 + 1) / 2)^6)\n fx3ξ3(ξ1, ξ2, ξ3) =\n @.(6 * ((ξ1 + 1) / 2)^6 * ((ξ2 + 1) / 2)^6 * ((ξ3 + 1) / 2)^5)\n\n e2c = Array{FT, 3}(undef, dim, 8, 1)\n e2c[:, :, 1] = [\n -1 1 -1 1 -1 1 -1 1\n -1 -1 1 1 -1 -1 1 1\n -1 -1 -1 -1 1 1 1 1\n ]\n\n nelem = size(e2c, 3)\n\n # Compute exact metrics\n (x1ξ1, x1ξ2, x1ξ3, x2ξ1, x2ξ2, x2ξ3, x3ξ1, x3ξ2, x3ξ3) = let\n (vgeo, _) =\n Grids.computegeometry(e2c, D, ξ, ω, (x...) -> identity(x))\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n ξ1 = vgeo[:, :, :, _x1, :]\n ξ2 = vgeo[:, :, :, _x2, :]\n ξ3 = vgeo[:, :, :, _x3, :]\n (\n fx1ξ1(ξ1, ξ2, ξ3),\n fx1ξ2(ξ1, ξ2, ξ3),\n fx1ξ3(ξ1, ξ2, ξ3),\n fx2ξ1(ξ1, ξ2, ξ3),\n fx2ξ2(ξ1, ξ2, ξ3),\n fx2ξ3(ξ1, ξ2, ξ3),\n fx3ξ1(ξ1, ξ2, ξ3),\n fx3ξ2(ξ1, ξ2, ξ3),\n fx3ξ3(ξ1, ξ2, ξ3),\n )\n end\n J = (\n x1ξ1 .* (x2ξ2 .* x3ξ3 - x2ξ3 .* x3ξ2) +\n x2ξ1 .* (x3ξ2 .* x1ξ3 - x3ξ3 .* x1ξ2) +\n x3ξ1 .* (x1ξ2 .* x2ξ3 - x1ξ3 .* x2ξ2)\n )\n\n ξ1x1 = (x2ξ2 .* x3ξ3 - x2ξ3 .* x3ξ2) ./ J\n ξ1x2 = (x3ξ2 .* x1ξ3 - x3ξ3 .* x1ξ2) ./ J\n ξ1x3 = (x1ξ2 .* x2ξ3 - x1ξ3 .* x2ξ2) ./ J\n ξ2x1 = (x2ξ3 .* x3ξ1 - x2ξ1 .* x3ξ3) ./ J\n ξ2x2 = (x3ξ3 .* x1ξ1 - x3ξ1 .* x1ξ3) ./ J\n ξ2x3 = (x1ξ3 .* x2ξ1 - x1ξ1 .* x2ξ3) ./ J\n ξ3x1 = (x2ξ1 .* x3ξ2 - x2ξ2 .* x3ξ1) ./ J\n ξ3x2 = (x3ξ1 .* x1ξ2 - x3ξ2 .* x1ξ1) ./ J\n ξ3x3 = (x1ξ1 .* x2ξ2 - x1ξ2 .* x2ξ1) ./ J\n\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, f)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n @test (@view vgeo[:, :, :, _M, :]) ≈\n J .* reshape(kron(reverse(ω)...), Nq..., 1)\n @test (@view vgeo[:, :, :, _ξ1x1, :]) ≈ ξ1x1\n @test (@view vgeo[:, :, :, _ξ1x2, :]) ≈ ξ1x2\n @test (@view vgeo[:, :, :, _ξ1x3, :]) ≈ ξ1x3\n @test (@view vgeo[:, :, :, _ξ2x1, :]) ≈ ξ2x1\n @test (@view vgeo[:, :, :, _ξ2x2, :]) ≈ ξ2x2\n @test (@view vgeo[:, :, :, _ξ2x3, :]) ≈ ξ2x3\n @test (@view vgeo[:, :, :, _ξ3x1, :]) ≈ ξ3x1\n @test (@view vgeo[:, :, :, _ξ3x2, :]) ≈ ξ3x2\n @test (@view vgeo[:, :, :, _ξ3x3, :]) ≈ ξ3x3\n n1 = @view sgeo.n1[:, :, :]\n n2 = @view sgeo.n2[:, :, :]\n n3 = @view sgeo.n3[:, :, :]\n sM = @view sgeo.sωJ[:, :, :]\n for d in 1:dim\n for f in (2d - 1):(2d)\n @test all(\n hypot.(\n n1[1:Nfp[d], f, :],\n n2[1:Nfp[d], f, :],\n n3[1:Nfp[d], f, :],\n ) .≈ 1,\n )\n end\n end\n d, f = 1, 1\n Mf = kron(1, ω[3], ω[2])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[1, :, :, :] .* ξ1x1[1, :, :, :])[:] .* Mf,\n (-J[1, :, :, :] .* ξ1x2[1, :, :, :])[:] .* Mf,\n (-J[1, :, :, :] .* ξ1x3[1, :, :, :])[:] .* Mf,\n ]\n d, f = 1, 2\n Mf = kron(1, ω[3], ω[2])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[Nq[d], :, :, :] .* ξ1x1[Nq[d], :, :, :])[:] .* Mf,\n (J[Nq[d], :, :, :] .* ξ1x2[Nq[d], :, :, :])[:] .* Mf,\n (J[Nq[d], :, :, :] .* ξ1x3[Nq[d], :, :, :])[:] .* Mf,\n ]\n d, f = 2, 3\n Mf = kron(1, ω[3], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[:, 1, :, :] .* ξ2x1[:, 1, :, :])[:] .* Mf,\n (-J[:, 1, :, :] .* ξ2x2[:, 1, :, :])[:] .* Mf,\n (-J[:, 1, :, :] .* ξ2x3[:, 1, :, :])[:] .* Mf,\n ]\n d, f = 2, 4\n Mf = kron(1, ω[3], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[:, Nq[d], :, :] .* ξ2x1[:, Nq[d], :, :])[:] .* Mf,\n (J[:, Nq[d], :, :] .* ξ2x2[:, Nq[d], :, :])[:] .* Mf,\n (J[:, Nq[d], :, :] .* ξ2x3[:, Nq[d], :, :])[:] .* Mf,\n ]\n d, f = 3, 5\n Mf = kron(1, ω[2], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[:, :, 1, :] .* ξ3x1[:, :, 1, :])[:] .* Mf,\n (-J[:, :, 1, :] .* ξ3x2[:, :, 1, :])[:] .* Mf,\n (-J[:, :, 1, :] .* ξ3x3[:, :, 1, :])[:] .* Mf,\n ]\n d, f = 3, 6\n Mf = kron(1, ω[2], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[:, :, Nq[d], :] .* ξ3x1[:, :, Nq[d], :])[:] .* Mf,\n (J[:, :, Nq[d], :] .* ξ3x2[:, :, Nq[d], :])[:] .* Mf,\n (J[:, :, Nq[d], :] .* ξ3x3[:, :, Nq[d], :])[:] .* Mf,\n ]\n end\n #}}}\n\n # Constant preserving test\n #{{{\n for FT in (Float32, Float64), N in ((5, 5, 5), (3, 4, 5), (4, 4, 5))\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(j -> Elements.lglpoints(FT, N[j]), dim)\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n f(ξ1, ξ2, ξ3) = @.( (\n ξ2 + ξ1 * ξ3 - (ξ1^2 * ξ2^2 * ξ3^2) / 4,\n ξ3 - ((ξ1 * ξ2 * ξ3 + 1) / 2)^3 + 1,\n ξ1 + ((ξ1 + 1) / 2)^6 * ((ξ2 + 1) / 2)^6 * ((ξ3 + 1) / 2)^6,\n ))\n\n e2c = Array{FT, 3}(undef, dim, 8, 1)\n e2c[:, :, 1] = [\n -1 1 -1 1 -1 1 -1 1\n -1 -1 1 1 -1 -1 1 1\n -1 -1 -1 -1 1 1 1 1\n ]\n\n nelem = size(e2c, 3)\n\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, f)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n (Cx1, Cx2, Cx3) = (zeros(FT, Nq...), zeros(FT, Nq...), zeros(FT, Nq...))\n\n J =\n (@view vgeo[:, :, :, _M, :]) ./\n reshape(kron(reverse(ω)...), Nq..., 1)\n ξ1x1 = @view vgeo[:, :, :, _ξ1x1, :]\n ξ1x2 = @view vgeo[:, :, :, _ξ1x2, :]\n ξ1x3 = @view vgeo[:, :, :, _ξ1x3, :]\n ξ2x1 = @view vgeo[:, :, :, _ξ2x1, :]\n ξ2x2 = @view vgeo[:, :, :, _ξ2x2, :]\n ξ2x3 = @view vgeo[:, :, :, _ξ2x3, :]\n ξ3x1 = @view vgeo[:, :, :, _ξ3x1, :]\n ξ3x2 = @view vgeo[:, :, :, _ξ3x2, :]\n ξ3x3 = @view vgeo[:, :, :, _ξ3x3, :]\n\n e = 1\n for k in 1:Nq[3]\n for j in 1:Nq[2]\n Cx1[:, j, k] += D[1] * (J[:, j, k, e] .* ξ1x1[:, j, k, e])\n Cx2[:, j, k] += D[1] * (J[:, j, k, e] .* ξ1x2[:, j, k, e])\n Cx3[:, j, k] += D[1] * (J[:, j, k, e] .* ξ1x3[:, j, k, e])\n end\n end\n\n for k in 1:Nq[3]\n for i in 1:Nq[1]\n Cx1[i, :, k] += D[2] * (J[i, :, k, e] .* ξ2x1[i, :, k, e])\n Cx2[i, :, k] += D[2] * (J[i, :, k, e] .* ξ2x2[i, :, k, e])\n Cx3[i, :, k] += D[2] * (J[i, :, k, e] .* ξ2x3[i, :, k, e])\n end\n end\n\n\n for j in 1:Nq[2]\n for i in 1:Nq[1]\n Cx1[i, j, :] += D[3] * (J[i, j, :, e] .* ξ3x1[i, j, :, e])\n Cx2[i, j, :] += D[3] * (J[i, j, :, e] .* ξ3x2[i, j, :, e])\n Cx3[i, j, :] += D[3] * (J[i, j, :, e] .* ξ3x3[i, j, :, e])\n end\n end\n @test maximum(abs.(Cx1)) ≤ 300 * eps(FT)\n @test maximum(abs.(Cx2)) ≤ 300 * eps(FT)\n @test maximum(abs.(Cx3)) ≤ 300 * eps(FT)\n end\n #}}}\n\n #N = 0 test\n #{{{\n let\n for FT in (Float32, Float64)\n N = (4, 4, 0)\n Nq = N .+ 1\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n Nq[j] == 1 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n fx1(ξ1, ξ2, ξ3) = ξ1 + (1 + ξ1)^2 * (1 + ξ2)^2 + ξ3 / 10\n fx1ξ1(ξ1, ξ2, ξ3) = 1 + 2 * (1 + ξ1) * (1 + ξ2)^2\n fx1ξ2(ξ1, ξ2, ξ3) = (1 + ξ1)^2 * 2 * (1 + ξ2)\n fx1ξ3(ξ1, ξ2, ξ3) = 1 / 10\n\n fx2(ξ1, ξ2, ξ3) = ξ2 - (1 + ξ1)^2 + (2 + ξ3) / 2\n fx2ξ1(ξ1, ξ2, ξ3) = -2 * (1 + ξ1)\n fx2ξ2(ξ1, ξ2, ξ3) = 1\n fx2ξ3(ξ1, ξ2, ξ3) = 1 / 2\n\n fx3(ξ1, ξ2, ξ3) = ξ3 + (1 + ξ1)^2 * (1 + ξ2)^2 / 10\n fx3ξ1(ξ1, ξ2, ξ3) = 2 * (1 + ξ1) * (1 + ξ2)^2 / 10\n fx3ξ2(ξ1, ξ2, ξ3) = (1 + ξ1)^2 * 2 * (1 + ξ2) / 10\n fx3ξ3(ξ1, ξ2, ξ3) = 1\n\n e2c = Array{FT, 3}(undef, 3, 8, 1)\n e2c[:, :, 1] = [\n -1 +1 -1 +1 -1 +1 -1 +1\n -1 -1 +1 +1 -1 -1 +1 +1\n -1 -1 -1 -1 +1 +1 +1 +1\n ]\n nelem = size(e2c, 3)\n\n # Create the metrics\n (x1, x2, x3, x1ξ1, x1ξ2, x1ξ3, x2ξ1, x2ξ2, x2ξ3, x3ξ1, x3ξ2, x3ξ3) =\n let\n vgeo = VolumeGeometry(FT, Nq, nelem)\n Metrics.creategrid!(vgeo, e2c, ξ)\n x1 = reshape(vgeo.x1, (Nq..., nelem))\n x2 = reshape(vgeo.x2, (Nq..., nelem))\n x3 = reshape(vgeo.x3, (Nq..., nelem))\n (\n fx1.(x1, x2, x3),\n fx2.(x1, x2, x3),\n fx3.(x1, x2, x3),\n fx1ξ1.(x1, x2, x3),\n fx1ξ2.(x1, x2, x3),\n fx1ξ3.(x1, x2, x3),\n fx2ξ1.(x1, x2, x3),\n fx2ξ2.(x1, x2, x3),\n fx2ξ3.(x1, x2, x3),\n fx3ξ1.(x1, x2, x3),\n fx3ξ2.(x1, x2, x3),\n fx3ξ3.(x1, x2, x3),\n )\n end\n\n J = @.(\n x1ξ1 * (x2ξ2 * x3ξ3 - x3ξ2 * x2ξ3) +\n x2ξ1 * (x3ξ2 * x1ξ3 - x1ξ2 * x3ξ3) +\n x3ξ1 * (x1ξ2 * x2ξ3 - x2ξ2 * x1ξ3)\n )\n\n ξ1x1 = (x2ξ2 .* x3ξ3 - x2ξ3 .* x3ξ2) ./ J\n ξ1x2 = (x3ξ2 .* x1ξ3 - x3ξ3 .* x1ξ2) ./ J\n ξ1x3 = (x1ξ2 .* x2ξ3 - x1ξ3 .* x2ξ2) ./ J\n ξ2x1 = (x2ξ3 .* x3ξ1 - x2ξ1 .* x3ξ3) ./ J\n ξ2x2 = (x3ξ3 .* x1ξ1 - x3ξ1 .* x1ξ3) ./ J\n ξ2x3 = (x1ξ3 .* x2ξ1 - x1ξ1 .* x2ξ3) ./ J\n ξ3x1 = (x2ξ1 .* x3ξ2 - x2ξ2 .* x3ξ1) ./ J\n ξ3x2 = (x3ξ1 .* x1ξ2 - x3ξ2 .* x1ξ1) ./ J\n ξ3x3 = (x1ξ1 .* x2ξ2 - x1ξ2 .* x2ξ1) ./ J\n\n M = J .* reshape(kron(reverse(ω)...), Nq..., 1)\n\n meshwarp(ξ1, ξ2, ξ3) =\n (fx1(ξ1, ξ2, ξ3), fx2(ξ1, ξ2, ξ3), fx3(ξ1, ξ2, ξ3))\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n vgeo = reshape(\n vgeo.array,\n Nq...,\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n @test x1 ≈ vgeo[:, :, :, _x1, :]\n @test x2 ≈ vgeo[:, :, :, _x2, :]\n @test x3 ≈ vgeo[:, :, :, _x3, :]\n\n @test M ≈ vgeo[:, :, :, _M, :]\n\n @test (@view vgeo[:, :, :, _ξ1x1, :]) ≈ ξ1x1\n @test (@view vgeo[:, :, :, _ξ1x2, :]) ≈ ξ1x2\n @test (@view vgeo[:, :, :, _ξ1x3, :]) ≈ ξ1x3\n\n @test (@view vgeo[:, :, :, _ξ2x1, :]) ≈ ξ2x1\n @test (@view vgeo[:, :, :, _ξ2x2, :]) ≈ ξ2x2\n @test (@view vgeo[:, :, :, _ξ2x3, :]) ≈ ξ2x3\n\n @test (@view vgeo[:, :, :, _ξ3x1, :]) ≈ ξ3x1\n @test (@view vgeo[:, :, :, _ξ3x2, :]) ≈ ξ3x2\n @test (@view vgeo[:, :, :, _ξ3x3, :]) ≈ ξ3x3\n\n # check the normals?\n sM = @view sgeo.sωJ[:, :, :]\n n1 = @view sgeo.n1[:, :, :]\n n2 = @view sgeo.n2[:, :, :]\n n3 = @view sgeo.n3[:, :, :]\n @test all(\n hypot.(\n n1[1:Nfp[1], 1:2, :],\n n2[1:Nfp[1], 1:2, :],\n n3[1:Nfp[1], 1:2, :],\n ) .≈ 1,\n )\n @test all(\n hypot.(\n n1[1:Nfp[2], 3:4, :],\n n2[1:Nfp[2], 3:4, :],\n n3[1:Nfp[2], 3:4, :],\n ) .≈ 1,\n )\n @test all(\n hypot.(\n n1[1:Nfp[3], 5:6, :],\n n2[1:Nfp[3], 5:6, :],\n n3[1:Nfp[3], 5:6, :],\n ) .≈ 1,\n )\n\n d, f = 1, 1\n Mf = kron(1, ω[3], ω[2])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[1, :, :, :] .* ξ1x1[1, :, :, :])[:] .* Mf,\n (-J[1, :, :, :] .* ξ1x2[1, :, :, :])[:] .* Mf,\n (-J[1, :, :, :] .* ξ1x3[1, :, :, :])[:] .* Mf,\n ]\n d, f = 1, 2\n Mf = kron(1, ω[3], ω[2])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[Nq[d], :, :, :] .* ξ1x1[Nq[d], :, :, :])[:] .* Mf,\n (J[Nq[d], :, :, :] .* ξ1x2[Nq[d], :, :, :])[:] .* Mf,\n (J[Nq[d], :, :, :] .* ξ1x3[Nq[d], :, :, :])[:] .* Mf,\n ]\n d, f = 2, 3\n Mf = kron(1, ω[3], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[:, 1, :, :] .* ξ2x1[:, 1, :, :])[:] .* Mf,\n (-J[:, 1, :, :] .* ξ2x2[:, 1, :, :])[:] .* Mf,\n (-J[:, 1, :, :] .* ξ2x3[:, 1, :, :])[:] .* Mf,\n ]\n d, f = 2, 4\n Mf = kron(1, ω[3], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[:, Nq[d], :, :] .* ξ2x1[:, Nq[d], :, :])[:] .* Mf,\n (J[:, Nq[d], :, :] .* ξ2x2[:, Nq[d], :, :])[:] .* Mf,\n (J[:, Nq[d], :, :] .* ξ2x3[:, Nq[d], :, :])[:] .* Mf,\n ]\n\n # for these faces we need the N = 1 metrics\n (x1ξ1, x1ξ2, x1ξ3, x2ξ1, x2ξ2, x2ξ3, x3ξ1, x3ξ2, x3ξ3) = let\n @assert Nq[1] != 1 && Nq[2] != 1 && Nq[3] == 1\n Nq_N1 = max.(2, Nq)\n vgeo_N1 = VolumeGeometry(FT, Nq_N1, nelem)\n Metrics.creategrid!(\n vgeo_N1,\n e2c,\n (ξ[1], ξ[2], Elements.lglpoints(FT, 1)[1]),\n )\n x1 = reshape(vgeo_N1.x1, (Nq_N1..., nelem))\n x2 = reshape(vgeo_N1.x2, (Nq_N1..., nelem))\n x3 = reshape(vgeo_N1.x3, (Nq_N1..., nelem))\n (\n fx1ξ1.(x1, x2, x3),\n fx1ξ2.(x1, x2, x3),\n fx1ξ3.(x1, x2, x3),\n fx2ξ1.(x1, x2, x3),\n fx2ξ2.(x1, x2, x3),\n fx2ξ3.(x1, x2, x3),\n fx3ξ1.(x1, x2, x3),\n fx3ξ2.(x1, x2, x3),\n fx3ξ3.(x1, x2, x3),\n )\n end\n J = @.(\n x1ξ1 * (x2ξ2 * x3ξ3 - x3ξ2 * x2ξ3) +\n x2ξ1 * (x3ξ2 * x1ξ3 - x1ξ2 * x3ξ3) +\n x3ξ1 * (x1ξ2 * x2ξ3 - x2ξ2 * x1ξ3)\n )\n ξ3x1 = (x2ξ1 .* x3ξ2 - x2ξ2 .* x3ξ1) ./ J\n ξ3x2 = (x3ξ1 .* x1ξ2 - x3ξ2 .* x1ξ1) ./ J\n ξ3x3 = (x1ξ1 .* x2ξ2 - x1ξ2 .* x2ξ1) ./ J\n\n d, f = 3, 5\n Mf = kron(1, ω[2], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (-J[:, :, 1, :] .* ξ3x1[:, :, 1, :])[:] .* Mf,\n (-J[:, :, 1, :] .* ξ3x2[:, :, 1, :])[:] .* Mf,\n (-J[:, :, 1, :] .* ξ3x3[:, :, 1, :])[:] .* Mf,\n ]\n\n d, f = 3, 6\n Mf = kron(1, ω[2], ω[1])\n @test [\n (sM[1:Nfp[d], f, :] .* n1[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n2[1:Nfp[d], f, :])[:],\n (sM[1:Nfp[d], f, :] .* n3[1:Nfp[d], f, :])[:],\n ] ≈ [\n (J[:, :, 2, :] .* ξ3x1[:, :, 2, :])[:] .* Mf,\n (J[:, :, 2, :] .* ξ3x2[:, :, 2, :])[:] .* Mf,\n (J[:, :, 2, :] .* ξ3x3[:, :, 2, :])[:] .* Mf,\n ]\n end\n end\n #}}}\n\n # Constant preserving test for N = 0\n #{{{\n let\n for FT in (Float64, Float32),\n N in ((4, 4, 0), (0, 0, 2), (0, 3, 4), (2, 0, 3))\n\n Nq = N .+ 1\n\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n Nq[j] == 1 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n rng = MersenneTwister(777)\n fx1(ξ1, ξ2, ξ3) = ξ1 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n fx2(ξ1, ξ2, ξ3) = ξ2 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n fx3(ξ1, ξ2, ξ3) = ξ3 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n\n e2c = Array{FT, 3}(undef, 3, 8, 1)\n e2c[:, :, 1] = [\n -1 +1 -1 +1 -1 +1 -1 +1\n -1 -1 +1 +1 -1 -1 +1 +1\n -1 -1 -1 -1 +1 +1 +1 +1\n ]\n nelem = size(e2c, 3)\n\n meshwarp(ξ1, ξ2, ξ3) =\n (fx1(ξ1, ξ2, ξ3), fx2(ξ1, ξ2, ξ3), fx2(ξ1, ξ2, ξ3))\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n\n vgeo = reshape(\n vgeo.array,\n prod(Nq),\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n M = vgeo[:, _M, :]\n ξ1x1 = vgeo[:, _ξ1x1, :]\n ξ2x1 = vgeo[:, _ξ2x1, :]\n ξ3x1 = vgeo[:, _ξ3x1, :]\n ξ1x2 = vgeo[:, _ξ1x2, :]\n ξ2x2 = vgeo[:, _ξ2x2, :]\n ξ3x2 = vgeo[:, _ξ3x2, :]\n ξ1x3 = vgeo[:, _ξ1x3, :]\n ξ2x3 = vgeo[:, _ξ2x3, :]\n ξ3x3 = vgeo[:, _ξ3x3, :]\n\n M = vgeo[:, _M, :]\n ξ1x1 = vgeo[:, _ξ1x1, :]\n ξ2x1 = vgeo[:, _ξ2x1, :]\n ξ1x2 = vgeo[:, _ξ1x2, :]\n ξ2x2 = vgeo[:, _ξ2x2, :]\n\n I1 = Matrix(I, Nq[1], Nq[1])\n I2 = Matrix(I, Nq[2], Nq[2])\n I3 = Matrix(I, Nq[3], Nq[3])\n D1 = kron(I3, I2, D[1])\n D2 = kron(I3, D[2], I1)\n D3 = kron(D[3], I2, I1)\n\n # Face interpolation operators\n L = (\n kron(I3, I2, I1[1, :]'),\n kron(I3, I2, I1[Nq[1], :]'),\n kron(I3, I2[1, :]', I1),\n kron(I3, I2[Nq[2], :]', I1),\n kron(I3[1, :]', I2, I1),\n kron(I3[Nq[3], :]', I2, I1),\n )\n sM = ntuple(f -> sgeo.sωJ[1:Nfp[cld(f, 2)], f, :], nface)\n n1 = ntuple(f -> sgeo.n1[1:Nfp[cld(f, 2)], f, :], nface)\n n2 = ntuple(f -> sgeo.n2[1:Nfp[cld(f, 2)], f, :], nface)\n n3 = ntuple(f -> sgeo.n3[1:Nfp[cld(f, 2)], f, :], nface)\n\n # If constant preserving then:\n # \\sum_{j} = D' * M * ξjxk = \\sum_{f} L_f' * sM_f * n1_f\n @test D1' * (M .* ξ1x1) + D2' * (M .* ξ2x1) + D3' * (M .* ξ3x1) ≈\n mapreduce((L, sM, n1) -> L' * (sM .* n1), +, L, sM, n1)\n @test D1' * (M .* ξ1x2) + D2' * (M .* ξ2x2) + D3' * (M .* ξ3x2) ≈\n mapreduce((L, sM, n2) -> L' * (sM .* n2), +, L, sM, n2)\n @test D1' * (M .* ξ1x3) + D2' * (M .* ξ2x3) + D3' * (M .* ξ3x3) ≈\n mapreduce((L, sM, n3) -> L' * (sM .* n3), +, L, sM, n3)\n end\n end\n #}}}\n\n # Constant preserving test with all N = 0\n #{{{\n let\n for FT in (Float64, Float32)\n N = (0, 0, 0)\n Nq = N .+ 1\n\n Np = prod(Nq)\n Nfp = div.(Np, Nq)\n\n dim = length(N)\n nface = 2dim\n\n # Create element operators for each polynomial order\n ξω = ntuple(\n j ->\n Nq[j] == 1 ? Elements.glpoints(FT, N[j]) :\n Elements.lglpoints(FT, N[j]),\n dim,\n )\n ξ, ω = ntuple(j -> map(x -> x[j], ξω), 2)\n D = ntuple(j -> Elements.spectralderivative(ξ[j]), dim)\n\n rng = MersenneTwister(777)\n fx1(ξ1, ξ2, ξ3) = ξ1 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n fx2(ξ1, ξ2, ξ3) = ξ2 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n fx3(ξ1, ξ2, ξ3) = ξ3 + (ξ1 * ξ2 * ξ3 * rand(rng) + rand(rng)) / 10\n\n e2c = Array{FT, 3}(undef, 3, 8, 1)\n e2c[:, :, 1] = [\n -1 +1 -1 +1 -1 +1 -1 +1\n -1 -1 +1 +1 -1 -1 +1 +1\n -1 -1 -1 -1 +1 +1 +1 +1\n ]\n nelem = size(e2c, 3)\n\n meshwarp(ξ1, ξ2, ξ3) =\n (fx1(ξ1, ξ2, ξ3), fx2(ξ1, ξ2, ξ3), fx2(ξ1, ξ2, ξ3))\n (vgeo, sgeo, _) = Grids.computegeometry(e2c, D, ξ, ω, meshwarp)\n\n vgeo = reshape(\n vgeo.array,\n prod(Nq),\n # - 1 after fieldcount is to remove the `array` field from the array allocation\n fieldcount(GeometricFactors.VolumeGeometry) - 1,\n nelem,\n )\n\n M = vgeo[:, _M, :]\n ξ1x1 = vgeo[:, _ξ1x1, :]\n ξ2x1 = vgeo[:, _ξ2x1, :]\n ξ3x1 = vgeo[:, _ξ3x1, :]\n ξ1x2 = vgeo[:, _ξ1x2, :]\n ξ2x2 = vgeo[:, _ξ2x2, :]\n ξ3x2 = vgeo[:, _ξ3x2, :]\n ξ1x3 = vgeo[:, _ξ1x3, :]\n ξ2x3 = vgeo[:, _ξ2x3, :]\n ξ3x3 = vgeo[:, _ξ3x3, :]\n\n M = vgeo[:, _M, :]\n ξ1x1 = vgeo[:, _ξ1x1, :]\n ξ2x1 = vgeo[:, _ξ2x1, :]\n ξ1x2 = vgeo[:, _ξ1x2, :]\n ξ2x2 = vgeo[:, _ξ2x2, :]\n\n I1 = Matrix(I, Nq[1], Nq[1])\n I2 = Matrix(I, Nq[2], Nq[2])\n I3 = Matrix(I, Nq[3], Nq[3])\n D1 = kron(I3, I2, D[1])\n D2 = kron(I3, D[2], I1)\n D3 = kron(D[3], I2, I1)\n\n # Face interpolation operators\n L = (\n kron(I3, I2, I1[1, :]'),\n kron(I3, I2, I1[Nq[1], :]'),\n kron(I3, I2[1, :]', I1),\n kron(I3, I2[Nq[2], :]', I1),\n kron(I3[1, :]', I2, I1),\n kron(I3[Nq[3], :]', I2, I1),\n )\n sM = ntuple(f -> sgeo.sωJ[1:Nfp[cld(f, 2)], f, :], nface)\n n1 = ntuple(f -> sgeo.n1[1:Nfp[cld(f, 2)], f, :], nface)\n n2 = ntuple(f -> sgeo.n2[1:Nfp[cld(f, 2)], f, :], nface)\n n3 = ntuple(f -> sgeo.n3[1:Nfp[cld(f, 2)], f, :], nface)\n\n # If constant preserving then \\sum_{f} L_f' * sM_f * n1_f ≈ 0\n @test abs(mapreduce(\n (L, sM, n1) -> L' * (sM .* n1),\n +,\n L,\n sM,\n n1,\n )[1]) < 10 * eps(FT)\n @test abs(mapreduce(\n (L, sM, n2) -> L' * (sM .* n2),\n +,\n L,\n sM,\n n2,\n )[1]) < 10 * eps(FT)\n @test abs(mapreduce(\n (L, sM, n3) -> L' * (sM .* n3),\n +,\n L,\n sM,\n n3,\n )[1]) < 10 * eps(FT)\n end\n end\n #}}}\nend\n" }, { "path": "test/Numerics/Mesh/filter.jl", "content": "using Test\nimport ClimateMachine\nusing ClimateMachine.VariableTemplates: @vars, Vars\nusing ClimateMachine.Mesh.Grids:\n EveryDirection, HorizontalDirection, VerticalDirection\nusing ClimateMachine.MPIStateArrays: weightedsum\nusing ClimateMachine.BalanceLaws\n\nimport GaussQuadrature\nusing MPI\nusing LinearAlgebra\n\nClimateMachine.init()\n\n@testset \"Exponential and Cutoff filter matrix\" begin\n let\n # Values computed with:\n # https://github.com/tcew/nodal-dg/blob/master/Codes1.1/Codes1D/Filter1D.m\n #! format: off\n W = [0x3fe98f3cd0d725e8 0x3fddfd863c6c9a44 0xbfe111110d0fd334 0x3fddbe357bce0b5c 0xbfc970267f929618\n 0x3fb608a150f6f927 0x3fe99528b1a1cd8d 0x3fcd41d41f8bae45 0xbfc987d5fabab8d5 0x3fb5da1cd858af87\n 0xbfb333332eb1cd92 0x3fc666666826f178 0x3fe999999798faaa 0x3fc666666826f176 0xbfb333332eb1cd94\n 0x3fb5da1cd858af84 0xbfc987d5fabab8d4 0x3fcd41d41f8bae46 0x3fe99528b1a1cd8e 0x3fb608a150f6f924\n 0xbfc970267f929618 0x3fddbe357bce0b5c 0xbfe111110d0fd333 0x3fddfd863c6c9a44 0x3fe98f3cd0d725e8]\n #! format: on\n W = reinterpret.(Float64, W)\n\n N = size(W, 1) - 1\n\n topology = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_SELF,\n -1.0:2.0:1.0,\n )\n\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topology;\n polynomialorder = N,\n FloatType = Float64,\n DeviceArray = Array,\n )\n\n filter = ClimateMachine.Mesh.Filters.ExponentialFilter(grid, 0, 32)\n nf = length(filter.filter_matrices)\n @test all(ntuple(i -> filter.filter_matrices[i] ≈ W, nf))\n end\n\n let\n # Values computed with:\n # https://github.com/tcew/nodal-dg/blob/master/Codes1.1/Codes1D/Filter1D.m\n #! format: off\n W = [0x3fd822e5f54ecb62 0x3fedd204a0f08ef8 0xbfc7d3aa58fd6968 0xbfbf74682ac4d276\n 0x3fc7db36e726d8c1 0x3fe59d16feee478b 0x3fc6745bfbb91e20 0xbfa30fbb7a645448\n 0xbfa30fbb7a645455 0x3fc6745bfbb91e26 0x3fe59d16feee478a 0x3fc7db36e726d8c4\n 0xbfbf74682ac4d280 0xbfc7d3aa58fd6962 0x3fedd204a0f08ef7 0x3fd822e5f54ecb62]\n #! format: on\n W = reinterpret.(Float64, W)\n\n N = size(W, 1) - 1\n\n topology = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_SELF,\n -1.0:2.0:1.0,\n )\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topology;\n polynomialorder = N,\n FloatType = Float64,\n DeviceArray = Array,\n )\n\n filter = ClimateMachine.Mesh.Filters.ExponentialFilter(grid, 1, 4)\n nf = length(filter.filter_matrices)\n @test all(ntuple(i -> filter.filter_matrices[i] ≈ W, nf))\n end\n\n let\n T = Float64\n N = (5, 3)\n Nc = (4, 2)\n\n topology = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_SELF,\n -1.0:2.0:1.0,\n )\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topology;\n polynomialorder = N,\n FloatType = T,\n DeviceArray = Array,\n )\n\n ξ = ClimateMachine.Mesh.Grids.referencepoints(grid)\n ξ1 = ξ[1]\n ξ2 = ξ[2]\n a1, b1 = GaussQuadrature.legendre_coefs(T, N[1])\n a2, b2 = GaussQuadrature.legendre_coefs(T, N[2])\n V1 = GaussQuadrature.orthonormal_poly(ξ1, a1, b1)\n V2 = GaussQuadrature.orthonormal_poly(ξ2, a2, b2)\n\n Σ1 = ones(T, N[1] + 1)\n Σ2 = ones(T, N[2] + 1)\n Σ1[(Nc[1]:N[1]) .+ 1] .= 0\n Σ2[(Nc[2]:N[2]) .+ 1] .= 0\n\n W1 = V1 * Diagonal(Σ1) / V1\n W2 = V2 * Diagonal(Σ2) / V2\n\n filter = ClimateMachine.Mesh.Filters.CutoffFilter(grid, Nc)\n @test filter.filter_matrices[1] ≈ W1\n @test filter.filter_matrices[2] ≈ W2\n end\n\n let\n T = Float64\n N = (5, 3)\n Nc = (4, 2)\n\n topology = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_SELF,\n -1.0:2.0:1.0,\n )\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topology;\n polynomialorder = N,\n FloatType = T,\n DeviceArray = Array,\n )\n\n ξ = ClimateMachine.Mesh.Grids.referencepoints(grid)\n ξ1 = ξ[1]\n ξ2 = ξ[2]\n a1, b1 = GaussQuadrature.legendre_coefs(T, N[1])\n a2, b2 = GaussQuadrature.legendre_coefs(T, N[2])\n V1 = GaussQuadrature.orthonormal_poly(ξ1, a1, b1)\n V2 = GaussQuadrature.orthonormal_poly(ξ2, a2, b2)\n\n Σ1 = ones(T, N[1] + 1)\n Σ2 = ones(T, N[2] + 1)\n Σ1[(Nc[1]:N[1]) .+ 1] .= 0\n Σ2[(Nc[2]:N[2]) .+ 1] .= 0\n\n W1 = V1 * Diagonal(Σ1) / V1\n W2 = V2 * Diagonal(Σ2) / V2\n\n filter =\n ClimateMachine.Mesh.Filters.MassPreservingCutoffFilter(grid, Nc)\n @test filter.filter_matrices[1] ≈ W1\n @test filter.filter_matrices[2] ≈ W2\n end\n\nend\n\nstruct FilterTestModel{N} <: ClimateMachine.BalanceLaws.BalanceLaw end\nClimateMachine.BalanceLaws.vars_state(::FilterTestModel, ::Auxiliary, FT) =\n @vars()\nClimateMachine.BalanceLaws.init_state_auxiliary!(::FilterTestModel, _...) =\n nothing\n\n# Legendre Polynomials\nl0(r) = 1\nl1(r) = r\nl2(r) = (3 * r^2 - 1) / 2\nl3(r) = (5 * r^3 - 3r) / 2\n\nlow(x, y, z) = l0(x) * l0(y) + 4 * l1(x) * l1(y) + 5 * l1(z) + 6 * l1(z) * l1(x)\n\nhigh(x, y, z) = l2(x) * l3(y) + l3(x) + l2(y) + l3(z) * l1(y)\n\nfiltered(::EveryDirection, dim, x, y, z) = high(x, y, z)\nfiltered(::VerticalDirection, dim, x, y, z) =\n (dim == 2) ? l2(x) * l3(y) + l2(y) : l3(z) * l1(y)\nfiltered(::HorizontalDirection, dim, x, y, z) =\n (dim == 2) ? l2(x) * l3(y) + l3(x) : l2(x) * l3(y) + l3(x) + l2(y)\n\nClimateMachine.BalanceLaws.vars_state(::FilterTestModel{4}, ::Prognostic, FT) =\n @vars(q1::FT, q2::FT, q3::FT, q4::FT)\nfunction ClimateMachine.BalanceLaws.init_state_prognostic!(\n ::FilterTestModel{4},\n state::Vars,\n aux::Vars,\n localgeo,\n filter_direction,\n dim,\n)\n (x, y, z) = localgeo.coord\n state.q1 = low(x, y, z) + high(x, y, z)\n state.q2 = low(x, y, z) + high(x, y, z)\n state.q3 = low(x, y, z) + high(x, y, z)\n state.q4 = low(x, y, z) + high(x, y, z)\n\n if !isnothing(filter_direction)\n state.q1 -= filtered(filter_direction, dim, x, y, z)\n state.q3 -= filtered(filter_direction, dim, x, y, z)\n end\nend\n\n@testset \"Exponential and Cutoff filter application\" begin\n N = 3\n Ne = (1, 1, 1)\n\n @testset for FT in (Float64, Float32)\n @testset for dim in 2:3\n @testset for direction in (\n EveryDirection,\n HorizontalDirection,\n VerticalDirection,\n )\n brickrange = ntuple(\n j -> range(FT(-1); length = Ne[j] + 1, stop = 1),\n dim,\n )\n topl = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_WORLD,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n )\n\n grid =\n ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = N,\n )\n\n filter = ClimateMachine.Mesh.Filters.CutoffFilter(grid, 2)\n\n model = FilterTestModel{4}()\n dg = ClimateMachine.DGMethods.DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n @testset for target in ((1, 3), (:q1, :q3))\n Q = ClimateMachine.DGMethods.init_ode_state(\n dg,\n nothing,\n dim,\n )\n ClimateMachine.Mesh.Filters.apply!(\n Q,\n target,\n grid,\n filter,\n direction = direction(),\n )\n P = ClimateMachine.DGMethods.init_ode_state(\n dg,\n direction(),\n dim,\n )\n @test Array(Q.data) ≈ Array(P.data)\n end\n end\n end\n end\nend\n\n@testset \"Mass Preserving Cutoff filter application\" begin\n N = 3\n Ne = (1, 1, 1)\n\n @testset for FT in (Float64, Float32)\n @testset for dim in 2:3\n @testset for direction in (\n EveryDirection,\n HorizontalDirection,\n VerticalDirection,\n )\n brickrange = ntuple(\n j -> range(FT(-1); length = Ne[j] + 1, stop = 1),\n dim,\n )\n topl = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_WORLD,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n )\n\n grid =\n ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = N,\n )\n\n filter = ClimateMachine.Mesh.Filters.MassPreservingCutoffFilter(\n grid,\n 2,\n )\n\n model = FilterTestModel{4}()\n dg = ClimateMachine.DGMethods.DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n @testset for target in ((1, 3), (:q1, :q3))\n Q = ClimateMachine.DGMethods.init_ode_state(\n dg,\n nothing,\n dim,\n )\n ClimateMachine.Mesh.Filters.apply!(\n Q,\n target,\n grid,\n filter,\n direction = direction(),\n )\n P = ClimateMachine.DGMethods.init_ode_state(\n dg,\n direction(),\n dim,\n )\n @test Array(Q.data) ≈ Array(P.data)\n end\n end\n end\n end\nend\n\nClimateMachine.BalanceLaws.vars_state(\n ::FilterTestModel{1},\n ::Prognostic,\n FT,\n) where {N} = @vars(q::FT)\nfunction ClimateMachine.BalanceLaws.init_state_prognostic!(\n ::FilterTestModel{1},\n state::Vars,\n aux::Vars,\n localgeo,\n)\n (x, y, z) = localgeo.coord\n state.q = abs(x) - 0.1\nend\n\n@testset \"TMAR filter application\" begin\n\n N = 4\n Ne = (2, 2, 2)\n\n @testset for FT in (Float64, Float32)\n @testset for dim in 2:3\n brickrange =\n ntuple(j -> range(FT(-1); length = Ne[j] + 1, stop = 1), dim)\n topl = ClimateMachine.Mesh.Topologies.BrickTopology(\n MPI.COMM_WORLD,\n brickrange,\n periodicity = ntuple(j -> true, dim),\n )\n\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = N,\n )\n\n model = FilterTestModel{1}()\n dg = ClimateMachine.DGMethods.DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n @testset for target in ((1,), (:q,), :)\n Q = ClimateMachine.DGMethods.init_ode_state(dg)\n\n initialsumQ = weightedsum(Q)\n @test minimum(Q.realdata) < 0\n\n ClimateMachine.Mesh.Filters.apply!(\n Q,\n target,\n grid,\n ClimateMachine.Mesh.Filters.TMARFilter(),\n )\n\n sumQ = weightedsum(Q)\n\n @test minimum(Q.realdata) >= 0\n @test isapprox(initialsumQ, sumQ; rtol = 10 * eps(FT))\n end\n end\n end\nend\n\nfunction cubedshellwarp(a, b, c, R = max(abs(a), abs(b), abs(c)))\n\n function f(sR, ξ, η)\n X, Y = tan(π * ξ / 4), tan(π * η / 4)\n x1 = sR / sqrt(X^2 + Y^2 + 1)\n x2, x3 = X * x1, Y * x1\n x1, x2, x3\n end\n\n fdim = argmax(abs.((a, b, c)))\n if fdim == 1 && a < 0\n # (-R, *, *) : Face I from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(-R, b / a, c / a)\n elseif fdim == 2 && b < 0\n # ( *,-R, *) : Face II from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(-R, a / b, c / b)\n elseif fdim == 1 && a > 0\n # ( R, *, *) : Face III from Ronchi, Iacono, Paolucci (1996)\n x1, x2, x3 = f(R, b / a, c / a)\n elseif fdim == 2 && b > 0\n # ( *, R, *) : Face IV from Ronchi, Iacono, Paolucci (1996)\n x2, x1, x3 = f(R, a / b, c / b)\n elseif fdim == 3 && c > 0\n # ( *, *, R) : Face V from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(R, b / c, a / c)\n elseif fdim == 3 && c < 0\n # ( *, *,-R) : Face VI from Ronchi, Iacono, Paolucci (1996)\n x3, x2, x1 = f(-R, b / c, a / c)\n else\n error(\"invalid case for cubedshellwarp: $a, $b, $c\")\n end\n\n return x1, x2, x3\nend\n\n@testset \"Mass Preserving Cutoff Filter Conservation Test\" begin\n N = 3\n Ne = (1, 1, 1)\n dim = 3\n # dim, direction\n @testset for FT in (Float64,)\n Rrange = [FT(1.0), FT(1.2)]\n\n topl = ClimateMachine.Mesh.Grids.StackedCubedSphereTopology(\n MPI.COMM_WORLD,\n 1,\n Rrange,\n boundary = (5, 6),\n )\n\n grid = ClimateMachine.Mesh.Grids.DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ClimateMachine.array_type(),\n polynomialorder = (N, N),\n meshwarp = cubedshellwarp,\n )\n\n\n mp_filter =\n ClimateMachine.Mesh.Filters.MassPreservingCutoffFilter(grid, 2)\n\n reg_filter = ClimateMachine.Mesh.Filters.CutoffFilter(grid, 2)\n\n model = FilterTestModel{4}()\n dg = ClimateMachine.DGMethods.DGModel(\n model,\n grid,\n nothing,\n nothing,\n nothing;\n state_gradient_flux = nothing,\n )\n\n # test mp filter\n filter = mp_filter\n Q = ClimateMachine.DGMethods.init_ode_state(dg, nothing, dim)\n sum_before_1 = weightedsum(Q, 1)\n sum_before_2 = weightedsum(Q, 2)\n sum_before_3 = weightedsum(Q, 3)\n target = 1:3\n ClimateMachine.Mesh.Filters.apply!(Q, target, grid, filter)\n sum_after_1 = weightedsum(Q, 1)\n sum_after_2 = weightedsum(Q, 2)\n sum_after_3 = weightedsum(Q, 3)\n\n @test sum_before_1 ≈ sum_after_1\n @test sum_before_2 ≈ sum_after_2\n @test sum_before_3 ≈ sum_after_3\n\n\n # test regular filter\n filter = reg_filter\n Q = ClimateMachine.DGMethods.init_ode_state(dg, nothing, dim)\n sum_before_1 = weightedsum(Q, 1)\n sum_before_2 = weightedsum(Q, 2)\n sum_before_3 = weightedsum(Q, 3)\n target = 1:3\n ClimateMachine.Mesh.Filters.apply!(Q, target, grid, filter)\n sum_after_1 = weightedsum(Q, 1)\n sum_after_2 = weightedsum(Q, 2)\n sum_after_3 = weightedsum(Q, 3)\n\n @test !(sum_before_1 ≈ sum_after_1)\n @test !(sum_before_2 ≈ sum_after_2)\n @test !(sum_before_3 ≈ sum_after_3)\n\n end\nend\n" }, { "path": "test/Numerics/Mesh/filter_TMAR.jl", "content": "# This tutorial uses the TMAR Filter from [Light2016](@cite)\n#\n# to reproduce the tutorial in section 4b. It is a shear swirling\n# flow deformation of a transported quantity from LeVeque 1996. The exact\n# solution at the final time is the same as the initial condition.\n\nusing MPI\nusing Test\nusing ClimateMachine\nClimateMachine.init()\nusing Logging\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing LinearAlgebra\nusing Printf\nusing Dates\nusing ClimateMachine.GenericCallbacks:\n EveryXWallTimeSeconds, EveryXSimulationSteps\nusing ClimateMachine.VTK: writevtk, writepvtu\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"test\",\n \"Numerics\",\n \"DGMethods\",\n \"advection_diffusion\",\n \"advection_diffusion_model.jl\",\n))\n\nBase.@kwdef struct SwirlingFlow{FT} <: AdvectionDiffusionProblem\n period::FT = 5\nend\n\ninit_velocity_diffusion!(::SwirlingFlow, aux::Vars, geom::LocalGeometry) =\n nothing\n\ncosbell(τ, q) = τ ≤ 1 ? ((1 + cospi(τ)) / 2)^q : zero(τ)\n\nfunction initial_condition!(::SwirlingFlow, state, aux, localgeo, t)\n FT = eltype(state)\n x, y, _ = aux.coord\n x0, y0 = FT(1 // 4), FT(1 // 4)\n τ = 4 * hypot(x - x0, y - y0)\n state.ρ = cosbell(τ, 3)\nend;\n\nhas_variable_coefficients(::SwirlingFlow) = true\nfunction update_velocity_diffusion!(\n problem::SwirlingFlow,\n ::AdvectionDiffusion,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n x, y, _ = aux.coord\n sx, cx = sinpi(x), cospi(x)\n sy, cy = sinpi(y), cospi(y)\n ct = cospi(t / problem.period)\n\n u = 2 * sx^2 * sy * cy * ct\n v = -2 * sy^2 * sx * cx * ct\n aux.advection.u = SVector(u, v, 0)\nend;\n\nfunction do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, testname)\n ## name of the file that this MPI rank will write\n filename = @sprintf(\n \"%s/%s_mpirank%04d_step%04d\",\n vtkdir,\n testname,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n\n writevtk(filename, Q, dg, statenames)\n\n ## generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s_step%04d\", vtkdir, testname, vtkstep)\n\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"%s_mpirank%04d_step%04d\", testname, i - 1, vtkstep)\n end\n\n writepvtu(pvtuprefix, prefixes, statenames, eltype(Q))\n\n @info \"Done writing VTK: $pvtuprefix\"\n end\nend;\n\nfunction test_run(\n mpicomm,\n ArrayType,\n topl,\n problem,\n dt,\n N,\n timeend,\n FT,\n vtkdir,\n outputtime,\n)\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n bcs = (HomogeneousBC{0}(),)\n model = AdvectionDiffusion{2}(problem, bcs, diffusion = false)\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n initialsumQ = weightedsum(Q)\n\n ## We integrate so that the final solution is equal to the initial solution\n Qe = copy(Q)\n\n rhs! = function (dQdt, Q, ::Nothing, t; increment = false)\n Filters.apply!(Q, :, grid, TMARFilter())\n dg(dQdt, Q, nothing, t; increment = false)\n end\n\n odesolver = SSPRK33ShuOsher(rhs!, Q; dt = dt, t0 = 0)\n\n cbTMAR = EveryXSimulationSteps(1) do\n Filters.apply!(Q, :, grid, TMARFilter())\n end\n\n ## create output directory on first rank of communicator\n if MPI.Comm_rank(mpicomm) == 0\n mkpath(vtkdir)\n end\n MPI.Barrier(mpicomm)\n\n vtkstep = 0\n ## output initial step\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, \"nonnegative\")\n ## setup the output callback\n cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do\n vtkstep += 1\n minQ, maxQ = minimum(Q), maximum(Q)\n sumQ = weightedsum(Q)\n\n sumerror = (initialsumQ - sumQ) / initialsumQ\n\n @info @sprintf \"\"\"Step = %d\n minimum(Q) = %.16e\n maximum(Q) = %.16e\n sum error = %.16e\n \"\"\" vtkstep minQ maxQ sumerror\n\n do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, \"nonnegative\")\n end\n\n callbacks = (cbTMAR, cbvtk)\n solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)\n\n minQ, maxQ = minimum(Q), maximum(Q)\n finalsumQ = weightedsum(Q)\n sumerror = (initialsumQ - finalsumQ) / initialsumQ\n error = euclidean_distance(Q, Qe)\n\n @test minQ ≥ 0\n\n @info @sprintf \"\"\"Finished\n minimum(Q) = %.16e\n maximum(Q) = %.16e\n L2 error = %.16e\n sum error = %.16e\n \"\"\" minQ maxQ error sumerror\nend;\n\nlet\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n FT = Float64\n dim = 2\n Ne = 20\n polynomialorder = 4\n\n problem = SwirlingFlow()\n\n brickrange = (\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = Ne + 1, stop = 1),\n range(FT(0); length = Ne + 1, stop = 1),\n )\n\n topology = BrickTopology(\n mpicomm,\n brickrange[1:dim],\n boundary = ntuple(d -> (1, 1), dim),\n )\n\n maxvelocity = 2\n elementsize = 1 / Ne\n dx = elementsize / polynomialorder^2\n CFL = 1\n dt = CFL * dx / maxvelocity\n\n vtkdir =\n abspath(joinpath(ClimateMachine.Settings.output_dir, \"vtk_nonnegative\"))\n outputtime = 0.0625\n dt = outputtime / ceil(Int64, outputtime / dt)\n\n timeend = problem.period\n\n @info @sprintf \"\"\"Starting\n FT = %s\n dim = %d\n Ne = %d\n polynomial order = %d\n final time = %.16e\n time step = %.16e\n \"\"\" FT dim Ne polynomialorder timeend dt\n\n test_run(\n mpicomm,\n ArrayType,\n topology,\n problem,\n dt,\n polynomialorder,\n timeend,\n FT,\n vtkdir,\n outputtime,\n )\nend;\n" }, { "path": "test/Numerics/Mesh/grid_integral.jl", "content": "using Test\nusing ClimateMachine\n\nlet\n for N in 1:10\n for FloatType in (Float64, Float32)\n (ξ, ω) = ClimateMachine.Mesh.Elements.lglpoints(FloatType, N)\n I∫ =\n ClimateMachine.Mesh.Grids.indefinite_integral_interpolation_matrix(\n ξ,\n ω,\n )\n for n in 1:N\n if N == 1\n @test sum(abs.(I∫ * ξ)) < 10 * eps(FloatType)\n else\n @test I∫ * ξ .^ n ≈\n (ξ .^ (n + 1) .- (-1) .^ (n + 1)) / (n + 1)\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/Mesh/interpolation.jl", "content": "using Dates\nusing LinearAlgebra\nusing Logging\nusing MPI\nusing Printf\nusing StaticArrays\nusing Statistics\nusing Test\nimport GaussQuadrature\nusing KernelAbstractions\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Atmos\nusing ClimateMachine.Atmos: vars_state\nusing ClimateMachine.Orientations\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Geometry\nusing ClimateMachine.Mesh.Interpolation\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Writers\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: R_d, planet_radius, grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n#-------------------------------------\nfcn(x, y, z) = sin(x) * cos(y) * cos(z) # sample function\n#-------------------------------------\nfunction Initialize_Brick_Interpolation_Test!(\n problem,\n bl,\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n)\n FT = eltype(state)\n # Dummy variables for initial condition function\n state.ρ = FT(0)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = FT(0)\n state.moisture.ρq_tot = FT(0)\nend\n#------------------------------------------------\nstruct TestSphereSetup{DT}\n p_ground::DT\n T_initial::DT\n domain_height::DT\n\n function TestSphereSetup(\n p_ground::DT,\n T_initial::DT,\n domain_height::DT,\n ) where {DT <: AbstractFloat}\n return new{DT}(p_ground, T_initial, domain_height)\n end\nend\n#----------------------------------------------------------------------------\nfunction (setup::TestSphereSetup)(problem, bl, state, aux, coords, t)\n # callable to set initial conditions\n FT = eltype(state)\n param_set = parameter_set(bl)\n _grav::FT = grav(param_set)\n _R_d::FT = R_d(param_set)\n\n z = altitude(bl, aux)\n\n scale_height::FT = _R_d * setup.T_initial / _grav\n p::FT = setup.p_ground * exp(-z / scale_height)\n e_int = internal_energy(param_set, setup.T_initial)\n e_pot = gravitational_potential(bl.orientation, aux)\n\n # TODO: Fix type instability: typeof(setup.T_initial) == typeof(p) fails\n state.ρ = air_density(param_set, FT(setup.T_initial), p)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n return nothing\nend\n#----------------------------------------------------------------------------\nfunction run_brick_interpolation_test(\n ::Type{DA},\n ::Type{FT},\n polynomialorders,\n toler::FT,\n) where {DA, FT <: AbstractFloat}\n mpicomm = MPI.COMM_WORLD\n root = 0\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n\n xmin, ymin, zmin = FT(0), FT(0), FT(0) # defining domain extent\n xmax, ymax, zmax = FT(2000), FT(400), FT(2000)\n xres = [FT(10), FT(10), FT(10)] # resolution of interpolation grid\n\n Ne = (20, 4, 20)\n #-------------------------\n _x, _y, _z = ClimateMachine.Mesh.Grids.vgeoid.x1id,\n ClimateMachine.Mesh.Grids.vgeoid.x2id,\n ClimateMachine.Mesh.Grids.vgeoid.x3id\n #-------------------------\n brickrange = (\n range(FT(xmin); length = Ne[1] + 1, stop = xmax),\n range(FT(ymin); length = Ne[2] + 1, stop = ymax),\n range(FT(zmin); length = Ne[3] + 1, stop = zmax),\n )\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n periodicity = (true, true, false),\n )\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = DA,\n polynomialorder = polynomialorders,\n )\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(0)),\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n orientation = SphericalOrientation(),\n init_state_prognostic = Initialize_Brick_Interpolation_Test!,\n source = (Gravity(),),\n )\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n\n Q = init_ode_state(dg, FT(0))\n\n #------------------------------\n x1 = @view grid.vgeo[:, _x:_x, :]\n x2 = @view grid.vgeo[:, _y:_y, :]\n x3 = @view grid.vgeo[:, _z:_z, :]\n #----calling interpolation function on state variable # st_idx--------------------------\n nvars = size(Q.data, 2)\n Q.data .= sin.(x1 ./ xmax) .* cos.(x2 ./ ymax) .* cos.(x3 ./ zmax)\n\n xbnd = Array{FT}(undef, 2, 3)\n\n xbnd[1, 1] = FT(xmin)\n xbnd[2, 1] = FT(xmax)\n xbnd[1, 2] = FT(ymin)\n xbnd[2, 2] = FT(ymax)\n xbnd[1, 3] = FT(zmin)\n xbnd[2, 3] = FT(zmax)\n #----------------------------------------------------------\n x1g = collect(range(xbnd[1, 1], xbnd[2, 1], step = xres[1]))\n nx1 = length(x1g)\n x2g = collect(range(xbnd[1, 2], xbnd[2, 2], step = xres[2]))\n nx2 = length(x2g)\n x3g = collect(range(xbnd[1, 3], xbnd[2, 3], step = xres[3]))\n nx3 = length(x3g)\n\n filename = \"test.nc\"\n varnames = (\"ρ\", \"ρu\", \"ρv\", \"ρw\", \"e\", \"other\")\n\n intrp_brck = InterpolationBrick(grid, xbnd, x1g, x2g, x3g) # sets up the interpolation structure\n iv = DA(Array{FT}(undef, intrp_brck.Npl, nvars)) # allocating space for the interpolation variable\n if pid == 0\n fiv = DA(Array{FT}(undef, nx1, nx2, nx3, nvars)) # allocating space for the full interpolation variables accumulated on proc# 0\n else\n fiv = DA(Array{FT}(undef, 0, 0, 0, 0))\n end\n interpolate_local!(intrp_brck, Q.data, iv) # interpolation\n\n accumulate_interpolated_data!(intrp_brck, iv, fiv) # write interpolation data to file\n #------------------------------\n err_inf_dom = zeros(FT, nvars)\n\n x1g = intrp_brck.x1g\n x2g = intrp_brck.x2g\n x3g = intrp_brck.x3g\n\n if pid == 0\n nx1 = length(x1g)\n nx2 = length(x2g)\n nx3 = length(x3g)\n x1 = Array{FT}(undef, nx1, nx2, nx3)\n x2 = similar(x1)\n x3 = similar(x1)\n\n fiv_cpu = Array(fiv)\n\n for k in 1:nx3, j in 1:nx2, i in 1:nx1\n x1[i, j, k] = x1g[i]\n x2[i, j, k] = x2g[j]\n x3[i, j, k] = x3g[k]\n end\n\n fex = sin.(x1 ./ xmax) .* cos.(x2 ./ ymax) .* cos.(x3 ./ zmax)\n\n for vari in 1:nvars\n err_inf_dom[vari] =\n maximum(abs.(fiv_cpu[:, :, :, vari] .- fex[:, :, :]))\n end\n end\n\n MPI.Bcast!(err_inf_dom, root, mpicomm)\n\n if maximum(err_inf_dom) > toler\n if pid == 0\n println(\"err_inf_domain = $(maximum(err_inf_dom)) is larger than prescribed tolerance of $toler\")\n end\n MPI.Barrier(mpicomm)\n end\n @test maximum(err_inf_dom) < toler\n return nothing\n #----------------\nend #function run_brick_interpolation_test\n#----------------------------------------------------------------------------\n\n#----------------------------------------------------------------------------\n# Cubed sphere, lat/long interpolation test\n#----------------------------------------------------------------------------\nfunction run_cubed_sphere_interpolation_test(\n ::Type{DA},\n ::Type{FT},\n polynomialorders,\n toler::FT,\n) where {DA, FT <: AbstractFloat}\n mpicomm = MPI.COMM_WORLD\n root = 0\n pid = MPI.Comm_rank(mpicomm)\n npr = MPI.Comm_size(mpicomm)\n\n domain_height = FT(30e3)\n numelem_horz = 6\n numelem_vert = 4\n #-------------------------\n _x, _y, _z = ClimateMachine.Mesh.Grids.vgeoid.x1id,\n ClimateMachine.Mesh.Grids.vgeoid.x2id,\n ClimateMachine.Mesh.Grids.vgeoid.x3id\n _ρ, _ρu, _ρv, _ρw = 1, 2, 3, 4\n #-------------------------\n _planet_radius::FT = planet_radius(param_set)\n\n vert_range = grid1d(\n _planet_radius,\n FT(_planet_radius + domain_height),\n nelem = numelem_vert,\n )\n\n lat_res = FT(1) # 1 degree resolution\n long_res = FT(1) # 1 degree resolution\n nel_vert_grd = 20 #100 #50 #10#50\n rad_res = FT((vert_range[end] - vert_range[1]) / FT(nel_vert_grd)) # 1000 m vertical resolution\n #----------------------------------------------------------\n _MSLP::FT = MSLP(param_set)\n setup = TestSphereSetup(_MSLP, FT(255), FT(30e3))\n\n topology = StackedCubedSphereTopology(mpicomm, numelem_horz, vert_range)\n\n grid = DiscontinuousSpectralElementGrid(\n topology,\n FloatType = FT,\n DeviceArray = DA,\n polynomialorder = polynomialorders,\n meshwarp = ClimateMachine.Mesh.Topologies.equiangular_cubed_sphere_warp,\n )\n\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = NoReferenceState(),\n turbulence = ConstantDynamicViscosity(FT(0)),\n moisture = DryModel(),\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = setup,\n source = (),\n )\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n\n #------------------------------\n x1 = @view grid.vgeo[:, _x:_x, :]\n x2 = @view grid.vgeo[:, _y:_y, :]\n x3 = @view grid.vgeo[:, _z:_z, :]\n\n xmax = _planet_radius\n ymax = _planet_radius\n zmax = _planet_radius\n\n nvars = size(Q.data, 2)\n\n Q.data .= sin.(x1 ./ xmax) .* cos.(x2 ./ ymax) .* cos.(x3 ./ zmax)\n #for ivar in 1:nvars\n # Q.data[:, ivar, :] .=\n # sin.(x1[:, 1, :] ./ xmax) .* cos.(x2[:, 1, :] ./ ymax) .*\n # cos.(x3[:, 1, :] ./ zmax)\n #end\n #------------------------------\n lat_min, lat_max = FT(-90.0), FT(90.0) # inclination/zeinth angle range\n long_min, long_max = FT(-180.0), FT(180.0) # azimuthal angle range\n rad_min, rad_max = vert_range[1], vert_range[end] # radius range\n\n\n lat_grd = collect(range(lat_min, lat_max, step = lat_res))\n n_lat = length(lat_grd)\n long_grd = collect(range(long_min, long_max, step = long_res))\n n_long = length(long_grd)\n rad_grd = collect(range(rad_min, rad_max, step = rad_res))\n n_rad = length(rad_grd)\n\n _ρu, _ρv, _ρw = 2, 3, 4\n\n filename = \"test.nc\"\n varnames = (\"ρ\", \"ρu\", \"ρv\", \"ρw\", \"e\")\n projectv = true\n\n intrp_cs = InterpolationCubedSphere(\n grid,\n collect(vert_range),\n numelem_horz,\n lat_grd,\n long_grd,\n rad_grd,\n ) # sets up the interpolation structure\n iv = DA(Array{FT}(undef, intrp_cs.Npl, nvars)) # allocating space for the interpolation variable\n if pid == 0\n fiv = DA(Array{FT}(undef, n_long, n_lat, n_rad, nvars)) # allocating space for the full interpolation variables accumulated on proc# 0\n else\n fiv = DA(Array{FT}(undef, 0, 0, 0, 0))\n end\n\n interpolate_local!(intrp_cs, Q.data, iv) # interpolation\n project_cubed_sphere!(intrp_cs, iv, (_ρu, _ρv, _ρw)) # project velocity onto unit vectors along rad, lat & long\n accumulate_interpolated_data!(intrp_cs, iv, fiv) # accumulate interpolated data on to proc# 0\n #----------------------------------------------------------\n # Testing\n err_inf_dom = zeros(FT, nvars)\n rad = Array(intrp_cs.rad_grd)\n lat = Array(intrp_cs.lat_grd)\n long = Array(intrp_cs.long_grd)\n fiv_cpu = Array(fiv)\n if pid == 0\n nrad = length(rad)\n nlat = length(lat)\n nlong = length(long)\n x1g = Array{FT}(undef, nrad, nlat, nlong)\n x2g = similar(x1g)\n x3g = similar(x1g)\n\n fex = zeros(FT, nlong, nlat, nrad, nvars)\n\n for vari in 1:nvars\n for i in 1:nlong, j in 1:nlat, k in 1:nrad\n x1g_ijk = rad[k] * cosd(lat[j]) * cosd(long[i]) # inclination -> latitude; azimuthal -> longitude.\n x2g_ijk = rad[k] * cosd(lat[j]) * sind(long[i]) # inclination -> latitude; azimuthal -> longitude.\n x3g_ijk = rad[k] * sind(lat[j])\n\n fex[i, j, k, vari] =\n fcn(x1g_ijk / xmax, x2g_ijk / ymax, x3g_ijk / zmax)\n end\n end\n\n if projectv\n for i in 1:nlong, j in 1:nlat, k in 1:nrad\n fex[i, j, k, _ρu] =\n -fex[i, j, k, _ρ] * sind(long[i]) +\n fex[i, j, k, _ρ] * cosd(long[i])\n\n fex[i, j, k, _ρv] =\n -fex[i, j, k, _ρ] * sind(lat[j]) * cosd(long[i]) -\n fex[i, j, k, _ρ] * sind(lat[j]) * sind(long[i]) +\n fex[i, j, k, _ρ] * cosd(lat[j])\n\n fex[i, j, k, _ρw] =\n fex[i, j, k, _ρ] * cosd(lat[j]) * cosd(long[i]) +\n fex[i, j, k, _ρ] * cosd(lat[j]) * sind(long[i]) +\n fex[i, j, k, _ρ] * sind(lat[j])\n end\n end\n\n for vari in 1:nvars\n err_inf_dom[vari] =\n maximum(abs.(fiv_cpu[:, :, :, vari] .- fex[:, :, :, vari]))\n end\n end\n\n MPI.Bcast!(err_inf_dom, root, mpicomm)\n\n if maximum(err_inf_dom) > toler\n if pid == 0\n println(\"err_inf_domain = $(maximum(err_inf_dom)) is larger than prescribed tolerance of $toler\")\n end\n MPI.Barrier(mpicomm)\n end\n @test maximum(err_inf_dom) < toler\n return nothing\nend\n#----------------------------------------------------------------------------\n@testset \"Interpolation tests\" begin\n DA = ClimateMachine.array_type()\n\n run_brick_interpolation_test(DA, Float32, (0), Float32(1E-1))\n run_brick_interpolation_test(DA, Float64, (0), Float64(1E-1))\n\n run_brick_interpolation_test(DA, Float32, (5), Float32(1E-6))\n run_brick_interpolation_test(DA, Float64, (5), Float64(1E-9))\n\n run_brick_interpolation_test(DA, Float32, (5, 6), Float32(1E-6))\n run_brick_interpolation_test(DA, Float64, (5, 6), Float64(1E-9))\n\n run_cubed_sphere_interpolation_test(DA, Float32, (0), Float32(2e-1))\n run_cubed_sphere_interpolation_test(DA, Float64, (0), Float64(2e-1))\n\n run_cubed_sphere_interpolation_test(DA, Float32, (5), Float32(2e-6))\n run_cubed_sphere_interpolation_test(DA, Float64, (5), Float64(2e-7))\n\n run_cubed_sphere_interpolation_test(DA, Float32, (5, 6), Float32(2e-6))\n run_cubed_sphere_interpolation_test(DA, Float64, (5, 6), Float64(2e-7))\nend\n#------------------------------------------------\n" }, { "path": "test/Numerics/Mesh/min_node_distance.jl", "content": "using Test\nusing MPI\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.VTK\nusing Logging\nusing Printf\nusing LinearAlgebra\n\nlet\n # boiler plate MPI stuff\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n # Mesh generation parameters\n Neh = 10\n Nev = 4\n\n Ns = ((4, (4, 3), (2, 3)), (4, (2, 3, 4), (4, 2, 3), (4, 2, 3)))\n @testset \"$(@__FILE__) DGModel matrix\" begin\n for FT in (Float64, Float32)\n for dim in (2, 3)\n for N in Ns[dim - 1]\n if dim == 2\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n elseif dim == 3\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n end\n\n topl = StackedBrickTopology(mpicomm, brickrange)\n\n function warpfun(ξ1, ξ2, ξ3)\n FT = eltype(ξ1)\n\n ξ1 ≥ FT(1 // 2) &&\n (ξ1 = FT(1 // 2) + 2 * (ξ1 - FT(1 // 2)))\n if dim == 2\n ξ2 ≥ FT(3 // 2) &&\n (ξ2 = FT(3 // 2) + 2 * (ξ2 - FT(3 // 2)))\n elseif dim == 3\n ξ2 ≥ FT(1 // 2) &&\n (ξ2 = FT(1 // 2) + 2 * (ξ2 - FT(1 // 2)))\n ξ3 ≥ FT(3 // 2) &&\n (ξ3 = FT(3 // 2) + 2 * (ξ3 - FT(3 // 2)))\n end\n (ξ1, ξ2, ξ3)\n end\n\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n # testname = \"grid_poly$(N)_dim$(dim)_$(ArrayType)_$(FT)\"\n # filename(rank) = @sprintf(\"%s_mpirank%04d\", testname, rank)\n # writevtk(filename(MPI.Comm_rank(mpicomm)), grid)\n # if MPI.Comm_rank(mpicomm) == 0\n # writepvtu(testname, filename.(0:MPI.Comm_size(mpicomm)-1), (), FT)\n # end\n\n ξ = referencepoints(grid)\n Δξ = ntuple(d -> ξ[d][2] - ξ[d][1], dim)\n\n hmnd = minimum(Δξ[1:(dim - 1)]) / (2Neh)\n vmnd = Δξ[end] / (2Nev)\n\n @test hmnd ≈ min_node_distance(grid, EveryDirection())\n @test vmnd ≈ min_node_distance(grid, VerticalDirection())\n @test hmnd ≈ min_node_distance(grid, HorizontalDirection())\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/Mesh/mpi_centroid.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.BrickMesh\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n\n (elemtovert, elemtocorner, facecode) = brickmesh(\n (2.0:5.0, 4.0:6.0),\n (false, true);\n part = MPI.Comm_rank(comm) + 1,\n numparts = MPI.Comm_size(comm),\n )\n\n code = centroidtocode(comm, elemtocorner)\n (d, nelem) = size(code)\n\n root = 0\n counts = MPI.Gather(Cint(length(code)), 0, comm)\n code_all = MPI.Gatherv!(\n code,\n MPI.Comm_rank(comm) == root ?\n VBuffer(similar(code, sum(counts)), counts) : nothing,\n root,\n comm,\n )\n\n if MPI.Comm_rank(comm) == root\n\n code_all = reshape(code_all, d, div(sum(counts), d))\n\n code_expect = UInt64[\n 0x0000000000000000 0x1555555555555555 0xffffffffffffffff 0x5555555555555555 0x6aaaaaaaaaaaaaaa 0xaaaaaaaaaaaaaaaa\n 0x0000000000000000 0x5555555555555555 0xffffffffffffffff 0x5555555555555555 0xaaaaaaaaaaaaaaaa 0xaaaaaaaaaaaaaaaa\n ]\n @test code_all == code_expect\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connect.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 3\n\n topology = BrickTopology(\n comm,\n (0:4, 5:9);\n boundary = ((1, 2), (3, 4)),\n periodicity = (false, true),\n connectivity = :face,\n )\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoface = [\n 1 2 2 1 1 1 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n\n globalelemtocoord = Array{Int}(undef, 2, 4, 16)\n globalelemtocoord[:, :, 1] = [0 1 0 1; 5 5 6 6]\n globalelemtocoord[:, :, 2] = [1 2 1 2; 5 5 6 6]\n globalelemtocoord[:, :, 3] = [1 2 1 2; 6 6 7 7]\n globalelemtocoord[:, :, 4] = [0 1 0 1; 6 6 7 7]\n globalelemtocoord[:, :, 5] = [0 1 0 1; 7 7 8 8]\n globalelemtocoord[:, :, 6] = [0 1 0 1; 8 8 9 9]\n globalelemtocoord[:, :, 7] = [1 2 1 2; 8 8 9 9]\n globalelemtocoord[:, :, 8] = [1 2 1 2; 7 7 8 8]\n globalelemtocoord[:, :, 9] = [2 3 2 3; 7 7 8 8]\n globalelemtocoord[:, :, 10] = [2 3 2 3; 8 8 9 9]\n globalelemtocoord[:, :, 11] = [3 4 3 4; 8 8 9 9]\n globalelemtocoord[:, :, 12] = [3 4 3 4; 7 7 8 8]\n globalelemtocoord[:, :, 13] = [3 4 3 4; 6 6 7 7]\n globalelemtocoord[:, :, 14] = [2 3 2 3; 6 6 7 7]\n globalelemtocoord[:, :, 15] = [2 3 2 3; 5 5 6 6]\n globalelemtocoord[:, :, 16] = [3 4 3 4; 5 5 6 6]\n\n globalelemtobndy = [\n 1 0 0 1 1 1 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 2 2 2 0 0 2\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n ]\n\n if crank == 0\n nrealelem = 5\n globalelems = [1, 2, 3, 4, 5, 6, 7, 8, 14, 15]\n elemtoelem_expect = [\n 1 1 4 2 3 4 7 8 9 10\n 2 10 9 3 8 6 7 8 9 10\n 6 7 2 1 4 6 7 8 9 10\n 4 3 8 5 6 6 7 8 9 10\n ]\n nabrtorank_expect = [1, 2]\n nabrtorecv_expect = UnitRange{Int}[1:3, 4:5]\n nabrtosend_expect = UnitRange{Int}[1:4, 5:6]\n elseif crank == 1\n nrealelem = 5\n globalelems = [6, 7, 8, 9, 10, 1, 2, 3, 5, 11, 12, 14, 15]\n elemtoelem_expect = [\n 1 1 9 3 2 2 7 8 3 10 11 12 13\n 2 5 4 11 10 6 7 8 9 1 2 12 13\n 9 3 8 12 4 6 7 8 9 10 11 12 13\n 6 7 2 5 13 6 7 8 9 10 11 12 13\n ]\n nabrtorank_expect = [0, 2]\n nabrtorecv_expect = UnitRange{Int}[1:4, 5:8]\n nabrtosend_expect = UnitRange{Int}[1:3, 4:5]\n elseif crank == 2\n nrealelem = 6\n globalelems = [11, 12, 13, 14, 15, 16, 2, 3, 9, 10]\n elemtoelem_expect = [\n 10 9 4 8 7 5 7 8 9 10\n 1 2 3 3 6 4 7 8 9 10\n 2 3 6 5 10 1 7 8 9 10\n 6 1 2 9 4 3 7 8 9 10\n ]\n nabrtorank_expect = [0, 1]\n nabrtorecv_expect = UnitRange{Int}[1:2, 3:4]\n nabrtosend_expect = UnitRange{Int}[1:2, 3:6]\n end\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n @test elemtocoord == globalelemtocoord[:, :, globalelems]\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connect_1d.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 5\n\n topology = BrickTopology(\n comm,\n (0:10,);\n boundary = ((1, 2),),\n periodicity = (true,),\n )\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoelem = [\n 10 1 2 3 4 5 6 7 8 9\n 2 3 4 5 6 7 8 9 10 1\n ]\n\n globalelemtoface = [\n 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n globalelemtobndy = zeros(Int, size(globalelemtoface))\n\n globalelemtocoord = Array{Int}(undef, 1, 2, 10)\n globalelemtocoord[:, :, 1] = [0 1]\n globalelemtocoord[:, :, 2] = [1 2]\n globalelemtocoord[:, :, 3] = [2 3]\n globalelemtocoord[:, :, 4] = [3 4]\n globalelemtocoord[:, :, 5] = [4 5]\n globalelemtocoord[:, :, 6] = [5 6]\n globalelemtocoord[:, :, 7] = [6 7]\n globalelemtocoord[:, :, 8] = [7 8]\n globalelemtocoord[:, :, 9] = [8 9]\n globalelemtocoord[:, :, 10] = [9 10]\n\n @assert csize == 5\n nrealelem = 2\n if crank == 0\n globalelems = [1, 2, 3, 10]\n elemtoelem_expect = [4 1 3 4; 2 3 3 4]\n nabrtorank_expect = [1, 4]\n elseif crank == 1\n globalelems = [3, 4, 2, 5]\n elemtoelem_expect = [3 1 3 4; 2 4 3 4]\n nabrtorank_expect = [0, 2]\n elseif crank == 2\n globalelems = [5, 6, 4, 7]\n elemtoelem_expect = [3 1 3 4; 2 4 3 4]\n nabrtorank_expect = [1, 3]\n elseif crank == 3\n globalelems = [7, 8, 6, 9]\n elemtoelem_expect = [3 1 3 4; 2 4 3 4]\n nabrtorank_expect = [2, 4]\n elseif crank == 4\n globalelems = [9, 10, 1, 8]\n elemtoelem_expect = [4 1 3 4; 2 3 3 4]\n nabrtorank_expect = [0, 3]\n end\n nabrtorecv_expect = UnitRange{Int}[1:1, 2:2]\n nabrtosend_expect = UnitRange{Int}[1:1, 2:2]\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n @test elemtocoord == globalelemtocoord[:, :, globalelems]\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == globalelemtoordr[:, globalelems]\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connect_ell.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 2\n\n FT = Float64\n Nx = 3\n Ny = 2\n x = range(FT(0); length = Nx + 1, stop = 1)\n y = range(FT(0); length = Ny + 1, stop = 1)\n\n topology = BrickTopology(\n comm,\n (x, y);\n boundary = ((1, 2), (3, 4)),\n periodicity = (true, true),\n connectivity = :face,\n )\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoface = [\n 2 2 2 2 2 2\n 1 1 1 1 1 1\n 4 4 4 4 4 4\n 3 3 3 3 3 3\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n globalelemtobndy = zeros(Int, size(globalelemtoface))\n\n if crank == 0\n nrealelem = 3\n globalelems = [1, 2, 3, 4, 5, 6]\n elemtoelem_expect = [\n 6 4 2 4 5 6\n 5 3 4 4 5 6\n 2 1 5 4 5 6\n 2 1 5 4 5 6\n ]\n nabrtorank_expect = [1]\n nabrtorecv_expect = UnitRange{Int}[1:3]\n nabrtosend_expect = UnitRange{Int}[1:3]\n elseif crank == 1\n nrealelem = 3\n globalelems = [4, 5, 6, 1, 2, 3]\n elemtoelem_expect = [\n 6 4 2 4 5 6\n 5 3 4 4 5 6\n 3 6 1 4 5 6\n 3 6 1 4 5 6\n ]\n nabrtorank_expect = [0]\n nabrtorecv_expect = UnitRange{Int}[1:3]\n nabrtosend_expect = UnitRange{Int}[1:3]\n end\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connect_sphere.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.MPIStateArrays\nusing KernelAbstractions\n\nfunction main()\n FT = Float64\n Nhorz = 3\n Nstack = 5\n N = 4\n DA = Array\n\n MPI.Initialized() || MPI.Init()\n\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n Rrange = FT.(accumulate(+, 1:(Nstack + 1)))\n topology = StackedCubedSphereTopology(\n MPI.COMM_SELF,\n Nhorz,\n Rrange;\n boundary = (1, 2),\n connectivity = :face,\n )\n grid = DiscontinuousSpectralElementGrid(\n topology;\n FloatType = FT,\n DeviceArray = DA,\n polynomialorder = N,\n meshwarp = Topologies.equiangular_cubed_sphere_warp,\n )\n\n let\n Np = (N + 1)^3\n activedofs = Array(grid.activedofs)\n vmaprecv = Array(grid.vmaprecv)\n active = union(1:(length(topology.realelems) * Np), vmaprecv)\n inactive = setdiff(1:(length(topology.elems) * Np), active)\n\n @test all(activedofs[active] .== true)\n @test all(activedofs[inactive] .== false)\n end\n\n #=\n @show elems = topology.elems\n @show realelems = topology.realelems\n @show ghostelems = topology.ghostelems\n @show sendelems = topology.sendelems\n @show elemtocoord = topology.elemtocoord\n @show elemtoelem = topology.elemtoelem\n @show elemtoface = topology.elemtoface\n @show elemtoordr = topology.elemtoordr\n @show elemtobndy = topology.elemtobndy\n @show nabrtorank = topology.nabrtorank\n @show nabrtorecv = topology.nabrtorecv\n @show nabrtosend = topology.nabrtosend\n =#\n\n # Check x1x2x3 matches before comm\n x1 = @view grid.vgeo[:, Grids._x1, :]\n x2 = @view grid.vgeo[:, Grids._x2, :]\n x3 = @view grid.vgeo[:, Grids._x3, :]\n\n interior_faces = vec(grid.elemtobndy .== 0)\n interior_vmap⁻ =\n reshape(grid.vmap⁻, (size(grid.vmap⁻, 1), :))[:, interior_faces]\n interior_vmap⁺ =\n reshape(grid.vmap⁺, (size(grid.vmap⁺, 1), :))[:, interior_faces]\n\n @test x1[interior_vmap⁻] ≈ x1[interior_vmap⁺]\n @test x2[interior_vmap⁻] ≈ x2[interior_vmap⁺]\n @test x3[interior_vmap⁻] ≈ x3[interior_vmap⁺]\n\n Np = (N + 1)^3\n x1x2x3 = MPIStateArray{FT}(\n topology.mpicomm,\n DA,\n Np,\n 3,\n length(topology.elems),\n realelems = topology.realelems,\n ghostelems = topology.ghostelems,\n vmaprecv = grid.vmaprecv,\n vmapsend = grid.vmapsend,\n nabrtorank = topology.nabrtorank,\n nabrtovmaprecv = grid.nabrtovmaprecv,\n nabrtovmapsend = grid.nabrtovmapsend,\n )\n x1x2x3.data[:, :, topology.realelems] .= @view grid.vgeo[\n :,\n [Grids._x1, Grids._x2, Grids._x3],\n topology.realelems,\n ]\n\n event = Event(array_device(x1x2x3))\n event = MPIStateArrays.begin_ghost_exchange!(x1x2x3, dependencies = event)\n event = MPIStateArrays.end_ghost_exchange!(x1x2x3, dependencies = event)\n wait(array_device(x1x2x3), event)\n\n # Check x1x2x3 matches after\n x1 = @view x1x2x3.data[:, 1, :]\n x2 = @view x1x2x3.data[:, 2, :]\n x3 = @view x1x2x3.data[:, 3, :]\n\n @test x1[interior_vmap⁻] ≈ x1[interior_vmap⁺]\n @test x2[interior_vmap⁻] ≈ x2[interior_vmap⁺]\n @test x3[interior_vmap⁻] ≈ x3[interior_vmap⁺]\n\n nothing\nend\nisinteractive() || main()\n" }, { "path": "test/Numerics/Mesh/mpi_connect_stacked.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 3\n\n topology = StackedBrickTopology(\n comm,\n (2:5, 4:6),\n periodicity = (false, true),\n boundary = ((1, 2), (3, 4)),\n connectivity = :face,\n )\n\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoface = [\n 1 1 2 2 2 2\n 1 1 1 1 2 2\n 4 4 4 4 4 4\n 3 3 3 3 3 3\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n\n globalelemtocoord = Array{Int}(undef, 2, 4, 6)\n globalelemtocoord[:, :, 1] = [2 3 2 3; 4 4 5 5]\n globalelemtocoord[:, :, 2] = [2 3 2 3; 5 5 6 6]\n globalelemtocoord[:, :, 3] = [3 4 3 4; 4 4 5 5]\n globalelemtocoord[:, :, 4] = [3 4 3 4; 5 5 6 6]\n globalelemtocoord[:, :, 5] = [4 5 4 5; 4 4 5 5]\n globalelemtocoord[:, :, 6] = [4 5 4 5; 5 5 6 6]\n\n\n globalelemtobndy = [\n 1 1 0 0 0 0\n 0 0 0 0 2 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n ]\n\n if crank == 0\n nrealelem = 2\n globalelems = [1, 2, 3, 4]\n elemtoelem_expect = [\n 1 2 3 4\n 3 4 3 4\n 2 1 3 4\n 2 1 3 4\n ]\n nabrtorank_expect = [1]\n nabrtorecv_expect = UnitRange{Int}[1:2]\n nabrtosend_expect = UnitRange{Int}[1:2]\n elseif crank == 1\n nrealelem = 2\n globalelems = [3, 4, 1, 2, 5, 6]\n elemtoelem_expect = [\n 3 4 1 2 5 6\n 5 6 3 4 1 2\n 2 1 3 4 5 6\n 2 1 3 4 5 6\n ]\n nabrtorank_expect = [0, 2]\n nabrtorecv_expect = UnitRange{Int}[1:2, 3:4]\n nabrtosend_expect = UnitRange{Int}[1:2, 3:4]\n elseif crank == 2\n nrealelem = 2\n globalelems = [5, 6, 3, 4]\n elemtoelem_expect = [\n 3 4 3 4\n 1 2 3 4\n 2 1 3 4\n 2 1 3 4\n ]\n nabrtorank_expect = [1]\n nabrtorecv_expect = UnitRange{Int}[1:2]\n nabrtosend_expect = UnitRange{Int}[1:2]\n end\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n @test elemtocoord == globalelemtocoord[:, :, globalelems]\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connect_stacked_3d.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction main()\n MPI.Init()\n\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 2\n\n topology = StackedBrickTopology(\n comm,\n (1:4, 5:8, 9:12),\n periodicity = (false, true, false),\n boundary = ((1, 2), (3, 4), (5, 6)),\n connectivity = :face,\n )\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoface = [\n 1 1 1 2 2 2 2 2 2 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 2 2 2 2 2 2 2 2 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n 5 6 6 5 6 6 5 6 6 5 6 6 5 6 6 5 6 6 5 6 6 5 6 6 5 6 6\n 5 5 6 5 5 6 5 5 6 5 5 6 5 5 6 5 5 6 5 5 6 5 5 6 5 5 6\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n\n globalelemtocoord = Array{Int}(undef, 3, 8, 27)\n\n globalelemtocoord[:, :, 1] =\n [1 2 1 2 1 2 1 2; 5 5 6 6 5 5 6 6; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 2] =\n [1 2 1 2 1 2 1 2; 5 5 6 6 5 5 6 6; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 3] =\n [1 2 1 2 1 2 1 2; 5 5 6 6 5 5 6 6; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 4] =\n [2 3 2 3 2 3 2 3; 5 5 6 6 5 5 6 6; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 5] =\n [2 3 2 3 2 3 2 3; 5 5 6 6 5 5 6 6; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 6] =\n [2 3 2 3 2 3 2 3; 5 5 6 6 5 5 6 6; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 7] =\n [2 3 2 3 2 3 2 3; 6 6 7 7 6 6 7 7; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 8] =\n [2 3 2 3 2 3 2 3; 6 6 7 7 6 6 7 7; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 9] =\n [2 3 2 3 2 3 2 3; 6 6 7 7 6 6 7 7; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 10] =\n [1 2 1 2 1 2 1 2; 6 6 7 7 6 6 7 7; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 11] =\n [1 2 1 2 1 2 1 2; 6 6 7 7 6 6 7 7; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 12] =\n [1 2 1 2 1 2 1 2; 6 6 7 7 6 6 7 7; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 13] =\n [1 2 1 2 1 2 1 2; 7 7 8 8 7 7 8 8; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 14] =\n [1 2 1 2 1 2 1 2; 7 7 8 8 7 7 8 8; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 15] =\n [1 2 1 2 1 2 1 2; 7 7 8 8 7 7 8 8; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 16] =\n [2 3 2 3 2 3 2 3; 7 7 8 8 7 7 8 8; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 17] =\n [2 3 2 3 2 3 2 3; 7 7 8 8 7 7 8 8; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 18] =\n [2 3 2 3 2 3 2 3; 7 7 8 8 7 7 8 8; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 19] =\n [3 4 3 4 3 4 3 4; 7 7 8 8 7 7 8 8; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 20] =\n [3 4 3 4 3 4 3 4; 7 7 8 8 7 7 8 8; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 21] =\n [3 4 3 4 3 4 3 4; 7 7 8 8 7 7 8 8; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 22] =\n [3 4 3 4 3 4 3 4; 6 6 7 7 6 6 7 7; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 23] =\n [3 4 3 4 3 4 3 4; 6 6 7 7 6 6 7 7; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 24] =\n [3 4 3 4 3 4 3 4; 6 6 7 7 6 6 7 7; 11 11 11 11 12 12 12 12]\n globalelemtocoord[:, :, 25] =\n [3 4 3 4 3 4 3 4; 5 5 6 6 5 5 6 6; 9 9 9 9 10 10 10 10]\n globalelemtocoord[:, :, 26] =\n [3 4 3 4 3 4 3 4; 5 5 6 6 5 5 6 6; 10 10 10 10 11 11 11 11]\n globalelemtocoord[:, :, 27] =\n [3 4 3 4 3 4 3 4; 5 5 6 6 5 5 6 6; 11 11 11 11 12 12 12 12]\n\n\n globalelemtobndy = [\n 1 1 1 0 0 0 0 0 0 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 2 2 2 2 2 2 2 2 2\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n 5 0 0 5 0 0 5 0 0 5 0 0 5 0 0 5 0 0 5 0 0 5 0 0 5 0 0\n 0 0 6 0 0 6 0 0 6 0 0 6 0 0 6 0 0 6 0 0 6 0 0 6 0 0 6\n ]\n\n if crank == 0\n nrealelem = 12\n globalelems = [\n 1,\n 2,\n 3, # 1\n 4,\n 5,\n 6, # 2\n 7,\n 8,\n 9, # 3\n 10,\n 11,\n 12, # 4\n 13,\n 14,\n 15, # 5\n 16,\n 17,\n 18, # 6\n 22,\n 23,\n 24, # 8\n 25,\n 26,\n 27,\n ] # 9\n\n elemtoelem_expect = [\n 1 2 3 1 2 3 10 11 12 4 5 6 7 8 9 16 17 18 19 20 21 22 23 24\n 4 5 6 22 23 24 19 20 21 7 8 9 13 14 15 16 17 18 1 2 3 4 5 6\n 13 14 15 16 17 18 4 5 6 1 2 3 13 14 15 16 17 18 19 20 21 22 23 24\n 10 11 12 7 8 9 16 17 18 13 14 15 13 14 15 16 17 18 19 20 21 22 23 24\n 1 1 2 2 4 5 3 7 8 4 10 11 5 14 15 6 17 18 7 20 21 8 23 24\n 2 3 1 5 6 2 8 9 3 11 12 4 13 14 5 16 17 6 19 20 7 22 23 8\n ]\n\n nabrtorank_expect = [1]\n nabrtorecv_expect = UnitRange{Int}[1:12]\n nabrtosend_expect = UnitRange{Int}[1:12]\n elseif crank == 1\n nrealelem = 15\n\n globalelems = [\n 13,\n 14,\n 15, # 5\n 16,\n 17,\n 18, # 6\n 19,\n 20,\n 21, # 7\n 22,\n 23,\n 24, # 8\n 25,\n 26,\n 27, # 9\n 1,\n 2,\n 3, # 1\n 4,\n 5,\n 6, # 2\n 7,\n 8,\n 9, # 3\n 10,\n 11,\n 12,\n ] # 4\n\n elemtoelem_expect = [\n 1 2 3 1 2 3 4 5 6 22 23 24 19 20 21 4 5 6 19 20 21 22 23 24 7 8 9\n 4 5 6 7 8 9 1 2 3 4 5 6 7 8 9 16 17 18 19 20 21 22 23 24 25 26 27\n 25 26 27 22 23 24 10 11 12 13 14 15 7 8 9 16 17 18 19 20 21 22 23 24 25 26 27\n 16 17 18 19 20 21 13 14 15 7 8 9 10 11 12 16 17 18 19 20 21 22 23 24 25 26 27\n 1 1 2 2 4 5 3 7 8 4 10 11 5 13 14 6 17 18 7 20 21 8 23 24 9 26 27\n 2 3 1 5 6 2 8 9 3 11 12 4 14 15 5 16 17 6 19 20 7 22 23 8 25 26 9\n ]\n\n nabrtorank_expect = [0]\n nabrtorecv_expect = UnitRange{Int}[1:12]\n nabrtosend_expect = UnitRange{Int}[1:12]\n end\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n\n @test elemtocoord == globalelemtocoord[:, :, globalelems]\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_connectfull.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.Topologies\n\nfunction test_connectmeshfull()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 3\n\n topology = BrickTopology(\n comm,\n (0:4, 5:9);\n boundary = ((1, 2), (3, 4)),\n periodicity = (false, true),\n connectivity = :full,\n )\n\n elems = topology.elems\n realelems = topology.realelems\n ghostelems = topology.ghostelems\n sendelems = topology.sendelems\n elemtocoord = topology.elemtocoord\n elemtoelem = topology.elemtoelem\n elemtoface = topology.elemtoface\n elemtoordr = topology.elemtoordr\n elemtobndy = topology.elemtobndy\n nabrtorank = topology.nabrtorank\n nabrtorecv = topology.nabrtorecv\n nabrtosend = topology.nabrtosend\n\n globalelemtoface = [\n 1 2 2 1 1 1 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n globalelemtoordr = ones(Int, size(globalelemtoface))\n\n globalelemtocoord = Array{Int}(undef, 2, 4, 16)\n globalelemtocoord[:, :, 1] = [0 1 0 1; 5 5 6 6]\n globalelemtocoord[:, :, 2] = [1 2 1 2; 5 5 6 6]\n globalelemtocoord[:, :, 3] = [1 2 1 2; 6 6 7 7]\n globalelemtocoord[:, :, 4] = [0 1 0 1; 6 6 7 7]\n globalelemtocoord[:, :, 5] = [0 1 0 1; 7 7 8 8]\n globalelemtocoord[:, :, 6] = [0 1 0 1; 8 8 9 9]\n globalelemtocoord[:, :, 7] = [1 2 1 2; 8 8 9 9]\n globalelemtocoord[:, :, 8] = [1 2 1 2; 7 7 8 8]\n globalelemtocoord[:, :, 9] = [2 3 2 3; 7 7 8 8]\n globalelemtocoord[:, :, 10] = [2 3 2 3; 8 8 9 9]\n globalelemtocoord[:, :, 11] = [3 4 3 4; 8 8 9 9]\n globalelemtocoord[:, :, 12] = [3 4 3 4; 7 7 8 8]\n globalelemtocoord[:, :, 13] = [3 4 3 4; 6 6 7 7]\n globalelemtocoord[:, :, 14] = [2 3 2 3; 6 6 7 7]\n globalelemtocoord[:, :, 15] = [2 3 2 3; 5 5 6 6]\n globalelemtocoord[:, :, 16] = [3 4 3 4; 5 5 6 6]\n\n globalelemtobndy = [\n 1 0 0 1 1 1 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 2 2 2 0 0 2\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n ]\n\n if crank == 0\n nrealelem = 5\n globalelems = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 14, 15]\n elemtoelem_expect = [\n 1 1 4 2 3 4 6 5 8 7 3 2\n 2 12 11 3 8 7 10 9 9 10 11 12\n 6 7 2 1 4 5 8 3 11 9 12 10\n 4 3 8 5 6 1 2 7 10 12 9 11\n ]\n\n elemtoface_expect = [\n 1 2 2 1 1 1 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 2 2 2 2\n 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n nabrtorank_expect = [1, 2]\n nabrtorecv_expect = UnitRange{Int}[1:5, 6:7]\n nabrtosend_expect = UnitRange{Int}[1:5, 6:7]\n\n elseif crank == 1\n nrealelem = 5\n globalelems = [6, 7, 8, 9, 10, 1, 2, 3, 4, 5, 11, 12, 13, 14, 15, 16]\n elemtoelem_expect = [\n 1 1 10 3 2 2 6 9 3 4 5 4 14 8 7 15\n 2 5 4 12 11 7 15 14 8 3 1 2 3 13 16 4\n 10 3 8 14 4 1 2 7 6 9 12 13 16 15 5 11\n 6 7 2 5 15 9 8 3 10 1 16 11 12 4 14 13\n ]\n elemtoface_expect = [\n 1 2 2 2 2 1 2 2 1 1 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n nabrtorank_expect = [0, 2]\n nabrtorecv_expect = UnitRange{Int}[1:5, 6:11]\n nabrtosend_expect = UnitRange{Int}[1:5, 6:9]\n\n elseif crank == 2\n nrealelem = 6\n globalelems = [11, 12, 13, 14, 15, 16, 2, 3, 7, 8, 9, 10]\n elemtoelem_expect = [\n 12 11 4 8 7 5 7 8 9 10 10 9\n 1 2 3 3 6 4 5 4 12 11 2 1\n 2 3 6 5 12 1 9 7 10 8 4 11\n 6 1 2 11 4 3 8 10 7 9 12 5\n ]\n elemtoface_expect = [\n 2 2 2 2 2 2 1 1 1 1 2 2\n 2 2 2 1 1 2 1 1 1 1 1 1\n 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n nabrtorank_expect = [0, 1]\n nabrtorecv_expect = UnitRange{Int}[1:2, 3:6]\n nabrtosend_expect = UnitRange{Int}[1:2, 3:8]\n end\n\n @test elems == 1:length(globalelems)\n @test realelems == 1:nrealelem\n @test ghostelems == (nrealelem + 1):length(globalelems)\n @test elemtocoord == globalelemtocoord[:, :, globalelems]\n @test elemtoface[:, realelems] ==\n globalelemtoface[:, globalelems[realelems]]\n @test elemtoelem == elemtoelem_expect\n @test elemtobndy == globalelemtobndy[:, globalelems]\n @test elemtoordr == ones(eltype(elemtoordr), size(elemtoordr))\n @test nabrtorank == nabrtorank_expect\n @test nabrtorecv == nabrtorecv_expect\n @test nabrtosend == nabrtosend_expect\n\n @test collect(realelems) ==\n sort(union(topology.exteriorelems, topology.interiorelems))\n @test unique(sort(sendelems)) == topology.exteriorelems\n @test length(intersect(topology.exteriorelems, topology.interiorelems)) == 0\nend\n\ntest_connectmeshfull()\n" }, { "path": "test/Numerics/Mesh/mpi_getpartition.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.BrickMesh\nusing Random\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n Nelemtotal = 113\n Random.seed!(1234)\n globalcode = randperm(Nelemtotal)\n\n @assert csize > 1\n\n bs = [\n (i == 1) ? (1:0) :\n BrickMesh.linearpartition(Nelemtotal, i - 1, csize - 1)\n for i in 1:csize\n ]\n as = [BrickMesh.linearpartition(Nelemtotal, i, csize) for i in 1:csize]\n\n codeb = globalcode[bs[crank + 1]]\n\n (so, ss, rs) = BrickMesh.getpartition(comm, codeb)\n\n codeb = codeb[so]\n\n\n codec = []\n for r in 0:(csize - 1)\n sendrange = ss[r + 1]:(ss[r + 2] - 1)\n rcounts = MPI.Gather(Cint(length(sendrange)), r, comm)\n c = MPI.Gatherv!(\n view(codeb, sendrange),\n crank == r ? VBuffer(similar(codeb, sum(rcounts)), rcounts) :\n nothing,\n r,\n comm,\n )\n if r == crank\n codec = c\n end\n end\n\n codea = (1:Nelemtotal)[as[crank + 1]]\n\n @test sort(codec) == codea\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_partition.jl", "content": "using Test\nusing MPI\nusing ClimateMachine.Mesh.BrickMesh\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n crank = MPI.Comm_rank(comm)\n csize = MPI.Comm_size(comm)\n\n @assert csize == 3\n\n (etv, etc, etb, fc) = brickmesh(\n (0:4, 5:9),\n (false, true),\n boundary = ((1, 2), (3, 4), (5, 6)),\n part = crank + 1,\n numparts = csize,\n )\n (etv, etc, etb, fc) = partition(comm, etv, etc, etb, fc)\n\n if crank == 0\n etv_expect = [\n 1 2 7 6 11\n 2 3 8 7 12\n 6 7 12 11 16\n 7 8 13 12 17\n ]\n @test etv == etv_expect\n @test etc[:, :, 1] == [0 1 0 1; 5 5 6 6]\n @test etc[:, :, 2] == [1 2 1 2; 5 5 6 6]\n @test etc[:, :, 3] == [1 2 1 2; 6 6 7 7]\n @test etc[:, :, 4] == [0 1 0 1; 6 6 7 7]\n @test etc[:, :, 5] == [0 1 0 1; 7 7 8 8]\n etb_expect = [\n 1 0 0 1 1\n 0 0 0 0 0\n 0 0 0 0 0\n 0 0 0 0 0\n ]\n @test etb == etb_expect\n fc_expect = Array{Int64, 1}[]\n @test fc == fc_expect\n elseif crank == 1\n etv_expect = [\n 16 17 12 13 18\n 17 18 13 14 19\n 21 22 17 18 23\n 22 23 18 19 24\n ]\n @test etv == etv_expect\n @test etc[:, :, 1] == [0 1 0 1; 8 8 9 9]\n @test etc[:, :, 2] == [1 2 1 2; 8 8 9 9]\n @test etc[:, :, 3] == [1 2 1 2; 7 7 8 8]\n @test etc[:, :, 4] == [2 3 2 3; 7 7 8 8]\n @test etc[:, :, 5] == [2 3 2 3; 8 8 9 9]\n etb_expect = [\n 1 0 0 0 0\n 0 0 0 0 0\n 0 0 0 0 0\n 0 0 0 0 0\n ]\n @test etb == etb_expect\n fc_expect = Array{Int64, 1}[[1, 4, 1, 2], [2, 4, 2, 3], [5, 4, 3, 4]]\n @test fc == fc_expect\n elseif crank == 2\n etv_expect = [\n 19 14 9 8 3 4\n 20 15 10 9 4 5\n 24 19 14 13 8 9\n 25 20 15 14 9 10\n ]\n @test etv == etv_expect\n @test etc[:, :, 1] == [3 4 3 4; 8 8 9 9]\n @test etc[:, :, 2] == [3 4 3 4; 7 7 8 8]\n @test etc[:, :, 3] == [3 4 3 4; 6 6 7 7]\n @test etc[:, :, 4] == [2 3 2 3; 6 6 7 7]\n @test etc[:, :, 5] == [2 3 2 3; 5 5 6 6]\n @test etc[:, :, 6] == [3 4 3 4; 5 5 6 6]\n etb_expect = [\n 0 0 0 0 0 0\n 2 2 2 0 0 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n ]\n @test etb == etb_expect\n fc_expect = Array{Int64, 1}[[1, 4, 4, 5]]\n @test fc == fc_expect\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/mpi_sortcolumns.jl", "content": "using Random\nusing Test\nusing MPI\nusing ClimateMachine.Mesh.BrickMesh\n\nfunction main()\n MPI.Init()\n comm = MPI.COMM_WORLD\n rank = MPI.Comm_rank(comm)\n\n Random.seed!(1234)\n d = 4\n A = rand(1:10, d, 3rank)\n B = BrickMesh.parallelsortcolumns(comm, A, rev = true)\n\n root = 0\n Acounts = MPI.Gather(Cint(length(A)), root, comm)\n A_all = MPI.Gatherv!(\n A,\n MPI.Comm_rank(comm) == root ?\n VBuffer(similar(A, sum(Acounts)), Acounts) : nothing,\n root,\n comm,\n )\n\n Bcounts = MPI.Gather(Cint(length(B)), root, comm)\n B_all = MPI.Gatherv!(\n B,\n MPI.Comm_rank(comm) == root ?\n VBuffer(similar(B, sum(Bcounts)), Bcounts) : nothing,\n root,\n comm,\n )\n\n if MPI.Comm_rank(comm) == root\n A_all = reshape(A_all, d, div(length(A_all), d))\n B_all = reshape(B_all, d, div(length(B_all), d))\n\n A_all = sortslices(A_all, dims = 2, rev = true)\n\n @test A_all == B_all\n end\nend\n\nmain()\n" }, { "path": "test/Numerics/Mesh/topology.jl", "content": "using Test\nusing ClimateMachine.Mesh.Topologies\nusing Combinatorics, MPI\n\nMPI.Initialized() || MPI.Init()\n\n@testset \"Equiangular cubed_sphere_warp tests\" begin\n import ClimateMachine.Mesh.Topologies: equiangular_cubed_sphere_warp\n\n # Create function alias for shorter formatting\n eacsw = equiangular_cubed_sphere_warp\n\n @testset \"check radius\" begin\n @test hypot(eacsw(3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(eacsw(-3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(eacsw(1.1, -2.2, 3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(eacsw(1.1, -2.2, -3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(eacsw(1.1, 3.0, 0.0)...) ≈ 3.0 rtol = eps()\n @test hypot(eacsw(1.1, -3.0, 0.0)...) ≈ 3.0 rtol = eps()\n end\n\n @testset \"check sign\" begin\n @test sign.(eacsw(3.0, -2.2, 1.3)) == sign.((3.0, -2.2, 1.3))\n @test sign.(eacsw(-3.0, -2.2, 1.3)) == sign.((-3.0, -2.2, 1.3))\n @test sign.(eacsw(1.1, -2.2, 3.0)) == sign.((1.1, -2.2, 3.0))\n @test sign.(eacsw(1.1, -2.2, -3.0)) == sign.((1.1, -2.2, -3.0))\n @test sign.(eacsw(1.1, 3.0, 0.0)) == sign.((1.1, 3.0, 0.0))\n @test sign.(eacsw(1.1, -3.0, 0.0)) == sign.((1.1, -3.0, 0.0))\n end\n\n @testset \"check continuity\" begin\n for (u, v) in zip(\n permutations([3.0, 2.999999999, 1.3]),\n permutations([2.999999999, 3.0, 1.3]),\n )\n @test all(eacsw(u...) .≈ eacsw(v...))\n end\n for (u, v) in zip(\n permutations([3.0, -2.999999999, 1.3]),\n permutations([2.999999999, -3.0, 1.3]),\n )\n @test all(eacsw(u...) .≈ eacsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, 2.999999999, 1.3]),\n permutations([-2.999999999, 3.0, 1.3]),\n )\n @test all(eacsw(u...) .≈ eacsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, -2.999999999, 1.3]),\n permutations([-2.999999999, -3.0, 1.3]),\n )\n @test all(eacsw(u...) .≈ eacsw(v...))\n end\n end\nend\n\n@testset \"Equiangular cubed_sphere_unwarp tests\" begin\n import ClimateMachine.Mesh.Topologies:\n cubed_sphere_warp, equiangular_cubed_sphere_unwarp\n\n # Create function aliases for shorter formatting\n eacsw = equiangular_cubed_sphere_warp\n eacsu = equiangular_cubed_sphere_unwarp\n\n for u in permutations([3.0, 2.999999999, 1.3])\n @test all(eacsu(eacsw(u...)...) .≈ u)\n end\n for u in permutations([3.0, -2.999999999, 1.3])\n @test all(eacsu(eacsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, 2.999999999, 1.3])\n @test all(eacsu(eacsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, -2.999999999, 1.3])\n @test all(eacsu(eacsw(u...)...) .≈ u)\n end\nend\n\n@testset \"Equidistant cubed_sphere_warp tests\" begin\n import ClimateMachine.Mesh.Topologies: equidistant_cubed_sphere_warp\n\n # Create function alias for shorter formatting\n edcsw = equidistant_cubed_sphere_warp\n\n @testset \"check radius\" begin\n @test hypot(edcsw(3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(edcsw(-3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(edcsw(1.1, -2.2, 3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(edcsw(1.1, -2.2, -3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(edcsw(1.1, 3.0, 0.0)...) ≈ 3.0 rtol = eps()\n @test hypot(edcsw(1.1, -3.0, 0.0)...) ≈ 3.0 rtol = eps()\n end\n\n @testset \"check sign\" begin\n @test sign.(edcsw(3.0, -2.2, 1.3)) == sign.((3.0, -2.2, 1.3))\n @test sign.(edcsw(-3.0, -2.2, 1.3)) == sign.((-3.0, -2.2, 1.3))\n @test sign.(edcsw(1.1, -2.2, 3.0)) == sign.((1.1, -2.2, 3.0))\n @test sign.(edcsw(1.1, -2.2, -3.0)) == sign.((1.1, -2.2, -3.0))\n @test sign.(edcsw(1.1, 3.0, 0.0)) == sign.((1.1, 3.0, 0.0))\n @test sign.(edcsw(1.1, -3.0, 0.0)) == sign.((1.1, -3.0, 0.0))\n end\n\n @testset \"check continuity\" begin\n for (u, v) in zip(\n permutations([3.0, 2.999999999, 1.3]),\n permutations([2.999999999, 3.0, 1.3]),\n )\n @test all(edcsw(u...) .≈ edcsw(v...))\n end\n for (u, v) in zip(\n permutations([3.0, -2.999999999, 1.3]),\n permutations([2.999999999, -3.0, 1.3]),\n )\n @test all(edcsw(u...) .≈ edcsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, 2.999999999, 1.3]),\n permutations([-2.999999999, 3.0, 1.3]),\n )\n @test all(edcsw(u...) .≈ edcsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, -2.999999999, 1.3]),\n permutations([-2.999999999, -3.0, 1.3]),\n )\n @test all(edcsw(u...) .≈ edcsw(v...))\n end\n end\nend\n\n@testset \"Equidistant cubed_sphere_unwarp tests\" begin\n import ClimateMachine.Mesh.Topologies:\n equidistant_cubed_sphere_warp, equidistant_cubed_sphere_unwarp\n\n # Create function aliases for shorter formatting\n edcsw = equidistant_cubed_sphere_warp\n edcsu = equidistant_cubed_sphere_unwarp\n\n for u in permutations([3.0, 2.999999999, 1.3])\n @test all(edcsu(edcsw(u...)...) .≈ u)\n end\n for u in permutations([3.0, -2.999999999, 1.3])\n @test all(edcsu(edcsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, 2.999999999, 1.3])\n @test all(edcsu(edcsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, -2.999999999, 1.3])\n @test all(edcsu(edcsw(u...)...) .≈ u)\n end\nend\n\n@testset \"Conformal cubed_sphere_warp tests\" begin\n import ClimateMachine.Mesh.Topologies: conformal_cubed_sphere_warp\n\n # Create function alias for shorter formatting\n ccsw = conformal_cubed_sphere_warp\n\n @testset \"check radius\" begin\n @test hypot(ccsw(3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(ccsw(-3.0, -2.2, 1.3)...) ≈ 3.0 rtol = eps()\n @test hypot(ccsw(1.1, -2.2, 3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(ccsw(1.1, -2.2, -3.0)...) ≈ 3.0 rtol = eps()\n @test hypot(ccsw(1.1, 3.0, 0.0)...) ≈ 3.0 rtol = eps()\n @test hypot(ccsw(1.1, -3.0, 0.0)...) ≈ 3.0 rtol = eps()\n end\n\n @testset \"check sign\" begin\n @test sign.(ccsw(3.0, -2.2, 1.3)) == sign.((3.0, -2.2, 1.3))\n @test sign.(ccsw(-3.0, -2.2, 1.3)) == sign.((-3.0, -2.2, 1.3))\n @test sign.(ccsw(1.1, -2.2, 3.0)) == sign.((1.1, -2.2, 3.0))\n @test sign.(ccsw(1.1, -2.2, -3.0)) == sign.((1.1, -2.2, -3.0))\n @test sign.(ccsw(1.1, 3.0, -2.2)) == sign.((1.1, 3.0, -2.2))\n @test sign.(ccsw(1.1, -3.0, -2.2)) == sign.((1.1, -3.0, -2.2))\n end\n\n @testset \"check continuity\" begin\n for (u, v) in zip(\n permutations([3.0, 2.999999999, 1.3]),\n permutations([2.999999999, 3.0, 1.3]),\n )\n @test all(ccsw(u...) .≈ ccsw(v...))\n end\n for (u, v) in zip(\n permutations([3.0, -2.999999999, 1.3]),\n permutations([2.999999999, -3.0, 1.3]),\n )\n @test all(ccsw(u...) .≈ ccsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, 2.999999999, 1.3]),\n permutations([-2.999999999, 3.0, 1.3]),\n )\n @test all(ccsw(u...) .≈ ccsw(v...))\n end\n for (u, v) in zip(\n permutations([-3.0, -2.999999999, 1.3]),\n permutations([-2.999999999, -3.0, 1.3]),\n )\n @test all(ccsw(u...) .≈ ccsw(v...))\n end\n end\nend\n\n@testset \"Conformal cubed_sphere_unwarp tests\" begin\n import ClimateMachine.Mesh.Topologies:\n conformal_cubed_sphere_warp, conformal_cubed_sphere_unwarp\n\n # Create function aliases for shorter formatting\n ccsw = conformal_cubed_sphere_warp\n ccsu = conformal_cubed_sphere_unwarp\n\n for u in permutations([3.0, 2.999999999, 1.3])\n @test all(ccsu(ccsw(u...)...) .≈ u)\n end\n for u in permutations([3.0, -2.999999999, 1.3])\n @test all(ccsu(ccsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, 2.999999999, 1.3])\n @test all(ccsu(ccsw(u...)...) .≈ u)\n end\n for u in permutations([-3.0, -2.999999999, 1.3])\n @test all(ccsu(ccsw(u...)...) .≈ u)\n end\nend\n\n@testset \"grid1d\" begin\n g = grid1d(0, 10, nelem = 10)\n @test eltype(g) == Float64\n @test length(g) == 11\n @test g[1] == 0\n @test g[end] == 10\n\n g = grid1d(10.0f0, 20.0f0, elemsize = 0.1)\n @test eltype(g) == Float32\n @test length(g) == 101\n @test g[1] == 10\n @test g[end] == 20\n\n g = grid1d(10.0f0, 20.0f0, InteriorStretching(0), elemsize = 0.1)\n @test eltype(g) == Float32\n @test length(g) == 101\n @test g[1] == 10\n @test g[end] == 20\n\n g = grid1d(\n 10.0f0,\n 20.0f0,\n SingleExponentialStretching(2.5f0),\n elemsize = 0.1,\n )\n @test eltype(g) == Float32\n @test length(g) == 101\n @test g[1] == 10\n @test g[end] == 20\nend\n\n@testset \"BrickTopology tests\" begin\n\n let\n comm = MPI.COMM_SELF\n\n\n elemrange = (0:10,)\n periodicity = (true,)\n\n topology = BrickTopology(\n comm,\n elemrange,\n periodicity = periodicity,\n connectivity = :face,\n )\n\n nelem = length(elemrange[1]) - 1\n\n for e in 1:nelem\n @test topology.elemtocoord[:, :, e] == [e - 1 e]\n end\n\n @test topology.elemtoelem ==\n [nelem collect(1:(nelem - 1))'; collect(2:nelem)' 1]\n @test topology.elemtoface == repeat(2:-1:1, outer = (1, nelem))\n\n @test topology.elemtoordr == ones(Int, size(topology.elemtoordr))\n @test topology.elemtobndy == zeros(Int, size(topology.elemtoordr))\n\n @test topology.elems == 1:nelem\n @test topology.realelems == 1:nelem\n @test topology.ghostelems == nelem .+ (1:0)\n\n @test length(topology.sendelems) == 0\n @test length(topology.exteriorelems) == 0\n @test collect(topology.realelems) == topology.interiorelems\n\n @test topology.nabrtorank == Int[]\n @test topology.nabrtorecv == UnitRange{Int}[]\n @test topology.nabrtosend == UnitRange{Int}[]\n end\n\n let\n comm = MPI.COMM_SELF\n topology = BrickTopology(\n comm,\n (0:4, 5:9),\n periodicity = (false, true),\n connectivity = :face,\n )\n\n nelem = 16\n\n @test topology.elemtocoord[:, :, 1] == [0 1 0 1; 5 5 6 6]\n @test topology.elemtocoord[:, :, 2] == [1 2 1 2; 5 5 6 6]\n @test topology.elemtocoord[:, :, 3] == [1 2 1 2; 6 6 7 7]\n @test topology.elemtocoord[:, :, 4] == [0 1 0 1; 6 6 7 7]\n @test topology.elemtocoord[:, :, 5] == [0 1 0 1; 7 7 8 8]\n @test topology.elemtocoord[:, :, 6] == [0 1 0 1; 8 8 9 9]\n @test topology.elemtocoord[:, :, 7] == [1 2 1 2; 8 8 9 9]\n @test topology.elemtocoord[:, :, 8] == [1 2 1 2; 7 7 8 8]\n @test topology.elemtocoord[:, :, 9] == [2 3 2 3; 7 7 8 8]\n @test topology.elemtocoord[:, :, 10] == [2 3 2 3; 8 8 9 9]\n @test topology.elemtocoord[:, :, 11] == [3 4 3 4; 8 8 9 9]\n @test topology.elemtocoord[:, :, 12] == [3 4 3 4; 7 7 8 8]\n @test topology.elemtocoord[:, :, 13] == [3 4 3 4; 6 6 7 7]\n @test topology.elemtocoord[:, :, 14] == [2 3 2 3; 6 6 7 7]\n @test topology.elemtocoord[:, :, 15] == [2 3 2 3; 5 5 6 6]\n @test topology.elemtocoord[:, :, 16] == [3 4 3 4; 5 5 6 6]\n\n @test topology.elemtoelem == [\n 1 1 4 2 3 4 6 5 8 7 10 9 14 3 2 15\n 2 15 14 3 8 7 10 9 12 11 5 6 7 13 16 8\n 6 7 2 1 4 5 8 3 14 9 12 13 16 15 10 11\n 4 3 8 5 6 1 2 7 10 15 16 11 12 9 14 13\n ]\n\n @test topology.elemtoface == [\n 1 2 2 1 1 1 2 2 2 2 2 2 2 2 2 2\n 1 1 1 1 1 1 1 1 1 1 2 2 2 1 1 2\n 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4 4\n 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3\n ]\n\n @test topology.elemtoordr == ones(Int, size(topology.elemtoordr))\n\n @test topology.elemtoelem[topology.elemtobndy .== 1] == 1:8\n\n @test topology.elems == 1:nelem\n @test topology.realelems == 1:nelem\n @test topology.ghostelems == nelem .+ (1:0)\n\n @test length(topology.sendelems) == 0\n @test length(topology.exteriorelems) == 0\n @test collect(topology.realelems) == topology.interiorelems\n\n @test topology.nabrtorank == Int[]\n @test topology.nabrtorecv == UnitRange{Int}[]\n @test topology.nabrtosend == UnitRange{Int}[]\n end\n\n let\n comm = MPI.COMM_SELF\n for px in (true, false)\n topology = BrickTopology(\n comm,\n (0:10,),\n periodicity = (px,),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology) == !px\n if px\n @test topology.bndytoelem == ()\n @test topology.bndytoface == ()\n else\n @test topology.bndytoelem == ([1, 10],)\n @test topology.bndytoface == ([1, 2],)\n end\n end\n for py in (true, false), px in (true, false)\n topology = BrickTopology(\n comm,\n (0:10, 0:3),\n periodicity = (px, py),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology) == !(px && py)\n if px && py\n @test topology.bndytoelem == ()\n @test topology.bndytoface == ()\n else\n @test sort(unique(topology.bndytoface[1])) == vcat(\n px ? Int64[] : Int64[1, 2],\n py ? Int64[] : Int64[3, 4],\n )\n end\n end\n for pz in (true, false), py in (true, false), px in (true, false)\n topology = BrickTopology(\n comm,\n (0:10, 0:3, -1:3),\n periodicity = (px, py, pz),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology) == !(px && py && pz)\n if px && py && pz\n @test topology.bndytoelem == ()\n @test topology.bndytoface == ()\n else\n @test sort(unique(topology.bndytoface[1])) == vcat(\n px ? Int64[] : Int64[1, 2],\n py ? Int64[] : Int64[3, 4],\n pz ? Int64[] : Int64[5, 6],\n )\n end\n end\n end\nend\n\n@testset \"StackedBrickTopology tests\" begin\n let\n comm = MPI.COMM_SELF\n topology = StackedBrickTopology(\n comm,\n (2:5, 4:6),\n periodicity = (false, true),\n boundary = ((1, 2), (3, 4)),\n connectivity = :face,\n )\n\n nelem = 6\n\n @test topology.elemtocoord[:, :, 1] == [2 3 2 3; 4 4 5 5]\n @test topology.elemtocoord[:, :, 2] == [2 3 2 3; 5 5 6 6]\n @test topology.elemtocoord[:, :, 3] == [3 4 3 4; 4 4 5 5]\n @test topology.elemtocoord[:, :, 4] == [3 4 3 4; 5 5 6 6]\n @test topology.elemtocoord[:, :, 5] == [4 5 4 5; 4 4 5 5]\n @test topology.elemtocoord[:, :, 6] == [4 5 4 5; 5 5 6 6]\n\n @test topology.elemtoelem == [\n 1 2 1 2 3 4\n 3 4 5 6 1 2\n 2 1 4 3 6 5\n 2 1 4 3 6 5\n ]\n\n @test topology.elemtoface == [\n 1 1 2 2 2 2\n 1 1 1 1 2 2\n 4 4 4 4 4 4\n 3 3 3 3 3 3\n ]\n\n @test topology.elemtoordr == ones(Int, size(topology.elemtoordr))\n\n @test topology.elemtobndy == [\n 1 1 0 0 0 0\n 0 0 0 0 2 2\n 0 0 0 0 0 0\n 0 0 0 0 0 0\n ]\n\n @test topology.elemtoelem[topology.elemtobndy .== 1] == 1:2\n @test topology.elemtoelem[topology.elemtobndy .== 2] == 1:2\n\n @test topology.elems == 1:nelem\n @test topology.realelems == 1:nelem\n @test topology.ghostelems == nelem .+ (1:0)\n\n @test length(topology.sendelems) == 0\n @test length(topology.exteriorelems) == 0\n @test collect(topology.realelems) == topology.interiorelems\n\n @test topology.nabrtorank == Int[]\n @test topology.nabrtorecv == UnitRange{Int}[]\n @test topology.nabrtosend == UnitRange{Int}[]\n\n @test topology.bndytoelem == ([1, 2], [5, 6])\n @test topology.bndytoface == ([1, 1], [2, 2])\n end\n let\n comm = MPI.COMM_SELF\n for py in (true, false), px in (true, false)\n topology = StackedBrickTopology(\n comm,\n (0:10, 0:3),\n periodicity = (px, py),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology) == !(px && py)\n if px && py\n @test topology.bndytoelem == ()\n @test topology.bndytoface == ()\n else\n @test sort(unique(topology.bndytoface[1])) == vcat(\n px ? Int64[] : Int64[1, 2],\n py ? Int64[] : Int64[3, 4],\n )\n end\n end\n for pz in (true, false), py in (true, false), px in (true, false)\n topology = StackedBrickTopology(\n comm,\n (0:10, 0:3, -1:3),\n periodicity = (px, py, pz),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology) == !(px && py && pz)\n if px && py && pz\n @test topology.bndytoelem == ()\n @test topology.bndytoface == ()\n else\n @test sort(unique(topology.bndytoface[1])) == vcat(\n px ? Int64[] : Int64[1, 2],\n py ? Int64[] : Int64[3, 4],\n pz ? Int64[] : Int64[5, 6],\n )\n end\n end\n end\nend\n\n@testset \"StackedCubedSphereTopology tests\" begin\n topology = StackedCubedSphereTopology(\n MPI.COMM_SELF,\n 3,\n 1.0:3.0,\n boundary = (2, 1),\n connectivity = :face,\n )\n @test Topologies.hasboundary(topology)\n @test map(unique, topology.bndytoface) == ([6], [5])\nend\n\n@testset \"CubedShellTopology tests\" begin\n topology =\n CubedShellTopology(MPI.COMM_SELF, 3, Float64, connectivity = :face)\n @test !Topologies.hasboundary(topology)\nend\n" }, { "path": "test/Numerics/ODESolvers/callbacks.jl", "content": "using MPI\nusing Test\nusing ClimateMachine.ODESolvers: AbstractODESolver\nusing ClimateMachine.GenericCallbacks\n\nmutable struct PseudoSolver <: AbstractODESolver\n t::Float64\n steps::Int\n PseudoSolver() = new(42.0, 0)\nend\ngettime(ps::PseudoSolver) = ps.t\n\n\nmutable struct MyCallback\n initialized::Bool\n calls::Int\n finished::Bool\nend\nMyCallback() = MyCallback(false, 0, false)\n\nGenericCallbacks.init!(cb::MyCallback, _...) = cb.initialized = true\nGenericCallbacks.call!(cb::MyCallback, _...) = (cb.calls += 1; nothing)\nGenericCallbacks.fini!(cb::MyCallback, _...) = cb.finished = true\n\nMPI.Init()\nmpicomm = MPI.COMM_WORLD\nps = PseudoSolver()\n\nwtcb = GenericCallbacks.EveryXWallTimeSeconds(MyCallback(), 2, mpicomm)\nstcb = GenericCallbacks.EveryXSimulationTime(MyCallback(), 0.5)\nsscb = GenericCallbacks.EveryXSimulationSteps(MyCallback(), 10)\n\nfn_calls = 0\nfncb = () -> (global fn_calls += 1)\n\nwtfn_calls = 0\nwtfncb = GenericCallbacks.EveryXWallTimeSeconds(\n AtInit(() -> global wtfn_calls += 1),\n 2,\n mpicomm,\n)\n\nstfn_calls = 0\nstfncb = GenericCallbacks.EveryXSimulationTime(\n AtInitAndFini(() -> global stfn_calls += 1),\n 0.5,\n)\n\nssfn_calls = 0\nssfncb = GenericCallbacks.EveryXSimulationSteps(\n AtInit(() -> global ssfn_calls += 1),\n 5,\n)\n\ncallbacks = ((wtcb, stcb, sscb), fncb, (wtfncb, stfncb, ssfncb))\n\n@testset \"GenericCallbacks\" begin\n GenericCallbacks.init!(callbacks, ps, nothing, nothing, ps.t)\n\n @test wtcb.callback.initialized\n @test stcb.callback.initialized\n @test sscb.callback.initialized\n\n @test wtcb.callback.calls == 0\n @test stcb.callback.calls == 0\n @test sscb.callback.calls == 0\n @test fn_calls == 0\n @test wtfn_calls >= 1\n @test stfn_calls == 1\n @test ssfn_calls == 1\n\n @test GenericCallbacks.call!(callbacks, ps, nothing, nothing, ps.t) in\n (0, nothing)\n\n @test wtcb.callback.calls >= 0\n @test stcb.callback.calls == 0\n @test sscb.callback.calls == 0\n @test fn_calls == 1\n @test wtfn_calls >= 1\n @test stfn_calls == 1\n @test ssfn_calls == 1\n\n ps.t += 0.5\n @test GenericCallbacks.call!(callbacks, ps, nothing, nothing, ps.t) in\n (0, nothing)\n\n @test wtcb.callback.calls >= 0\n @test stcb.callback.calls == 1\n @test sscb.callback.calls == 0\n @test fn_calls == 2\n @test wtfn_calls >= 1\n @test stfn_calls == 2\n @test ssfn_calls == 1\n\n sleep(max((2.1 + wtcb.lastcbtime_ns / 1e9 - time_ns() / 1e9), 0.0))\n ps.t += 0.5\n @test GenericCallbacks.call!(callbacks, ps, nothing, nothing, ps.t) in\n (0, nothing)\n\n @test wtcb.callback.calls >= 1\n @test stcb.callback.calls == 2\n @test sscb.callback.calls == 0\n @test fn_calls == 3\n @test wtfn_calls >= 2\n @test stfn_calls == 3\n @test ssfn_calls == 1\n\n for i in 1:7\n @test GenericCallbacks.call!(callbacks, ps, nothing, nothing, ps.t) in\n (0, nothing)\n end\n\n @test wtcb.callback.calls >= 1\n @test stcb.callback.calls == 2\n @test sscb.callback.calls == 1\n @test fn_calls == 10\n @test wtfn_calls >= 2\n @test stfn_calls == 3\n @test ssfn_calls == 3\n\n GenericCallbacks.fini!(callbacks, ps, nothing, nothing, ps.t)\n @test wtcb.callback.finished\n @test stcb.callback.finished\n @test sscb.callback.finished\n @test wtfn_calls >= 2\n @test stfn_calls == 4\n @test ssfn_calls == 3\nend\n" }, { "path": "test/Numerics/ODESolvers/ode_tests_basic.jl", "content": "using Test\nusing ClimateMachine\nusing LinearAlgebra\nimport OrdinaryDiffEq: SSPRK73\n\ninclude(\"ode_tests_common.jl\")\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\n\na = 100\nb = 1\nc = 1 / 100\nΔ = sqrt(4 * a * c - b^2)\n\nα1, α2 = 1 / 4, 3 / 4\nβ1, β2, β3 = 1 / 3, 3 / 6, 1 / 6\n\nfunction rhs!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = $cos(t) * (a + b * Q + c * Q^2)\n end\nend\nfunction rhs_linear!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += $cos(t) * b * Q\n else\n @. dQ = $cos(t) * b * Q\n end\nend\nstruct ODETestBasicLinBE <: AbstractBackwardEulerSolver end\nODESolvers.Δt_is_adjustable(::ODETestBasicLinBE) = true\n(::ODETestBasicLinBE)(Q, Qhat, α, p, t) = @. Q = Qhat / (1 - α * $cos(t) * b)\nfunction rhs_nonlinear!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += $cos(t) * (a + c * Q^2)\n else\n @. dQ = $cos(t) * (a + c * Q^2)\n end\nend\nfunction rhs_fast!(dQ, Q, ::Any, t; increment = false)\n if increment\n @. dQ += α1 * $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = α1 * $cos(t) * (a + b * Q + c * Q^2)\n end\nend\nfunction rhs_fast_linear!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += α1 * $cos(t) * Q\n else\n @. dQ = α1 * $cos(t) * Q\n end\nend\nfunction rhs_slow!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += α2 * $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = α2 * $cos(t) * (a + b * Q + c * Q^2)\n end\nend\nfunction rhs1!(dQ, Q, ::Any, t; increment = false)\n if increment\n @. dQ += β1 * $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = β1 * $cos(t) * (a + b * Q + c * Q^2)\n end\nend\nfunction rhs2!(dQ, Q, ::Any, t; increment = false)\n if increment\n @. dQ += β2 * $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = β2 * $cos(t) * (a + b * Q + c * Q^2)\n end\nend\nfunction rhs3!(dQ, Q, ::Nothing, t; increment = false)\n if increment\n @. dQ += β3 * $cos(t) * (a + b * Q + c * Q^2)\n else\n @. dQ = β3 * $cos(t) * (a + b * Q + c * Q^2)\n end\nend\n\nfunction exactsolution(t, q0, t0)\n k = @. 2 * atan((2 * c * q0 + b) / Δ) / Δ - sin(t0)\n solution = @. (Δ * tan((k + sin(t)) * Δ / 2) - b) / (2 * c)\n return ArrayType(solution)\nend\n\nq0 = ArrayType === Array ? [1.0] : range(-1.0, 1.0, length = 303)\nt0 = 0.1\nfinaltime = 1.2\n\nQinit = exactsolution(t0, q0, t0)\nQ = similar(Qinit)\nQexact = exactsolution(finaltime, q0, t0)\n\n@testset \"Convergence/limited\" begin\n @testset \"Explicit methods\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (method, expected_order) in explicit_methods\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = method(rhs!, Q; dt = dt, t0 = t0)\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.7)\n end\n end\n\n @testset \"IMEX methods (LowStorageVariant)\" begin\n dts = [2.0^(-k) for k in 4:5]\n errors = similar(dts)\n for (method, order) in imex_methods_lowstorage_compatible\n for split_explicit_implicit in (false, true)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n rhs_arg! = split_explicit_implicit ? rhs_nonlinear! : rhs!\n solver = method(\n rhs_arg!,\n rhs_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver();\n isadjustable = true,\n ),\n Q;\n dt = dt,\n t0 = t0,\n split_explicit_implicit = split_explicit_implicit,\n variant = LowStorageVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n if split_explicit_implicit\n if method === ARK1ForwardBackwardEuler\n expected_order = 1\n else\n expected_order = 2\n end\n else\n expected_order = order\n end\n @test isapprox(rates[end], expected_order; rtol = 0.3)\n end\n end\n end\n\n @testset \"IMEX methods (NaiveVariant)\" begin\n dts = [2.0^(-k) for k in 4:5]\n errors = similar(dts)\n for (method, expected_order) in imex_methods_naivestorage_compatible\n for split_explicit_implicit in (false, true)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n rhs_arg! = split_explicit_implicit ? rhs_nonlinear! : rhs!\n solver = method(\n rhs_arg!,\n rhs_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver();\n isadjustable = true,\n ),\n Q;\n dt = dt,\n t0 = t0,\n split_explicit_implicit = split_explicit_implicit,\n variant = NaiveVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; rtol = 0.3)\n end\n end\n end\n\n @testset \"IMEX methods with direct solver\" begin\n dts = [2.0^(-k) for k in 4:5]\n errors = similar(dts)\n for (method, order) in imex_methods_naivestorage_compatible\n for split_explicit_implicit in (false, true)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n rhs_arg! = split_explicit_implicit ? rhs_nonlinear! : rhs!\n solver = method(\n rhs_arg!,\n rhs_linear!,\n ODETestBasicLinBE(),\n Q;\n dt = dt,\n t0 = t0,\n split_explicit_implicit = split_explicit_implicit,\n variant = NaiveVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n expected_order = order\n @test isapprox(rates[end], expected_order; rtol = 0.3)\n end\n end\n end\n\n @testset \"MRRK methods with 2 rates\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in fast_mrrk_methods\n for nsubsteps in (1, 3)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = MultirateRungeKutta(\n (\n slow_method(rhs_slow!, Q; dt = dt),\n fast_method(rhs_fast!, Q; dt = dt / nsubsteps),\n );\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n min_order = min(slow_expected_order, fast_expected_order)\n max_order = max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rates[end], min_order; atol = 0.5) ||\n isapprox(rates[end], max_order; atol = 0.5) ||\n min_order <= rates[end] <= max_order\n )\n end\n end\n end\n end\n\n @testset \"MRRK methods with IMEX\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in\n imex_methods_lowstorage_compatible\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = MultirateRungeKutta(\n (\n slow_method(rhs_slow!, Q; dt = dt),\n fast_method(\n rhs_fast!,\n rhs_fast_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver();\n isadjustable = true,\n ),\n Q;\n dt = dt,\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n ),\n ),\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n min_order = min(slow_expected_order, fast_expected_order)\n max_order = max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rates[end], min_order; atol = 0.5) ||\n isapprox(rates[end], max_order; atol = 0.5) ||\n min_order <= rates[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MRRK methods with 3 rates\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (rate3_method, rate3_order) in slow_mrrk_methods\n for (rate2_method, rate2_order) in slow_mrrk_methods\n for (rate1_method, rate1_order) in fast_mrrk_methods\n for nsubsteps in (1, 2)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = MultirateRungeKutta(\n (\n rate3_method(rhs3!, Q, dt = dt),\n rate2_method(rhs2!, Q, dt = dt / nsubsteps),\n rate1_method(\n rhs1!,\n Q;\n dt = dt / nsubsteps^2,\n ),\n );\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test 3.8 <= rates[end]\n end\n end\n end\n end\n end\n\n @testset \"MIS methods\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (method, expected_order) in mis_methods\n for fast_method in (LSRK54CarpenterKennedy,)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = method(\n rhs_slow!,\n rhs_fast!,\n fast_method,\n 4,\n Q;\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.6)\n end\n end\n end\n\n @testset \"MRI GARK methods with 2 rates\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (slow_method, expected_order) in mrigark_erk_methods\n for (fast_method, _) in fast_mrigark_methods\n for nsubsteps in (1, 3)\n for (n, dt) in enumerate(dts)\n dt /= 4 # Need a smaller dt to get convergence rate\n Q .= Qinit\n fastsolver =\n fast_method(rhs_fast!, Q; dt = dt / nsubsteps)\n solver = slow_method(\n rhs_slow!,\n fastsolver,\n Q;\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n\n errors[n] = norm(Q - Qexact)\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.5)\n end\n end\n end\n end\n\n @testset \"MRI GARK methods with 3 rates\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (rate3_method, rate3_order) in mrigark_erk_methods\n for (rate2_method, rate2_order) in mrigark_erk_methods\n for (rate1_method, _) in fast_mrigark_methods\n for nsubsteps in (1, 2)\n for (n, dt) in enumerate(dts)\n dt /= 4 # Need a smaller dt to get convergence rate\n Q .= Qinit\n solver1 =\n rate1_method(rhs1!, Q; dt = dt / nsubsteps^2)\n solver2 = rate2_method(\n rhs2!,\n solver1,\n Q;\n dt = dt / nsubsteps,\n )\n solver3 = rate3_method(\n rhs3!,\n solver2,\n Q;\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver3; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n expected_order = min(rate3_order, rate2_order)\n @test isapprox(rates[end], expected_order; atol = 0.5)\n end\n end\n end\n end\n end\n\n @testset \"MRI GARK implicit methods with 2 rates and linear solver\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (slow_method, expected_order) in mrigark_irk_methods\n for (fast_method, _) in fast_mrigark_methods\n for nsubsteps in (1, 3)\n for (n, dt) in enumerate(dts)\n dt /= 4\n Q .= Qinit\n nsteps = ceil(Int, (finaltime - t0) / dt)\n dt = (finaltime - t0) / nsteps\n fastsolver =\n fast_method(rhs_nonlinear!, Q; dt = dt / nsubsteps)\n solver = slow_method(\n rhs_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver();\n isadjustable = true,\n ),\n fastsolver,\n Q;\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n\n errors[n] = norm(Q - Qexact)\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.5)\n end\n end\n end\n end\n\n @testset \"MRI GARK implicit methods with 2 rates and custom solver\" begin\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (slow_method, expected_order) in mrigark_irk_methods\n for (fast_method, _) in fast_mrigark_methods\n for nsubsteps in (1, 3)\n for (n, dt) in enumerate(dts)\n dt /= 4\n Q .= Qinit\n nsteps = ceil(Int, (finaltime - t0) / dt)\n dt = (finaltime - t0) / nsteps\n fastsolver =\n fast_method(rhs_nonlinear!, Q; dt = dt / nsubsteps)\n solver = slow_method(\n rhs_linear!,\n ODETestBasicLinBE(),\n fastsolver,\n Q;\n dt = dt,\n t0 = t0,\n )\n solve!(Q, solver; timeend = finaltime)\n\n errors[n] = norm(Q - Qexact)\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.5)\n end\n end\n end\n end\nend\n\n@testset \"Explicit methods composition of solve!\" begin\n halftime = 1.0\n finaltime = 2.0\n dt = 0.075\n for (method, _) in explicit_methods\n Q .= Qinit\n solver1 = method(rhs!, Q; dt = dt, t0 = t0)\n solve!(Q, solver1; timeend = finaltime)\n\n Q2 = similar(Q)\n Q2 .= Qinit\n solver2 = method(rhs!, Q2; dt = dt, t0 = t0)\n solve!(Q2, solver2; timeend = halftime, adjustfinalstep = false)\n solve!(Q2, solver2; timeend = finaltime)\n\n @test Q2 == Q\n end\nend\n\n@testset \"DiffEq methods\" begin\n alg = SSPRK73()\n expected_order = 3\n dts = [2.0^(-k) for k in 3:4]\n errors = similar(dts)\n for (n, dt) in enumerate(dts)\n Q .= Qinit\n solver = DiffEqJLSolver(rhs!, alg, Q; dt = dt, t0 = t0)\n solve!(Q, solver; timeend = finaltime)\n errors[n] = norm(Q - Qexact)\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.7)\nend\n" }, { "path": "test/Numerics/ODESolvers/ode_tests_common.jl", "content": "using ClimateMachine.ODESolvers\nusing ClimateMachine.SystemSolvers\n\nconst slow_mrrk_methods =\n ((LSRK54CarpenterKennedy, 4), (LSRK144NiegemannDiehlBusch, 4))\nconst fast_mrrk_methods = (\n (LSRK54CarpenterKennedy, 4),\n (LSRK144NiegemannDiehlBusch, 4),\n (SSPRK33ShuOsher, 3),\n (SSPRK34SpiteriRuuth, 3),\n)\nconst explicit_methods = (\n (LSRK54CarpenterKennedy, 4),\n (LSRK144NiegemannDiehlBusch, 4),\n (LS3NRK44Classic, 4),\n (LS3NRK33Heuns, 3),\n (SSPRK22Heuns, 2),\n (SSPRK22Ralstons, 2),\n (SSPRK33ShuOsher, 3),\n (SSPRK34SpiteriRuuth, 3),\n (LSRKEulerMethod, 1),\n)\n\nconst imex_methods_lowstorage_compatible = (\n # Low-storage variant methods have an assumption that the\n # explicit and implicit rhs/time-scaling coefficients (B/C vectors)\n # in the Butcher tables are the same.\n (ARK1ForwardBackwardEuler, 1),\n (ARK2ImplicitExplicitMidpoint, 2),\n (ARK2GiraldoKellyConstantinescu, 2),\n (ARK437L2SA1KennedyCarpenter, 4),\n (ARK548L2SA2KennedyCarpenter, 5),\n (DBM453VoglEtAl, 3),\n)\nconst imex_methods_naivestorage_compatible = (\n imex_methods_lowstorage_compatible...,\n # Some methods can only be used with the `NaiveVariant` storage\n # scheme since, in general, ARK methods can have different time-scaling/rhs-scaling\n # coefficients (C/B vectors in the Butcher tables). For future reference,\n # any other ARK-type methods that have more general Butcher tables\n # (but with same number of stages) should be tested here:\n (Trap2LockWoodWeller, 2),\n)\n\nconst mis_methods =\n ((MIS2, 2), (MIS3C, 2), (MIS4, 3), (MIS4a, 3), (TVDMISA, 2), (TVDMISB, 2))\n\n\nconst mrigark_erk_methods = ((MRIGARKERK33aSandu, 3), (MRIGARKERK45aSandu, 4))\n\nconst mrigark_irk_methods = (\n (MRIGARKESDIRK24LSA, 2),\n (MRIGARKESDIRK23LSA, 2),\n (MRIGARKIRK21aSandu, 2),\n (MRIGARKESDIRK34aSandu, 3),\n (MRIGARKESDIRK46aSandu, 4),\n)\n\nconst fast_mrigark_methods =\n ((LSRK54CarpenterKennedy, 4), (LSRK144NiegemannDiehlBusch, 4))\n\nstruct DivideLinearSolver <: AbstractSystemSolver end\nfunction SystemSolvers.prefactorize(\n linearoperator!,\n ::DivideLinearSolver,\n args...,\n)\n linearoperator!\nend\nfunction SystemSolvers.linearsolve!(\n linearoperator!,\n preconditioner,\n ::DivideLinearSolver,\n Qtt,\n Qhat,\n args...,\n)\n @. Qhat = 1 / Qhat\n linearoperator!(Qtt, Qhat, args...)\n @. Qtt = 1 / Qtt\nend\n" }, { "path": "test/Numerics/ODESolvers/ode_tests_convergence.jl", "content": "using Test\nusing ClimateMachine\nusing StaticArrays\nusing LinearAlgebra\nusing KernelAbstractions\nusing ClimateMachine.MPIStateArrays: array_device\n\ninclude(\"ode_tests_common.jl\")\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\n\n@testset \"ODE Solvers\" begin\n @testset \"Convergence/extensive\" begin\n @testset \"1-rate ODE\" begin\n function rhs!(dQ, Q, ::Nothing, time; increment)\n if increment\n dQ .+= Q * cos(time)\n else\n dQ .= Q * cos(time)\n end\n end\n exactsolution(q0, time) = q0 * exp(sin(time))\n\n @testset \"Explicit methods\" begin\n finaltime = 20.0\n dts = [2.0^(-k) for k in 6:7]\n errors = similar(dts)\n q0 =\n ArrayType === Array ? [1.0] : range(-1.0, 1.0, length = 303)\n for (method, expected_order) in explicit_methods\n for (n, dt) in enumerate(dts)\n Q = ArrayType(q0)\n solver = method(rhs!, Q; dt = dt, t0 = 0.0)\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] =\n maximum(@. abs(Q - exactsolution(q0, finaltime)))\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(rates[end], expected_order; atol = 0.17)\n end\n end\n end\n\n @testset \"Two-rate ODE with a linear stiff part\" begin\n c = 100.0\n function rhs_full!(dQ, Q, ::Nothing, time; increment)\n if increment\n dQ .+= im * c * Q .+ exp(im * time)\n else\n dQ .= im * c * Q .+ exp(im * time)\n end\n end\n\n function rhs_nonlinear!(dQ, Q, ::Nothing, time; increment)\n if increment\n dQ .+= exp(im * time)\n else\n dQ .= exp(im * time)\n end\n end\n rhs_slow! = rhs_nonlinear!\n\n function rhs_linear!(dQ, Q, ::Nothing, time; increment)\n if increment\n dQ .+= im * c * Q\n else\n dQ .= im * c * Q\n end\n end\n rhs_fast! = rhs_linear!\n\n function exactsolution(q0, time)\n q0 * exp(im * c * time) +\n (exp(im * time) - exp(im * c * time)) / (im * (1 - c))\n end\n\n @testset \"IMEX methods (LowStorageVariant)\" begin\n finaltime = pi / 2\n dts = [2.0^(-k) for k in 13:14]\n errors = similar(dts)\n\n q0 = ArrayType <: Array ? [1.0] : range(-1.0, 1.0, length = 303)\n for (method, expected_order) in\n imex_methods_lowstorage_compatible\n for split_explicit_implicit in (false, true)\n for (n, dt) in enumerate(dts)\n Q = ArrayType{ComplexF64}(q0)\n rhs! =\n split_explicit_implicit ? rhs_nonlinear! :\n rhs_full!\n solver = method(\n rhs!,\n rhs_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver(),\n isadjustable = true,\n ),\n Q;\n dt = dt,\n t0 = 0.0,\n split_explicit_implicit = split_explicit_implicit,\n variant = LowStorageVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] = maximum(@. abs(\n Q - exactsolution(q0, finaltime),\n ))\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test errors[1] < 2.0\n @test isapprox(rates[end], expected_order; atol = 0.35)\n end\n end\n end\n\n @testset \"IMEX methods (NaiveVariant)\" begin\n finaltime = pi / 2\n dts = [2.0^(-k) for k in 13:14]\n errors = similar(dts)\n\n q0 = ArrayType <: Array ? [1.0] : range(-1.0, 1.0, length = 303)\n for (method, expected_order) in\n imex_methods_naivestorage_compatible\n for split_explicit_implicit in (false, true)\n for (n, dt) in enumerate(dts)\n Q = ArrayType{ComplexF64}(q0)\n rhs! =\n split_explicit_implicit ? rhs_nonlinear! :\n rhs_full!\n solver = method(\n rhs!,\n rhs_linear!,\n LinearBackwardEulerSolver(\n DivideLinearSolver(),\n isadjustable = true,\n ),\n Q;\n dt = dt,\n t0 = 0.0,\n split_explicit_implicit = split_explicit_implicit,\n variant = NaiveVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] = maximum(@. abs(\n Q - exactsolution(q0, finaltime),\n ))\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test errors[1] < 2.0\n @test isapprox(rates[end], expected_order; atol = 0.35)\n end\n end\n end\n\n @testset \"MRRK methods (no substeps)\" begin\n finaltime = pi / 2\n dts = [2.0^(-k) for k in 10:11]\n errors = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in fast_mrrk_methods\n q0 =\n ArrayType === Array ? [1.0] :\n range(-1.0, 1.0, length = 303)\n for (n, dt) in enumerate(dts)\n Q = ArrayType{ComplexF64}(q0)\n solver = MultirateRungeKutta(\n (\n slow_method(rhs_slow!, Q),\n fast_method(rhs_fast!, Q),\n );\n dt = dt,\n t0 = 0.0,\n )\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] = maximum(@. abs(\n Q - exactsolution(q0, finaltime),\n ))\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n min_order =\n min(slow_expected_order, fast_expected_order)\n max_order =\n max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rates[end], min_order; atol = 0.1) ||\n isapprox(rates[end], max_order; atol = 0.1) ||\n min_order <= rates[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MRRK methods (with substeps)\" begin\n finaltime = pi / 2\n dts = [2.0^(-k) for k in 14:15]\n errors = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in fast_mrrk_methods\n q0 =\n ArrayType === Array ? [1.0] :\n range(-1.0, 1.0, length = 303)\n for (n, fast_dt) in enumerate(dts)\n slow_dt = c * fast_dt\n Q = ArrayType{ComplexF64}(q0)\n solver = MultirateRungeKutta((\n slow_method(rhs_slow!, Q; dt = slow_dt),\n fast_method(rhs_fast!, Q; dt = fast_dt),\n ))\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] = maximum(@. abs(\n Q - exactsolution(q0, finaltime),\n ))\n end\n\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n min_order =\n min(slow_expected_order, fast_expected_order)\n max_order =\n max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rates[end], min_order; atol = 0.1) ||\n isapprox(rates[end], max_order; atol = 0.1) ||\n min_order <= rates[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MIS methods (with substeps)\" begin\n finaltime = pi / 2\n dts = [2.0^(-k) for k in 8:9]\n errors = similar(dts)\n for (mis_method, mis_expected_order) in mis_methods\n for fast_method in (LSRK54CarpenterKennedy,)\n q0 =\n ArrayType === Array ? [1.0] :\n range(-1.0, 1.0, length = 303)\n for (n, dt) in enumerate(dts)\n Q = ArrayType{ComplexF64}(q0)\n solver = mis_method(\n rhs_slow!,\n rhs_fast!,\n fast_method,\n 4,\n Q;\n dt = dt,\n t0 = 0.0,\n )\n solve!(Q, solver; timeend = finaltime)\n Q = Array(Q)\n errors[n] = maximum(@. abs(\n Q - exactsolution(q0, finaltime),\n ))\n end\n rates = log2.(errors[1:(end - 1)] ./ errors[2:end])\n @test isapprox(\n rates[end],\n mis_expected_order;\n atol = 0.1,\n )\n end\n end\n end\n end\n\n #=\n Test problem (4.2) from [Roberts2018](@cite)\n\n Note: The actual rates are all over the place with this test and passing largely\n depends on final dt size\n =#\n @testset \"2-rate ODE from Roberts2018\" begin\n ω = 100\n λf = -10\n λs = -1\n ξ = 1 // 10\n α = 1\n ηfs = ((1 - ξ) / α) * (λf - λs)\n ηsf = -ξ * α * (λf - λs)\n Ω = @SMatrix [\n λf ηfs\n ηsf λs\n ]\n\n function rhs_fast!(dQ, Q, param, t; increment)\n @inbounds begin\n increment || (dQ .= 0)\n yf = Q[1]\n ys = Q[2]\n gf = (-3 + yf^2 - cos(ω * t)) / 2yf\n gs = (-2 + ys^2 - cos(t)) / 2ys\n dQ[1] += Ω[1, 1] * gf + Ω[1, 2] * gs - ω * sin(ω * t) / 2yf\n end\n end\n\n function rhs_slow!(dQ, Q, param, t; increment)\n @inbounds begin\n increment || (dQ .= 0)\n yf = Q[1]\n ys = Q[2]\n gf = (-3 + yf^2 - cos(ω * t)) / 2yf\n gs = (-2 + ys^2 - cos(t)) / 2ys\n dQ[2] += Ω[2, 1] * gf + Ω[2, 2] * gs - sin(t) / 2ys\n end\n end\n\n exactsolution(t) = [sqrt(3 + cos(ω * t)); sqrt(2 + cos(t))]\n\n @testset \"MRRK methods (no substeps)\" begin\n finaltime = 5π / 2\n dts = [2.0^(-k) for k in 7:8]\n error = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in fast_mrrk_methods\n for (n, dt) in enumerate(dts)\n Q = exactsolution(0)\n solver = MultirateRungeKutta(\n (\n slow_method(rhs_slow!, Q),\n fast_method(rhs_fast!, Q),\n );\n dt = dt,\n t0 = 0.0,\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n min_order =\n min(slow_expected_order, fast_expected_order)\n max_order =\n max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rate[end], min_order; atol = 0.3) ||\n isapprox(rate[end], max_order; atol = 0.3) ||\n min_order <= rate[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MRRK methods (with substeps)\" begin\n finaltime = 5π / 2\n dts = [2.0^(-k) for k in 8:9]\n error = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in fast_mrrk_methods\n for (n, fast_dt) in enumerate(dts)\n Q = exactsolution(0)\n slow_dt = ω * fast_dt\n solver = MultirateRungeKutta((\n slow_method(rhs_slow!, Q; dt = slow_dt),\n fast_method(rhs_fast!, Q; dt = fast_dt),\n ))\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n min_order =\n min(slow_expected_order, fast_expected_order)\n max_order =\n max(slow_expected_order, fast_expected_order)\n @test (\n isapprox(rate[end], min_order; atol = 0.3) ||\n isapprox(rate[end], max_order; atol = 0.3) ||\n min_order <= rate[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MRRK methods with IMEX fast solver\" begin\n function rhs_zero!(dQ, Q, param, t; increment)\n if !increment\n dQ .= 0\n end\n end\n\n finaltime = 5π / 2\n dts = [2.0^(-k) for k in 8:9]\n error = similar(dts)\n for (slow_method, slow_expected_order) in slow_mrrk_methods\n for (fast_method, fast_expected_order) in\n imex_methods_lowstorage_compatible\n for (n, fast_dt) in enumerate(dts)\n Q = exactsolution(0)\n slow_dt = ω * fast_dt\n solver = MultirateRungeKutta((\n slow_method(rhs_slow!, Q; dt = slow_dt),\n fast_method(\n rhs_fast!,\n rhs_zero!,\n LinearBackwardEulerSolver(\n DivideLinearSolver(),\n isadjustable = true,\n ),\n Q;\n dt = fast_dt,\n split_explicit_implicit = false,\n variant = LowStorageVariant(),\n ),\n ))\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n min_order =\n min(slow_expected_order, fast_expected_order)\n max_order =\n max(slow_expected_order, fast_expected_order)\n atol =\n fast_method == ARK2GiraldoKellyConstantinescu ?\n 0.5 : 0.37\n @test (\n isapprox(rate[end], min_order; atol = atol) ||\n isapprox(rate[end], max_order; atol = atol) ||\n min_order <= rate[end] <= max_order\n )\n end\n end\n end\n\n @testset \"MIS methods (with substeps)\" begin\n finaltime = 5π / 2\n dts = [2.0^(-k) for k in 9:10]\n error = similar(dts)\n for (mis_method, mis_expected_order) in mis_methods\n for fast_method in (LSRK54CarpenterKennedy,)\n for (n, dt) in enumerate(dts)\n Q = exactsolution(0)\n solver = mis_method(\n rhs_slow!,\n rhs_fast!,\n fast_method,\n 4,\n Q;\n dt = dt,\n t0 = 0.0,\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n @test isapprox(\n rate[end],\n mis_expected_order;\n atol = 0.1,\n )\n end\n end\n end\n end\n\n # Simple 3-rate problem based on test of [Roberts2018](@cite)\n #\n # NOTE: Since we have no theory to say this ODE solver is accurate, the rates\n # suggest that things are really only 2nd order.\n #\n # TODO: This is not great, but no theory to say we should be accurate!\n @testset \"3-rate ODE\" begin\n ω1, ω2, ω3 = 10000, 100, 1\n λ1, λ2, λ3 = -100, -10, -1\n β1, β2, β3 = 2, 3, 4\n\n ξ12 = λ2 / λ1\n ξ13 = λ3 / λ1\n ξ23 = λ3 / λ2\n\n α12, α13, α23 = 1, 1, 1\n\n η12 = ((1 - ξ12) / α12) * (λ1 - λ2)\n η13 = ((1 - ξ13) / α13) * (λ1 - λ3)\n η23 = ((1 - ξ23) / α23) * (λ2 - λ3)\n\n η21 = ξ12 * α12 * (λ2 - λ1)\n η31 = ξ13 * α13 * (λ3 - λ1)\n η32 = ξ23 * α23 * (λ3 - λ2)\n\n Ω = @SMatrix [\n λ1 η12 η13\n η21 λ2 η23\n η31 η32 λ3\n ]\n\n function rhs1!(dQ, Q, param, t; increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[1] += Ω[1, :]' * g - ω1 * sin(ω1 * t) / 2y1\n end\n end\n function rhs2!(dQ, Q, param, t; increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[2] += Ω[2, :]' * g - ω2 * sin(ω2 * t) / 2y2\n end\n end\n function rhs3!(dQ, Q, param, t; increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[3] += Ω[3, :]' * g - ω3 * sin(ω3 * t) / 2y3\n end\n end\n function rhs12!(dQ, Q, param, t; increment)\n rhs1!(dQ, Q, param, t; increment = increment)\n rhs2!(dQ, Q, param, t; increment = true)\n end\n\n exactsolution(t) = [\n sqrt(β1 + cos(ω1 * t)),\n sqrt(β2 + cos(ω2 * t)),\n sqrt(β3 + cos(ω3 * t)),\n ]\n\n @testset \"MRRK methods (no substeps)\" begin\n finaltime = π / 2\n dts = [2.0^(-k) for k in 9:10]\n error = similar(dts)\n for (rate3_method, rate3_order) in slow_mrrk_methods\n for (rate2_method, rate2_order) in slow_mrrk_methods\n for (rate1_method, rate1_order) in fast_mrrk_methods\n for (n, dt) in enumerate(dts)\n Q = exactsolution(0)\n solver = MultirateRungeKutta(\n (\n rate3_method(rhs3!, Q),\n rate2_method(rhs2!, Q),\n rate1_method(rhs1!, Q),\n );\n dt = dt,\n t0 = 0.0,\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n min_order =\n min(rate3_order, rate2_order, rate1_order)\n max_order =\n max(rate3_order, rate2_order, rate1_order)\n @test 2 <= rate[end]\n end\n end\n end\n end\n\n @testset \"MRRK methods (with substeps)\" begin\n finaltime = π / 2\n dts = [2.0^(-k) for k in 16:17]\n error = similar(dts)\n for (rate3_method, rate3_order) in slow_mrrk_methods\n for (rate2_method, rate2_order) in slow_mrrk_methods\n for (rate1_method, rate1_order) in fast_mrrk_methods\n for (n, dt1) in enumerate(dts)\n Q = exactsolution(0)\n dt2 = (ω1 / ω2) * dt1\n dt3 = (ω2 / ω3) * dt2\n solver = MultirateRungeKutta((\n rate3_method(rhs3!, Q; dt = dt3),\n rate2_method(rhs2!, Q; dt = dt2),\n rate1_method(rhs1!, Q; dt = dt1),\n ))\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n min_order =\n min(rate3_order, rate2_order, rate1_order)\n max_order =\n max(rate3_order, rate2_order, rate1_order)\n\n @test 2 <= rate[end]\n end\n end\n end\n end\n end\n\n #=\n Test problem (8.2) from [Sandu2019](@cite) for MRI-GARK Schemes\n =#\n @testset \"2-rate problem\" begin\n ω = 20\n λf = -10\n λs = -1\n ξ = 1 // 10\n α = 1\n ηfs = ((1 - ξ) / α) * (λf - λs)\n ηsf = -ξ * α * (λf - λs)\n Ω = @SMatrix [\n λf ηfs\n ηsf λs\n ]\n\n function rhs_fast!(dQ, Q, param, t; increment)\n @kernel function knl!(dQ, Q, t, increment)\n @inbounds begin\n increment || (dQ .= 0)\n yf = Q[1]\n ys = Q[2]\n gf = (-3 + yf^2 - cos(ω * t)) / 2yf\n gs = (-2 + ys^2 - cos(t)) / 2ys\n dQ[1] +=\n Ω[1, 1] * gf + Ω[1, 2] * gs - ω * sin(ω * t) / 2yf\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n dQ,\n Q,\n t,\n increment;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n\n function rhs_slow!(dQ, Q, param, t; increment)\n @kernel function knl!(dQ, Q, t, increment)\n @inbounds begin\n increment || (dQ .= 0)\n yf = Q[1]\n ys = Q[2]\n gf = (-3 + yf^2 - cos(ω * t)) / 2yf\n gs = (-2 + ys^2 - cos(t)) / 2ys\n dQ[2] += Ω[2, 1] * gf + Ω[2, 2] * gs - sin(t) / 2ys\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n dQ,\n Q,\n t,\n increment;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n struct ODETestConvNonLinBE <: AbstractBackwardEulerSolver end\n ODESolvers.Δt_is_adjustable(::ODETestConvNonLinBE) = true\n function (::ODETestConvNonLinBE)(Q, Qhat, α, p, t)\n @kernel function knl!(Q, Qhat, α, p, t)\n @inbounds begin\n # Slow RHS has zero tendency of yf so just copy Qhat\n Q[1] = yf = Qhat[1]\n\n gf = (-3 + yf^2 - cos(ω * t)) / 2yf\n\n # solves: Q = Qhat[2] + α * rhs_slow(Q, t)\n # (simplifies to a quadratic equation)\n a = 2 - α * Ω[2, 2]\n b = -2 * (Qhat[2] + α * Ω[2, 1] * gf)\n c = α * (Ω[2, 2] * (2 + cos(t)) + sin(t))\n Q[2] = (-b + sqrt(b^2 - 4 * a * c)) / (2a)\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n Q,\n Qhat,\n α,\n p,\n t;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n exactsolution(t) =\n ArrayType([sqrt(3 + cos(ω * t)); sqrt(2 + cos(t))])\n\n finaltime = 1\n dts = [2.0^(-k) for k in 7:9]\n error = similar(dts)\n @testset \"Explicit\" begin\n for (mri_method, mri_expected_order) in mrigark_erk_methods\n for (fast_method, fast_expected_order) in\n fast_mrigark_methods\n for (n, slow_dt) in enumerate(dts)\n Q = exactsolution(0)\n fast_dt = slow_dt / ω\n fastsolver = fast_method(rhs_fast!, Q; dt = fast_dt)\n solver = mri_method(\n rhs_slow!,\n fastsolver,\n Q,\n dt = slow_dt,\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n order = mri_expected_order\n @test isapprox(\n rate[end],\n mri_expected_order;\n atol = 0.3,\n )\n end\n end\n end\n\n @testset \"Implicit\" begin\n for (mri_method, mri_expected_order) in mrigark_irk_methods\n for (fast_method, fast_expected_order) in\n fast_mrigark_methods\n for (n, slow_dt) in enumerate(dts)\n Q = exactsolution(0)\n fast_dt = slow_dt / ω\n fastsolver = fast_method(rhs_fast!, Q; dt = fast_dt)\n solver = mri_method(\n rhs_slow!,\n ODETestConvNonLinBE(),\n fastsolver,\n Q,\n dt = slow_dt,\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n order = mri_expected_order\n @test isapprox(\n rate[end],\n mri_expected_order;\n atol = 0.3,\n )\n end\n end\n end\n\n @testset \"IMEX\" begin\n for (method, expected_order) in\n imex_methods_naivestorage_compatible\n for (n, dt) in enumerate(dts)\n Q = exactsolution(0)\n solver = method(\n rhs_fast!,\n rhs_slow!,\n ODETestConvNonLinBE(),\n Q;\n dt = dt,\n t0 = 0.0,\n split_explicit_implicit = true,\n variant = NaiveVariant(),\n )\n solve!(Q, solver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n order = expected_order\n @test isapprox(rate[end], expected_order; atol = 0.3)\n end\n end\n end\n\n # 3-rate modification of the above test problem\n @testset \"3-rate problem\" begin\n ω1, ω2, ω3 = 20, 5, 1\n λ1, λ2, λ3 = -20, -5, -1\n β1, β2, β3 = 2, 2, 2\n\n ξ12 = λ2 / (λ1 + λ2)\n ξ13 = λ3 / (λ1 + λ3)\n ξ23 = λ3 / (λ2 + λ3)\n\n α12, α13, α23 = 1, 1, 1\n\n η12 = ((1 - ξ12) / α12) * (λ1 - λ2)\n η13 = 0#((1 - ξ13) / α13) * (λ1 - λ3)\n η23 = ((1 - ξ23) / α23) * (λ2 - λ3)\n\n η21 = ξ12 * α12 * (λ2 - λ1)\n η31 = 0#ξ13 * α13 * (λ3 - λ1)\n η32 = ξ23 * α23 * (λ3 - λ2)\n\n Ω = @SMatrix [\n λ1 η12 η13\n η21 λ2 η23\n η31 η32 λ3\n ]\n\n function rhs1!(dQ, Q, param, t; increment)\n @kernel function knl!(dQ, Q, t, increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[1] += Ω[1, :]' * g - ω1 * sin(ω1 * t) / 2y1\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n dQ,\n Q,\n t,\n increment;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n function rhs2!(dQ, Q, param, t; increment)\n @kernel function knl!(dQ, Q, t, increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[2] += Ω[2, :]' * g - ω2 * sin(ω2 * t) / 2y2\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n dQ,\n Q,\n t,\n increment;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n function rhs3!(dQ, Q, param, t; increment)\n @kernel function knl!(dQ, Q, t, increment)\n @inbounds begin\n increment || (dQ .= 0)\n y1, y2, y3 = Q[1], Q[2], Q[3]\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n (-β3 + y3^2 - cos(ω3 * t)) / 2y3,\n ]\n dQ[3] += Ω[3, :]' * g - ω3 * sin(ω3 * t) / 2y3\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n dQ,\n Q,\n t,\n increment;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n struct ODETestConvNonLinBE3Rate <: AbstractBackwardEulerSolver end\n ODESolvers.Δt_is_adjustable(::ODETestConvNonLinBE3Rate) = true\n function (::ODETestConvNonLinBE3Rate)(Q, Qhat, α, p, t)\n @kernel function knl!(Q, Qhat, α, p, t)\n @inbounds begin\n # Slower RHS has zero tendency of yf so just copy Qhat\n Q[1] = y1 = Qhat[1]\n Q[2] = y2 = Qhat[2]\n\n g = @SVector [\n (-β1 + y1^2 - cos(ω1 * t)) / 2y1,\n (-β2 + y2^2 - cos(ω2 * t)) / 2y2,\n ]\n\n # solves: Q = Qhat + α * rhs_slow(Q, t)\n # (simplifies to a quadratic equation)\n a = 2 - α * Ω[3, 3]\n b =\n -2 *\n (Qhat[3] + α * Ω[3, 1] * g[1] + α * Ω[3, 2] * g[2])\n c = α * (Ω[3, 3] * (β3 + cos(t)) + sin(t))\n Q[3] = (-b + sqrt(b^2 - 4 * a * c)) / (2a)\n end\n end\n event = Event(array_device(Q))\n event = knl!(array_device(Q), 1)(\n Q,\n Qhat,\n α,\n p,\n t;\n ndrange = 1,\n dependencies = (event,),\n )\n wait(array_device(Q), event)\n end\n\n exactsolution(t) = ArrayType([\n sqrt(β1 + cos(ω1 * t)),\n sqrt(β2 + cos(ω2 * t)),\n sqrt(β3 + cos(ω3 * t)),\n ])\n\n finaltime = 1\n dts = [2.0^(-k) for k in 6:7]\n error = similar(dts)\n @testset \"Explicit\" begin\n for (slow_method, slow_order) in mrigark_erk_methods\n for (mid_method, mid_order) in mrigark_erk_methods\n for (fast_method, fast_order) in fast_mrigark_methods\n for (n, slow_dt) in enumerate(dts)\n Q = exactsolution(0)\n mid_dt = slow_dt / ω2\n fast_dt = slow_dt / ω1\n fastsolver = fast_method(rhs1!, Q; dt = fast_dt)\n midsolver = mid_method(\n rhs2!,\n fastsolver,\n Q,\n dt = mid_dt,\n )\n slowsolver = slow_method(\n rhs3!,\n midsolver,\n Q,\n dt = slow_dt,\n )\n solve!(Q, slowsolver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n expected_order =\n min(slow_order, mid_order, fast_order)\n @test isapprox(\n rate[end],\n expected_order;\n atol = 0.3,\n )\n end\n end\n end\n end\n @testset \"Implicit-Explicit\" begin\n for (slow_method, slow_order) in mrigark_irk_methods\n for (mid_method, mid_order) in mrigark_erk_methods\n for (fast_method, fast_order) in fast_mrigark_methods\n for (n, slow_dt) in enumerate(dts)\n Q = exactsolution(0)\n mid_dt = slow_dt / ω2\n fast_dt = slow_dt / ω1\n fastsolver = fast_method(rhs1!, Q; dt = fast_dt)\n midsolver = mid_method(\n rhs2!,\n fastsolver,\n Q,\n dt = mid_dt,\n )\n slowsolver = slow_method(\n rhs3!,\n ODETestConvNonLinBE3Rate(),\n midsolver,\n Q,\n dt = slow_dt,\n )\n solve!(Q, slowsolver; timeend = finaltime)\n error[n] = norm(Q - exactsolution(finaltime))\n end\n\n rate = log2.(error[1:(end - 1)] ./ error[2:end])\n expected_order =\n min(slow_order, mid_order, fast_order)\n @test isapprox(\n rate[end],\n expected_order;\n atol = 0.4,\n )\n end\n end\n end\n end\n end\n end\nend\n" }, { "path": "test/Numerics/ODESolvers/runtests.jl", "content": "using Test, MPI\ninclude(joinpath(\"..\", \"..\", \"testhelpers.jl\"))\n\n@testset \"ODE Solvers\" begin\n runmpi(joinpath(@__DIR__, \"callbacks.jl\"))\nend\n" }, { "path": "test/Numerics/SystemSolvers/bandedsystem.jl", "content": "using ClimateMachine\nusing MPI\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing Logging\nusing LinearAlgebra\nusing Random\nusing StaticArrays\nusing ClimateMachine.BalanceLaws\nimport ClimateMachine.BalanceLaws: vars_state, number_states\nusing ClimateMachine.DGMethods: DGModel, DGFVModel, init_ode_state, create_state\nusing ClimateMachine.SystemSolvers: banded_matrix, banded_matrix_vector_product!\nusing ClimateMachine.DGMethods.FVReconstructions: FVConstant, FVLinear\nusing ClimateMachine.DGMethods.NumericalFluxes:\n RusanovNumericalFlux,\n CentralNumericalFluxSecondOrder,\n CentralNumericalFluxGradient\nusing ClimateMachine.MPIStateArrays: MPIStateArray, euclidean_distance\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.SystemSolvers: band_lu!, band_forward!, band_back!\n\nusing Test\n\ninclude(\"../DGMethods/advection_diffusion/advection_diffusion_model.jl\")\n\nstruct Pseudo1D{n, α, β, μ, δ} <: AdvectionDiffusionProblem end\n\nfunction init_velocity_diffusion!(\n ::Pseudo1D{n, α, β},\n aux::Vars,\n geom::LocalGeometry,\n) where {n, α, β}\n # Direction of flow is n with magnitude α\n aux.advection.u = α * n\n\n # diffusion of strength β in the n direction\n aux.diffusion.D = β * n * n'\nend\n\nstruct BigAdvectionDiffusion <: BalanceLaw end\nfunction vars_state(::BigAdvectionDiffusion, ::Prognostic, FT)\n @vars begin\n ρ::FT\n X::SVector{3, FT}\n end\nend\n\nfunction initial_condition!(\n ::Pseudo1D{n, α, β, μ, δ},\n state,\n aux,\n localgeo,\n t,\n) where {n, α, β, μ, δ}\n ξn = dot(n, localgeo.coord)\n # ξT = SVector(x) - ξn * n\n state.ρ = exp(-(ξn - μ - α * t)^2 / (4 * β * (δ + t))) / sqrt(1 + t / δ)\nend\n\nlet\n # boiler plate MPI stuff\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n Random.seed!(777 + MPI.Comm_rank(mpicomm))\n\n # Mesh generation parameters\n Neh = 10\n Nev = 4\n\n @testset \"$(@__FILE__) DGModel matrix\" begin\n for FT in (Float64, Float32)\n for dim in (2, 3)\n connectivity = dim == 3 ? :full : :face\n for single_column in (false, true)\n # Setup the topology\n if dim == 2\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n elseif dim == 3\n brickrange = (\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(0); length = Neh + 1, stop = 1),\n range(FT(1); length = Nev + 1, stop = 2),\n )\n end\n topl = StackedBrickTopology(\n mpicomm,\n brickrange,\n connectivity = connectivity,\n )\n\n # Warp mesh\n function warpfun(ξ1, ξ2, ξ3)\n # single column currently requires no geometry warping\n\n # Even if the warping is in only the horizontal, the way we\n # compute metrics causes problems for the single column approach\n # (possibly need to not use curl-invariant computation)\n if !single_column\n ξ1 = ξ1 + sin(2 * FT(π) * ξ1 * ξ2) / 10\n ξ2 = ξ2 + sin(2 * FT(π) * ξ1) / 5\n if dim == 3\n ξ3 = ξ3 + sin(8 * FT(π) * ξ1 * ξ2) / 10\n end\n end\n (ξ1, ξ2, ξ3)\n end\n\n d = dim == 2 ? FT[1, 10, 0] : FT[1, 1, 10]\n n = SVector{3, FT}(d ./ norm(d))\n\n α = FT(1)\n β = FT(1 // 100)\n μ = FT(-1 // 2)\n δ = FT(1 // 10)\n bcs = (HomogeneousBC{0}(),)\n model =\n AdvectionDiffusion{dim}(Pseudo1D{n, α, β, μ, δ}(), bcs)\n\n for (N, fvmethod) in (\n ((4, 4), nothing),\n ((4, 0), FVConstant()),\n ((4, 0), FVLinear()),\n )\n\n # create the actual grid\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n meshwarp = warpfun,\n )\n\n\n # the nonlinear model is needed so we can grab the state_auxiliary below\n dg =\n (fvmethod === nothing) ?\n DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n ) :\n DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n vdg =\n (fvmethod === nothing) ?\n DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n ) :\n DGFVModel(\n model,\n grid,\n fvmethod,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n direction = VerticalDirection(),\n state_auxiliary = dg.state_auxiliary,\n )\n\n A_banded = banded_matrix(\n vdg,\n MPIStateArray(dg),\n MPIStateArray(dg);\n single_column = single_column,\n )\n\n Q = init_ode_state(dg, FT(0))\n dQ1 = MPIStateArray(dg)\n dQ2 = MPIStateArray(dg)\n\n vdg(dQ1, Q, nothing, 0; increment = false)\n Q.data .= dQ1.data\n\n vdg(dQ1, Q, nothing, 0; increment = false)\n banded_matrix_vector_product!(A_banded, dQ2, Q)\n @test all(isapprox.(\n Array(dQ1.realdata),\n Array(dQ2.realdata),\n atol = 100 * eps(FT),\n ))\n\n big_Q = create_state(\n BigAdvectionDiffusion(),\n grid,\n Prognostic(),\n )\n big_dQ = create_state(\n BigAdvectionDiffusion(),\n grid,\n Prognostic(),\n )\n\n big_Q .= NaN\n big_dQ .= NaN\n\n big_Q.data[:, 1:size(Q, 2), :] .= Q.data\n\n vdg(big_dQ, big_Q, nothing, 0; increment = false)\n\n @test all(isapprox.(\n Array(big_dQ.realdata[:, 1:size(Q, 2), :]),\n Array(dQ1.realdata),\n atol = 100 * eps(FT),\n ))\n\n @test all(big_dQ[:, (size(Q, 2) + 1):end, :] .== 0)\n\n big_dQ.data[:, (size(Q, 2) + 1):end, :] .= -7\n\n vdg(big_dQ, big_Q, nothing, 0; increment = true)\n\n @test all(big_dQ[:, (size(Q, 2) + 1):end, :] .== -7)\n\n @test all(isapprox.(\n Array(big_dQ.realdata[:, 1:size(Q, 2), :]),\n Array(dQ1.realdata),\n atol = 100 * eps(FT),\n ))\n\n big_A = banded_matrix(\n vdg,\n similar(big_Q),\n similar(big_dQ);\n single_column = single_column,\n )\n @test all(isapprox.(\n Array(big_A.data),\n Array(A_banded.data),\n atol = 100 * eps(FT),\n ))\n\n α = FT(1 // 10)\n function op!(LQ, Q)\n vdg(LQ, Q, nothing, 0; increment = false)\n @. LQ = Q + α * LQ\n end\n\n A_banded = banded_matrix(\n op!,\n vdg,\n MPIStateArray(dg),\n MPIStateArray(dg);\n single_column = single_column,\n )\n\n Q = init_ode_state(dg, FT(0))\n dQ1 = MPIStateArray(vdg)\n dQ2 = MPIStateArray(vdg)\n\n op!(dQ1, Q)\n Q.data .= dQ1.data\n\n op!(dQ1, Q)\n banded_matrix_vector_product!(A_banded, dQ2, Q)\n @test all(isapprox.(\n Array(dQ1.realdata),\n Array(dQ2.realdata),\n atol = 100 * eps(FT),\n ))\n\n big_Q .= NaN\n big_Q.data[:, 1:size(Q, 2), :] .= Q.data\n\n big_dQ .= NaN\n op!(big_dQ, big_Q)\n big_dQ.data[:, (size(Q, 2) + 1):end, :] .= -7\n\n big_A = banded_matrix(\n op!,\n vdg,\n similar(big_Q),\n similar(big_dQ);\n single_column = single_column,\n )\n\n @test all(isapprox.(\n Array(big_A.data),\n Array(A_banded.data),\n atol = 100 * eps(FT),\n ))\n\n band_lu!(big_A)\n band_forward!(big_dQ, big_A)\n band_back!(big_dQ, big_A)\n\n @test all(isapprox.(\n Array(big_dQ.realdata[:, 1:size(Q, 2), :]),\n Array(Q.realdata),\n atol = 100 * eps(FT),\n ))\n @test all(big_dQ[:, (size(Q, 2) + 1):end, :] .== -7)\n end\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/SystemSolvers/bgmres.jl", "content": "using MPI\nusing Test\nusing LinearAlgebra\nusing Random\nusing StaticArrays\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.MPIStateArrays\nusing CUDA\nusing Random\nusing KernelAbstractions\n\nimport ClimateMachine.MPIStateArrays: array_device\n\nClimateMachine.init(; fix_rng_seed = true)\n\n@kernel function multiply_by_A!(x, A, y, n1, n2)\n I = @index(Global)\n for i in 1:n1\n tmp = zero(eltype(x))\n for j in 1:n2\n tmp += A[i, j, I] * y[j, I]\n end\n x[i, I] = tmp\n end\nend\n\nlet\n if CUDA.has_cuda_gpu()\n Arrays = [Array, CuArray]\n else\n Arrays = [Array]\n end\n\n for ArrayType in Arrays\n for T in [Float32, Float64]\n ϵ = eps(T)\n\n @testset \"($ArrayType, $T) Basic Test\" begin\n Random.seed!(42)\n\n # Test 1: Basic Functionality\n n = 100 # size of local (batch) matrix\n ni = 100 # batch size\n\n b = ArrayType(randn(n, ni)) # rhs\n x = ArrayType(randn(n, ni)) # initial guess\n x_ref = similar(x)\n\n A = ArrayType(randn((n, n, ni)) ./ sqrt(n))\n for i in 1:n\n A[i, i, :] .+= 10i\n end\n\n ss = size(b)[1]\n bgmres = BatchedGeneralizedMinimalResidual(\n b,\n n,\n ni,\n M = ss,\n atol = ϵ,\n rtol = ϵ,\n )\n\n # Define the linear operator\n function closure_linear_operator_multi!(A, n1, n2, n3)\n function linear_operator!(x, y)\n device = array_device(x)\n if isa(device, CPU)\n groupsize = Threads.nthreads()\n else # isa(device, CUDADevice)\n groupsize = 256\n end\n event = Event(device)\n event = multiply_by_A!(device, groupsize)(\n x,\n A,\n y,\n n1,\n n2,\n ndrange = n3,\n dependencies = (event,),\n )\n wait(device, event)\n nothing\n end\n end\n linear_operator! = closure_linear_operator_multi!(A, size(A)...)\n\n # Now solve\n linearsolve!(\n linear_operator!,\n nothing,\n bgmres,\n x,\n b;\n max_iters = Inf,\n )\n\n # reference solution\n for i in 1:ni\n x_ref[:, i] = A[:, :, i] \\ b[:, i]\n end\n\n @test norm(x - x_ref) < 1000ϵ\n end\n\n ###\n # Test 2: MPI State Array\n ###\n @testset \"($ArrayType, $T) MPIStateArray Test\" begin\n Random.seed!(43)\n\n n1 = 8\n n2 = 3\n n3 = 10\n mpicomm = MPI.COMM_WORLD\n mpi_b = MPIStateArray{T}(mpicomm, ArrayType, n1, n2, n3)\n mpi_x = MPIStateArray{T}(mpicomm, ArrayType, n1, n2, n3)\n mpi_A = ArrayType(randn(n1 * n2, n1 * n2, n3))\n\n mpi_b.data[:] .= ArrayType(randn(n1 * n2 * n3))\n mpi_x.data[:] .= ArrayType(randn(n1 * n2 * n3))\n\n bgmres = BatchedGeneralizedMinimalResidual(\n mpi_b,\n n1 * n2,\n n3,\n M = n1 * n2,\n atol = ϵ,\n rtol = ϵ,\n )\n\n # Now define the linear operator\n function closure_linear_operator_mpi!(A, n1, n2, n3)\n function linear_operator!(x, y)\n alias_x = reshape(x.data, (n1, n3))\n alias_y = reshape(y.data, (n1, n3))\n device = array_device(x)\n if isa(device, CPU)\n groupsize = Threads.nthreads()\n else # isa(device, CUDADevice)\n groupsize = 256\n end\n event = Event(device)\n event = multiply_by_A!(device, groupsize)(\n alias_x,\n A,\n alias_y,\n n1,\n n2,\n ndrange = n3,\n dependencies = (event,),\n )\n wait(device, event)\n nothing\n end\n end\n linear_operator! =\n closure_linear_operator_mpi!(mpi_A, size(mpi_A)...)\n\n # Now solve\n linearsolve!(\n linear_operator!,\n nothing,\n bgmres,\n mpi_x,\n mpi_b;\n max_iters = Inf,\n )\n\n # check all solutions\n norms = -zeros(n3)\n for cidx in 1:n3\n sol =\n Array(mpi_A[:, :, cidx]) \\\n Array(mpi_b.data[:, :, cidx])[:]\n norms[cidx] = norm(sol - Array(mpi_x.data[:, :, cidx])[:])\n end\n @test maximum(norms) < 8000ϵ\n end\n\n ###\n # Test 3: Columnwise test\n ###\n @testset \"(Array, $T) Columnwise Test\" begin\n\n Random.seed!(2424)\n function closure_linear_operator_columwise!(A, tup)\n function linear_operator!(y, x)\n alias_x = reshape(x, tup)\n alias_y = reshape(y, tup)\n for i6 in 1:tup[6]\n for i4 in 1:tup[4]\n for i2 in 1:tup[2]\n for i1 in 1:tup[1]\n tmp = alias_x[i1, i2, :, i4, :, i6][:]\n tmp2 = A[i1, i2, i4, i6] * tmp\n alias_y[i1, i2, :, i4, :, i6] .=\n reshape(tmp2, (tup[3], tup[5]))\n end\n end\n end\n end\n end\n end\n\n tup = (3, 4, 7, 6, 5, 2)\n B = [\n randn(tup[3] * tup[5], tup[3] * tup[5])\n for\n i1 in 1:tup[1],\n i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n columnwise_A = [\n B[i1, i2, i4, i6] + 10I\n for\n i1 in 1:tup[1],\n i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n # taking the inverse of A isn't great, but it is convenient\n columnwise_inv_A = [\n inv(columnwise_A[i1, i2, i4, i6])\n for\n i1 in 1:tup[1],\n i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n columnwise_linear_operator! =\n closure_linear_operator_columwise!(columnwise_A, tup)\n columnwise_inverse_linear_operator! =\n closure_linear_operator_columwise!(columnwise_inv_A, tup)\n\n mpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])\n b = randn(mpi_tup)\n x = copy(b)\n columnwise_inverse_linear_operator!(x, b)\n x +=\n randn((tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])) *\n 0.1\n\n reshape_tuple_f = tup\n permute_tuple = (5, 3, 1, 4, 2, 6)\n\n bgmres = BatchedGeneralizedMinimalResidual(\n b,\n tup[3] * tup[5],\n tup[1] * tup[2] * tup[4] * tup[6];\n M = tup[3] * tup[5],\n forward_reshape = tup,\n forward_permute = permute_tuple,\n atol = 10ϵ,\n rtol = 10ϵ,\n )\n\n x_exact = copy(x)\n linearsolve!(\n columnwise_linear_operator!,\n nothing,\n bgmres,\n x,\n b,\n max_iters = tup[3] * tup[5],\n )\n\n columnwise_inverse_linear_operator!(x_exact, b)\n @test norm(x - x_exact) / norm(x_exact) < 1000ϵ\n columnwise_linear_operator!(x_exact, x)\n @test norm(x_exact - b) / norm(b) < 1000ϵ\n end\n end\n end\nend\n" }, { "path": "test/Numerics/SystemSolvers/cg.jl", "content": "using CUDA\nusing MPI\nusing Test\nusing LinearAlgebra\nusing Random\nusing StaticArrays\nusing KernelAbstractions: CPU, CUDADevice\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.MPIStateArrays\nusing Random\n\nRandom.seed!(1235)\n\nlet\n ClimateMachine.init()\n mpicomm = MPI.COMM_WORLD\n ArrayType = ClimateMachine.array_type()\n n = 100\n T = Float64\n A = rand(n, n)\n scale = 2.0\n ϵ = 0.1\n # Matrix 1\n A = A' * A .* ϵ + scale * I\n\n err_thresh = sqrt(eps(T))\n\n # Matrix 2\n # A = Diagonal(collect(1:n) * 1.0)\n positive_definite = minimum(eigvals(A)) > eps(1.0)\n @test positive_definite\n\n b = ones(n) * 1.0\n mulbyA!(y, x) = (y .= A * x)\n\n tol = sqrt(eps(T))\n method(b, tol) = ConjugateGradient(b, max_iter = n)\n linearsolver = method(b, tol)\n\n x = ones(n) * 1.0\n x0 = copy(x)\n iters = linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)\n exact = A \\ b\n x0 = copy(x)\n\n @testset \"Array test\" begin\n @test norm(x - exact) / norm(exact) < err_thresh\n @test norm(A * x - b) / norm(b) < err_thresh\n end\n\n # Testing for CUDA CuArrays\n if CUDA.has_cuda_gpu()\n at_A = ArrayType(A)\n at_b = ArrayType(b)\n at_x = ArrayType(ones(n) * 1.0)\n mulbyat_A!(y, x) = (y .= at_A * x)\n at_method(b, tol) = ConjugateGradient(b, max_iter = n)\n linearsolver = at_method(at_b, tol)\n iters = linearsolve!(\n mulbyat_A!,\n nothing,\n linearsolver,\n at_x,\n at_b;\n max_iters = n,\n )\n exact = at_A \\ at_b\n\n @testset \"CuArray test\" begin\n @test norm(at_x - exact) / norm(exact) < err_thresh\n @test norm(at_A * at_x - at_b) / norm(at_b) < err_thresh\n end\n end\n\n mpi_b = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)\n mpi_x = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)\n mpi_A = ArrayType(randn(4^3, 4^3))\n mpi_A .= mpi_A' * mpi_A\n\n function mpi_mulby!(x, y)\n fy = y.data[:]\n fx = mpi_A * fy\n x.data[:] .= fx[:]\n return nothing\n end\n\n mpi_b.data[:] .= ArrayType(randn(4^3))\n mpi_x.data[:] .= ArrayType(randn(4^3))\n\n mpi_method(mpi_b, tol) = ConjugateGradient(mpi_b, max_iter = n)\n linearsolver = mpi_method(mpi_b, tol)\n iters = linearsolve!(\n mpi_mulby!,\n nothing,\n linearsolver,\n mpi_x,\n mpi_b;\n max_iters = n,\n )\n\n exact = mpi_A \\ mpi_b[:]\n\n @testset \"MPIStateArray test\" begin\n @test norm(mpi_x.data[:] - exact) / norm(exact) < err_thresh\n mpi_Ax = MPIStateArray{T}(mpicomm, ArrayType, 4, 4, 4)\n mpi_mulby!(mpi_Ax, mpi_x)\n @test norm(mpi_Ax - mpi_b) / norm(mpi_b) < err_thresh\n end\n\n # ## More Complex Example\n function closure_linear_operator!(A, tup)\n function linear_operator!(y, x)\n alias_x = reshape(x, tup)\n alias_y = reshape(y, tup)\n for i6 in 1:tup[6]\n for i4 in 1:tup[4]\n for i2 in 1:tup[2]\n for i1 in 1:tup[1]\n tmp = alias_x[i1, i2, :, i4, :, i6][:]\n tmp2 = A[i1, i2, i4, i6] * tmp\n alias_y[i1, i2, :, i4, :, i6] .=\n reshape(tmp2, (tup[3], tup[5]))\n end\n end\n end\n end\n end\n end\n\n tup = (3, 4, 7, 2, 20, 2)\n\n B = [\n randn(tup[3] * tup[5], tup[3] * tup[5])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n columnwise_A = [\n B[i1, i2, i4, i6] * B[i1, i2, i4, i6]' + 10I\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n columnwise_inv_A = [\n inv(columnwise_A[i1, i2, i4, i6])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n ]\n columnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup)\n columnwise_inverse_linear_operator! =\n closure_linear_operator!(columnwise_inv_A, tup)\n\n mpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6])\n b = randn(mpi_tup)\n x = randn(mpi_tup)\n\n linearsolver = ConjugateGradient(\n x,\n max_iter = tup[3] * tup[5],\n dims = (3, 5),\n reshape_tuple = tup,\n )\n\n iters =\n linearsolve!(columnwise_linear_operator!, nothing, linearsolver, x, b)\n x_exact = copy(x)\n columnwise_inverse_linear_operator!(x_exact, b)\n\n @testset \"Columnwise test\" begin\n @test norm(x - x_exact) / norm(x_exact) < err_thresh\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/SystemSolvers/columnwiselu.jl", "content": "using MPI\nusing Test\nusing LinearAlgebra\nusing Random\nusing KernelAbstractions, StaticArrays\n\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.SystemSolvers:\n band_lu_kernel!,\n band_forward_kernel!,\n band_back_kernel!,\n DGColumnBandedMatrix,\n lower_bandwidth,\n upper_bandwidth\n\nimport ClimateMachine.MPIStateArrays: array_device\n\nClimateMachine.init()\nconst ArrayType = ClimateMachine.array_type()\n\nfunction band_to_full(B, p, q)\n _, n = size(B)\n\n A = similar(B, n, n) # assume square\n fill!(A, 0)\n for j in 1:n, i in max(1, j - q):min(j + p, n)\n A[i, j] = B[q + i - j + 1, j]\n end\n A\nend\n\nfunction run_columnwiselu_test(FT, N)\n\n Nq = N .+ 1\n Nq1 = Nq[1]\n Nq2 = Nq[2]\n Nqv = Nq[3]\n nstate = 3\n nvertelem = 5\n nhorzelem = 4\n eband = 2\n\n m = n = Nqv * nstate * nvertelem\n p = lower_bandwidth(N[end], nstate, eband)\n q = upper_bandwidth(N[end], nstate, eband)\n\n Random.seed!(1234)\n AB = rand(FT, Nq1, Nq2, p + q + 1, n, nhorzelem)\n\n AB[:, :, q + 1, :, :] .+= 10 # Make A's diagonally dominate\n\n Random.seed!(5678)\n b = rand(FT, Nq1, Nq2, Nqv, nstate, nvertelem, nhorzelem)\n x = similar(b)\n\n perm = (4, 3, 5, 1, 2, 6)\n bp = reshape(PermutedDimsArray(b, perm), n, Nq1, Nq2, nhorzelem)\n xp = reshape(PermutedDimsArray(x, perm), n, Nq1, Nq2, nhorzelem)\n\n d_F = ArrayType(AB)\n d_F = DGColumnBandedMatrix{\n 3,\n N,\n nstate,\n nhorzelem,\n nvertelem,\n eband,\n false,\n typeof(d_F),\n }(\n d_F,\n )\n\n groupsize = (Nq1, Nq2)\n ndrange = (Nq1, Nq2, nhorzelem)\n\n event = Event(array_device(d_F.data))\n event = band_lu_kernel!(array_device(d_F.data), groupsize, ndrange)(\n d_F,\n dependencies = (event,),\n )\n wait(array_device(d_F.data), event)\n\n F = Array(d_F.data)\n\n for h in 1:nhorzelem, j in 1:Nq2, i in 1:Nq1\n B = AB[i, j, :, :, h]\n G = band_to_full(B, p, q)\n GLU = lu!(G, Val(false))\n\n H = band_to_full(F[i, j, :, :, h], p, q)\n\n @assert H ≈ G\n\n xp[:, i, j, h] .= GLU \\ bp[:, i, j, h]\n end\n\n b = reshape(b, Nq1 * Nq2 * Nqv, nstate, nvertelem * nhorzelem)\n x = reshape(x, Nq1 * Nq2 * Nqv, nstate, nvertelem * nhorzelem)\n\n d_x = ArrayType(b)\n\n event = Event(array_device(d_x))\n event = band_forward_kernel!(array_device(d_x), groupsize, ndrange)(\n d_x,\n d_F,\n dependencies = (event,),\n )\n\n event = band_back_kernel!(array_device(d_x), groupsize, ndrange)(\n d_x,\n d_F,\n dependencies = (event,),\n )\n wait(array_device(d_x), event)\n\n result = x ≈ Array(d_x)\n return result\nend\n\n@testset \"Columnwise LU test\" begin\n for FT in (Float64, Float32)\n for N in ((1, 1, 1), (1, 1, 2), (2, 2, 1), (2, 2, 0))\n result = run_columnwiselu_test(FT, N)\n @test result\n end\n end\nend\n" }, { "path": "test/Numerics/SystemSolvers/iterativesolvers.jl", "content": "using Test\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\n\nusing StaticArrays, LinearAlgebra, Random\n\n# this test setup is partly based on IterativeSolvers.jl, see e.g\n# https://github.com/JuliaMath/IterativeSolvers.jl/blob/master/test/cg.jl\n@testset \"SystemSolvers small full system\" begin\n n = 10\n\n methods = (\n (b, tol) -> GeneralizedConjugateResidual(2, b, rtol = tol),\n (b, tol) -> GeneralizedMinimalResidual(b, M = 3, rtol = tol),\n (b, tol) -> GeneralizedMinimalResidual(b, M = n, rtol = tol),\n )\n\n expected_iters = (\n Dict(Float32 => 7, Float64 => 11),\n Dict(Float32 => 5, Float64 => 17),\n Dict(Float32 => 4, Float64 => 10),\n )\n\n for (m, method) in enumerate(methods), T in [Float32, Float64]\n Random.seed!(44)\n\n A = @MMatrix rand(T, n, n)\n A = A' * A + I\n b = @MVector rand(T, n)\n\n mulbyA!(y, x) = (y .= A * x)\n\n tol = sqrt(eps(T))\n linearsolver = method(b, tol)\n\n x = @MVector rand(T, n)\n x0 = copy(x)\n iters =\n linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)\n\n @test iters == expected_iters[m][T]\n @test norm(A * x - b) / norm(A * x0 - b) <= tol\n\n # test for convergence in 0 iterations by\n # initializing with the exact solution\n x = A \\ b\n\n iters =\n linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)\n @test iters == 0\n @test norm(A * x - b) <= 100 * eps(T)\n\n\n newtol = 1000tol\n settolerance!(linearsolver, newtol)\n\n x = @MVector rand(T, n)\n x0 = copy(x)\n linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)\n\n @test norm(A * x - b) / norm(A * x0 - b) <= newtol\n @test norm(A * x - b) / norm(A * x0 - b) >= tol\n\n end\nend\n\n@testset \"SystemSolvers large full system\" begin\n n = 1000\n\n methods = (\n (b, tol) -> GeneralizedMinimalResidual(b, M = 15, rtol = tol),\n (b, tol) -> GeneralizedMinimalResidual(b, M = 20, rtol = tol),\n )\n\n expected_iters = (\n Dict(Float32 => (3, 3), Float64 => (9, 8)),\n Dict(Float32 => (3, 3), Float64 => (9, 8)),\n )\n\n for (m, method) in enumerate(methods), T in [Float32, Float64]\n for (i, α) in enumerate(T[1e-2, 5e-3])\n Random.seed!(44)\n A = rand(T, 200, 1000)\n A = α * A' * A + I\n b = rand(T, n)\n\n mulbyA!(y, x) = (y .= A * x)\n\n tol = sqrt(eps(T))\n linearsolver = method(b, tol)\n\n x = rand(T, n)\n x0 = copy(x)\n iters = linearsolve!(\n mulbyA!,\n nothing,\n linearsolver,\n x,\n b;\n max_iters = Inf,\n )\n\n @test iters == expected_iters[m][T][i]\n @test norm(A * x - b) / norm(A * x0 - b) <= tol\n\n newtol = 1000tol\n settolerance!(linearsolver, newtol)\n\n x = rand(T, n)\n x0 = copy(x)\n linearsolve!(mulbyA!, nothing, linearsolver, x, b; max_iters = Inf)\n\n @test norm(A * x - b) / norm(A * x0 - b) <= newtol\n end\n end\nend\n" }, { "path": "test/Numerics/SystemSolvers/poisson.jl", "content": "using MPI\nusing Test\nusing StaticArrays\nusing Logging, Printf\n\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.DGMethods\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n flux_first_order!,\n flux_second_order!,\n source!,\n boundary_conditions,\n boundary_state!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!\n\nimport ClimateMachine.DGMethods: numerical_boundary_flux_second_order!\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\n\nimport ClimateMachine.DGMethods.NumericalFluxes:\n NumericalFluxSecondOrder, numerical_flux_second_order!\n\nif !@isdefined integration_testing\n const integration_testing = parse(\n Bool,\n lowercase(get(ENV, \"JULIA_CLIMA_INTEGRATION_TESTING\", \"false\")),\n )\nend\n\nstruct PoissonModel{dim} <: BalanceLaw end\n\nvars_state(::PoissonModel, ::Auxiliary, T) = @vars(rhs_ϕ::T)\nvars_state(::PoissonModel, ::Prognostic, T) = @vars(ϕ::T)\nvars_state(::PoissonModel, ::Gradient, T) = @vars(ϕ::T)\nvars_state(::PoissonModel, ::GradientFlux, T) = @vars(∇ϕ::SVector{3, T})\n\nboundary_conditions(::PoissonModel) = ()\nboundary_state!(nf, bc, bl::PoissonModel, _...) = nothing\n\nfunction flux_first_order!(::PoissonModel, _...) end\n\nfunction flux_second_order!(\n ::PoissonModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n state_auxiliary::Vars,\n t::Real,\n)\n flux.ϕ = diffusive.∇ϕ\nend\n\nstruct PenaltyNumFluxDiffusive <: NumericalFluxSecondOrder end\n\n# There is no boundary since we are periodic\nnumerical_boundary_flux_second_order!(nf::PenaltyNumFluxDiffusive, _...) =\n nothing\n\nfunction numerical_flux_second_order!(\n ::PenaltyNumFluxDiffusive,\n bl::PoissonModel,\n fluxᵀn::Vars{S},\n n::SVector,\n state⁻::Vars{S},\n diff⁻::Vars{D},\n hyperdiff⁻::Vars{HD},\n aux⁻::Vars{A},\n state⁺::Vars{S},\n diff⁺::Vars{D},\n hyperdiff⁺::Vars{HD},\n aux⁺::Vars{A},\n t,\n) where {S, HD, D, A}\n\n numerical_flux_second_order!(\n CentralNumericalFluxSecondOrder(),\n bl,\n fluxᵀn,\n n,\n state⁻,\n diff⁻,\n hyperdiff⁻,\n aux⁻,\n state⁺,\n diff⁺,\n hyperdiff⁺,\n aux⁺,\n t,\n )\n\n Fᵀn = parent(fluxᵀn)\n FT = eltype(Fᵀn)\n tau = FT(1)\n Fᵀn .-= tau * (parent(state⁻) - parent(state⁺))\nend\n\nfunction compute_gradient_argument!(\n ::PoissonModel,\n transformstate::Vars,\n state::Vars,\n state_auxiliary::Vars,\n t::Real,\n)\n transformstate.ϕ = state.ϕ\nend\n\nfunction compute_gradient_flux!(\n ::PoissonModel,\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n state_auxiliary::Vars,\n t::Real,\n)\n diffusive.∇ϕ = ∇transform.ϕ\nend\n\nsource!(::PoissonModel, _...) = nothing\n\n# note, that the code assumes solutions with zero mean\nsol1d(x) = sin(2pi * x)^4 - 3 / 8\ndxx_sol1d(x) =\n -16 * pi^2 * sin(2pi * x)^2 * (sin(2pi * x)^2 - 3 * cos(2pi * x)^2)\n\nfunction nodal_init_state_auxiliary!(\n ::PoissonModel{dim},\n aux::Vars,\n tmp::Vars,\n g::LocalGeometry,\n) where {dim}\n aux.rhs_ϕ = 0\n @inbounds for d in 1:dim\n x1 = g.coord[d]\n x2 = g.coord[1 + mod(d, dim)]\n x3 = g.coord[1 + mod(d + 1, dim)]\n x23 = SVector(x2, x3)\n aux.rhs_ϕ -= dxx_sol1d(x1) * prod(sol1d, view(x23, 1:(dim - 1)))\n end\nend\n\nfunction init_state_prognostic!(\n ::PoissonModel{dim},\n state::Vars,\n aux::Vars,\n localgeo,\n t,\n) where {dim}\n coords = localgeo.coord\n state.ϕ = prod(sol1d, view(coords, 1:dim))\nend\n\nfunction test_run(\n mpicomm,\n ArrayType,\n ::Type{FT},\n dim,\n polynomialorder,\n brickrange,\n periodicity,\n linmethod,\n) where {FT}\n\n topology = BrickTopology(mpicomm, brickrange, periodicity = periodicity)\n grid = DiscontinuousSpectralElementGrid(\n topology,\n polynomialorder = polynomialorder,\n DeviceArray = ArrayType,\n FloatType = FT,\n )\n\n dg = DGModel(\n PoissonModel{dim}(),\n grid,\n CentralNumericalFluxFirstOrder(),\n PenaltyNumFluxDiffusive(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n Qrhs = dg.state_auxiliary\n Qexact = init_ode_state(dg, FT(0))\n\n linearoperator!(y, x) = dg(y, x, nothing, 0; increment = false)\n\n linearsolver = linmethod(Q)\n\n iters = linearsolve!(linearoperator!, nothing, linearsolver, Q, Qrhs)\n\n error = euclidean_distance(Q, Qexact)\n\n @info @sprintf \"\"\"Finished\n error = %.16e\n iters = %d\n \"\"\" error iters\n error, iters\nend\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n\n mpicomm = MPI.COMM_WORLD\n\n polynomialorder = 4\n base_num_elem = 4\n tol = 1e-9\n\n linmethods = (\n b -> GeneralizedConjugateResidual(3, b, rtol = tol),\n b -> GeneralizedMinimalResidual(b, M = 7, rtol = tol),\n )\n\n expected_result = Array{Float64}(undef, 2, 2, 3) # method, dim-1, lvl\n\n # GeneralizedConjugateResidual\n expected_result[1, 1, 1] = 5.0540243616448058e-02\n expected_result[1, 1, 2] = 1.4802275366044011e-03\n expected_result[1, 1, 3] = 3.3852821775121401e-05\n expected_result[1, 2, 1] = 1.4957957657736219e-02\n expected_result[1, 2, 2] = 4.7282369781541172e-04\n expected_result[1, 2, 3] = 1.4697449643351771e-05\n\n # GeneralizedMinimalResidual\n expected_result[2, 1, 1] = 5.0540243587512981e-02\n expected_result[2, 1, 2] = 1.4802275409186211e-03\n expected_result[2, 1, 3] = 3.3852820667079927e-05\n expected_result[2, 2, 1] = 1.4957957659220951e-02\n expected_result[2, 2, 2] = 4.7282369895963614e-04\n expected_result[2, 2, 3] = 1.4697449516628483e-05\n\n lvls = integration_testing ? size(expected_result)[end] : 1\n\n for (m, linmethod) in enumerate(linmethods), FT in (Float64,)\n result = Array{Tuple{FT, Int}}(undef, lvls)\n for dim in 2:3\n for l in 1:lvls\n Ne = ntuple(d -> 2^(l - 1) * base_num_elem, dim)\n brickrange =\n ntuple(d -> range(FT(0), length = Ne[d], stop = 1), dim)\n periodicity = ntuple(d -> true, dim)\n\n @info (ArrayType, FT, m, dim)\n result[l] = test_run(\n mpicomm,\n ArrayType,\n FT,\n dim,\n polynomialorder,\n brickrange,\n periodicity,\n linmethod,\n )\n\n @test isapprox(\n result[l][1],\n FT(expected_result[m, dim - 1, l]),\n rtol = sqrt(tol),\n )\n end\n\n if integration_testing\n @info begin\n msg = \"\"\n for l in 1:(lvls - 1)\n rate = log2(result[l][1]) - log2(result[l + 1][1])\n msg *= @sprintf(\"\\n rate for level %d = %e\\n\", l, rate)\n end\n msg\n end\n end\n end\n end\nend\n\nnothing\n" }, { "path": "test/Numerics/SystemSolvers/runtests.jl", "content": "using Test, MPI\ninclude(joinpath(\"..\", \"..\", \"testhelpers.jl\"))\n\ninclude(\"iterativesolvers.jl\")\n" }, { "path": "test/Numerics/runtests.jl", "content": "using Test, Pkg\n\n@testset \"Numerics\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"SystemSolvers\", \"ODESolvers\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"Numerics/$submodule\" in ARGS ||\n \"Numerics\" in ARGS\n include_test(submodule)\n end\n end\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_3D_spindown.jl", "content": "#!/usr/bin/env julia --project\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.MPIStateArrays: euclidean_distance, norm\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.BalanceLaws: parameter_set\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Grids: polynomialorders\nusing ClimateMachine.Ocean.HydrostaticBoussinesq\nusing ClimateMachine.Ocean.OceanProblems\n\nusing Test\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nfunction config_simple_box(\n ::Type{FT},\n N,\n resolution,\n dimensions;\n BC = nothing,\n) where {FT}\n if BC == nothing\n problem = SimpleBox{FT}(dimensions...)\n else\n problem = SimpleBox{FT}(dimensions...; BC = BC)\n end\n\n model = HydrostaticBoussinesqModel{FT}(\n param_set,\n problem;\n cʰ = FT(1),\n αᵀ = FT(0),\n κʰ = FT(0),\n κᶻ = FT(0),\n fₒ = FT(0),\n β = FT(0),\n )\n solver_type = ExplicitSolverType(solver_method = LSRK144NiegemannDiehlBusch)\n config = ClimateMachine.OceanBoxGCMConfiguration(\n \"hydrostatic_spindown\",\n N,\n resolution,\n param_set,\n model;\n periodicity = (true, true, false),\n boundary = ((0, 0), (0, 0), (1, 2)),\n )\n\n return config, solver_type\nend\n\nfunction run_hydrostatic_test(; imex::Bool = false, BC = nothing, refDat = ())\n FT = Float64\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution and size\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(8)\n resolution = (Nˣ, Nʸ, Nᶻ)\n\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n timestart = FT(0) # s\n timeout = FT(6 * 3600) # s\n timeend = FT(86400) # s\n dt = FT(120) # s\n\n if imex\n solver_type =\n ClimateMachine.IMEXSolverType(implicit_model = LinearHBModel)\n else\n solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n end\n\n driver, solver_type_driver_config =\n config_simple_box(FT, N, resolution, dimensions; BC = BC)\n\n grid = driver.grid\n N = polynomialorders(grid)\n Nvert = N[end]\n vert_filter = CutoffFilter(grid, Nvert - 1)\n exp_filter = ExponentialFilter(grid, 1, 8)\n modeldata = (vert_filter = vert_filter, exp_filter = exp_filter)\n\n solver = ClimateMachine.SolverConfiguration(\n timestart,\n timeend,\n driver,\n init_on_cpu = true,\n ode_dt = dt,\n ode_solver_type = solver_type,\n modeldata = modeldata,\n )\n\n output_interval = ceil(Int64, timeout / solver.dt)\n\n ClimateMachine.Settings.vtk = \"never\"\n # ClimateMachine.Settings.vtk = \"$(output_interval)steps\"\n\n ClimateMachine.Settings.diagnostics = \"never\"\n # ClimateMachine.Settings.diagnostics = \"$(output_interval)steps\"\n\n cb = ClimateMachine.StateCheck.sccreate(\n [(solver.Q, \"state\"), (solver.dg.state_auxiliary, \"aux\")],\n output_interval;\n prec = 12,\n )\n\n result = ClimateMachine.invoke!(solver; user_callbacks = [cb])\n\n Q_exact = ClimateMachine.DGMethods.init_ode_state(\n solver.dg,\n timeend,\n solver.init_args...;\n init_on_cpu = solver.init_on_cpu,\n )\n\n error = euclidean_distance(solver.Q, Q_exact) / norm(Q_exact)\n\n println(\"error = \", error)\n @test isapprox(error, FT(0.0); atol = 0.005)\n\n ## Check results against reference\n ClimateMachine.StateCheck.scprintref(cb)\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(cb, refDat)\n end\nend\n\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/3D_hydrostatic_spindown_refvals.jl\")\n\n run_hydrostatic_test(imex = false, refDat = refVals.explicit) # error = 0.0011289879366523504\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_initial_value_problem.jl", "content": "using ClimateMachine\n\nClimateMachine.init()\n\nusing ClimateMachine.Ocean.OceanProblems: InitialConditions, InitialValueProblem\n\n@testset \"$(@__FILE__)\" begin\n U = 0.1\n L = 0.2\n a = 0.3\n U = 0.4\n\n Ψ(x, L) = exp(-x^2 / 2 * L^2) # a Gaussian\n\n uᵢ(x, y, z) = +U * y / L * Ψ(x, L)\n vᵢ(x, y, z) = -U * x / L * Ψ(x, L)\n ηᵢ(x, y, z) = a * Ψ(x, L)\n θᵢ(x, y, z) = 20.0 + 1e-3 * z\n\n initial_conditions = InitialConditions(u = uᵢ, v = vᵢ, η = ηᵢ, θ = θᵢ)\n\n problem = InitialValueProblem{Float64}(\n dimensions = (π, 42, 1.1),\n initial_conditions = initial_conditions,\n )\n\n @test problem.Lˣ == Float64(π)\n @test problem.Lʸ == 42.0\n @test problem.H == 1.1\n\n @test problem.initial_conditions.u === uᵢ\n @test problem.initial_conditions.v === vᵢ\n @test problem.initial_conditions.η === ηᵢ\n @test problem.initial_conditions.θ === θᵢ\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_ocean_gyre_long.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../../../experiments/OceanBoxGCM/simple_box.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/test_ocean_gyre_refvals.jl\")\n\n # simulation time\n timestart = FT(0) # s\n timeend = FT(86400) # s\n timespan = (timestart, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(20)\n Nʸ = Int(20)\n Nᶻ = Int(20)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 4e6 # m\n Lʸ = 4e6 # m\n H = 1000 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), TemperatureFlux()),\n )\n\n run_simple_box(\n \"ocean_gyre_long\",\n resolution,\n dimensions,\n timespan,\n OceanGyre,\n imex = false,\n BC = BC,\n refDat = refVals.long,\n )\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_ocean_gyre_short.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../../../experiments/OceanBoxGCM/simple_box.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/test_ocean_gyre_refvals.jl\")\n\n # simulation time\n timestart = FT(0) # s\n timeend = FT(3600) # s\n timespan = (timestart, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(5)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 1000 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), TemperatureFlux()),\n )\n\n run_simple_box(\n \"ocean_gyre_short\",\n resolution,\n dimensions,\n timespan,\n OceanGyre,\n imex = false,\n BC = BC,\n Δt = 120,\n refDat = refVals.short,\n )\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_windstress_long.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../../../experiments/OceanBoxGCM/simple_box.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/test_windstress_refvals.jl\")\n\n # simulation time\n timestart = FT(0) # s\n timeend = FT(86400) # s\n timespan = (timestart, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(20)\n Nʸ = Int(20)\n Nᶻ = Int(50)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 4e6 # m\n Lʸ = 4e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), Insulating()),\n )\n\n run_simple_box(\n \"test_windstress_long_imex\",\n resolution,\n dimensions,\n timespan,\n HomogeneousBox,\n imex = true,\n BC = BC,\n Δt = 55,\n refDat = refVals.long,\n )\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesq/test_windstress_short.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"../../../experiments/OceanBoxGCM/simple_box.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/test_windstress_refvals.jl\")\n\n # simulation time\n timestart = FT(0) # s\n timeend = FT(3600) # s\n timespan = (timestart, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(5)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(NoSlip()), Insulating()),\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(KinematicStress()), Insulating()),\n )\n\n run_simple_box(\n \"test_windstress_short_imex\",\n resolution,\n dimensions,\n timespan,\n HomogeneousBox,\n imex = true,\n BC = BC,\n Δt = 60,\n refDat = refVals.imex,\n )\n\n run_simple_box(\n \"test_windstress_short_explicit\",\n resolution,\n dimensions,\n timespan,\n HomogeneousBox,\n imex = false,\n BC = BC,\n Δt = 180,\n refDat = ClimateMachine.array_type() == Array ? refVals.explicit_cpu :\n refVals.explicit_gpu,\n )\nend\n" }, { "path": "test/Ocean/HydrostaticBoussinesqModel/test_hydrostatic_boussinesq_model.jl", "content": "using ClimateMachine\n\nClimateMachine.init()\n\nusing ClimateMachine.CartesianDomains\n\nusing CLIMAParameters: AbstractEarthParameterSet\nstruct DefaultParameters <: AbstractEarthParameterSet end\n\nusing ClimateMachine.Ocean:\n HydrostaticBoussinesqSuperModel, current_time, current_step, Δt\n\n@testset \"$(@__FILE__)\" begin\n\n domain = RectangularDomain(\n Ne = (1, 1, 1),\n Np = 4,\n x = (-1, 1),\n y = (-1, 1),\n z = (-1, 0),\n )\n\n model = HydrostaticBoussinesqSuperModel(\n domain = domain,\n time_step = 0.1,\n parameters = DefaultParameters(),\n )\n\n @test model isa HydrostaticBoussinesqSuperModel\n @test Δt(model) == 0.1\n @test current_step(model) == 0\n @test current_time(model) == 0.0\nend\n" }, { "path": "test/Ocean/OceanProblems/test_initial_value_problem.jl", "content": "using ClimateMachine\n\nClimateMachine.init()\n\nusing ClimateMachine.Ocean.OceanProblems: InitialConditions, InitialValueProblem\n\n@testset \"$(@__FILE__)\" begin\n U = 0.1\n L = 0.2\n a = 0.3\n U = 0.4\n\n Ψ(x, L) = exp(-x^2 / 2 * L^2) # a Gaussian\n\n uᵢ(x, y, z) = +U * y / L * Ψ(x, L)\n vᵢ(x, y, z) = -U * x / L * Ψ(x, L)\n ηᵢ(x, y, z) = a * Ψ(x, L)\n θᵢ(x, y, z) = 20.0 + 1e-3 * z\n\n initial_conditions = InitialConditions(u = uᵢ, v = vᵢ, η = ηᵢ, θ = θᵢ)\n\n problem = InitialValueProblem{Float64}(\n dimensions = (π, 42, 1.1),\n initial_conditions = initial_conditions,\n )\n\n @test problem.Lˣ == Float64(π)\n @test problem.Lʸ == 42.0\n @test problem.H == 1.1\n\n @test problem.initial_conditions.u === uᵢ\n @test problem.initial_conditions.v === vᵢ\n @test problem.initial_conditions.η === ηᵢ\n @test problem.initial_conditions.θ === θᵢ\nend\n" }, { "path": "test/Ocean/ShallowWater/GyreDriver.jl", "content": "using MPI\nusing Test\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.VariableTemplates: flattenednames\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.Ocean.ShallowWater\nusing ClimateMachine.Ocean.ShallowWater:\n TurbulenceClosure, LinearDrag, ConstantViscosity\nusing ClimateMachine.Ocean\nusing ClimateMachine.Ocean.OceanProblems\n\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\nusing ClimateMachine.VTK\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nif !isempty(ARGS)\n stommel = Bool(parse(Int, ARGS[1]))\n linear = Bool(parse(Int, ARGS[2]))\n test = parse(Int, ARGS[3])\nelse\n stommel = true\n linear = true\n test = 1\nend\n\n###################\n# PARAM SELECTION #\n###################\nconst FT = Float64\n\nconst τₒ = 2e-4 # value includes τₒ, g, and ρ\nconst fₒ = 1e-4\nconst β = 1e-11\nconst λ = 1e-6\nconst ν = 1e4\n\nconst Lˣ = 1e6\nconst Lʸ = 1e6\nconst timeend = 100 * 24 * 60 * 60\nconst H = 1000\nconst c = sqrt(grav(param_set) * H)\n\nif stommel\n gyre = \"stommel\"\nelse\n gyre = \"munk\"\nend\n\nif linear\n advec = \"linear\"\nelse\n advec = \"nonlinear\"\nend\n\noutname = \"vtk_new_dt_\" * gyre * \"_\" * advec\n\nfunction setup_model(\n ::Type{FT},\n stommel,\n linear,\n τₒ,\n fₒ,\n β,\n γ,\n ν,\n Lˣ,\n Lʸ,\n H,\n) where {FT}\n problem = HomogeneousBox{FT}(\n Lˣ,\n Lʸ,\n H,\n τₒ = τₒ,\n BC = (OceanBC(Impenetrable(FreeSlip()), Insulating()),),\n )\n\n\n if stommel\n turbulence = LinearDrag{FT}(λ)\n else\n turbulence = ConstantViscosity{FT}(ν)\n end\n\n if linear\n advection = nothing\n else\n advection = NonLinearAdvectionTerm()\n end\n\n model = ShallowWaterModel{FT}(\n param_set,\n problem,\n turbulence,\n advection,\n c = c,\n fₒ = fₒ,\n β = β,\n )\nend\n\n#########################\n# Timestepping function #\n#########################\n\nfunction test_run(\n mpicomm,\n topl,\n ArrayType,\n N,\n dt,\n ::Type{FT},\n model,\n test,\n) where {FT}\n grid = DiscontinuousSpectralElementGrid(\n topl,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n dg = DGModel(\n model,\n grid,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q = init_ode_state(dg, FT(0))\n Qe = init_ode_state(dg, FT(timeend))\n\n lsrk = LSRK144NiegemannDiehlBusch(dg, Q; dt = dt, t0 = 0)\n\n nt_freq = floor(Int, 1 // 10 * timeend / dt)\n\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [(Q, \"2D state\")],\n nt_freq;\n prec = 12,\n )\n\n cb = (cbcs_dg,)\n\n if test > 2\n outprefix = @sprintf(\"ic_mpirank%04d_ic\", MPI.Comm_rank(mpicomm))\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n outprefix = @sprintf(\"exact_mpirank%04d\", MPI.Comm_rank(mpicomm))\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Qe)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Qe)))\n writevtk(outprefix, Qe, dg, statenames, dg.state_auxiliary, auxnames)\n\n vtkstep = [0]\n vtkpath = outname\n mkpath(vtkpath)\n cbvtk = GenericCallbacks.EveryXSimulationSteps(1000) do\n outprefix = @sprintf(\n \"%s/mpirank%04d_step%04d\",\n vtkpath,\n MPI.Comm_rank(mpicomm),\n vtkstep[1]\n )\n @debug \"doing VTK output\" outprefix\n statenames =\n flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxiliarynames =\n flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(\n outprefix,\n Q,\n dg,\n statenames,\n dg.state_auxiliary,\n auxiliarynames,\n )\n vtkstep[1] += 1\n nothing\n end\n cb = (cb..., cbvtk)\n end\n\n solve!(Q, lsrk; timeend = timeend, callbacks = cb)\n\n error = euclidean_distance(Q, Qe) / norm(Qe)\n @info @sprintf(\n \"\"\"Finished\n error = %.16e\n \"\"\",\n error\n )\n\n return error\nend\n\n\n################\n# Start Driver #\n################\n\nlet\n ClimateMachine.init()\n ArrayType = ClimateMachine.array_type()\n mpicomm = MPI.COMM_WORLD\n\n model = setup_model(FT, stommel, linear, τₒ, fₒ, β, λ, ν, Lˣ, Lʸ, H)\n\n if test == 1\n cellsrange = 10:10\n orderrange = 4:4\n testval = 1.6068814534535144e-03\n elseif test == 2\n cellsrange = 5:5:10\n orderrange = 3:4\n elseif test > 2\n cellsrange = 5:5:25\n orderrange = 6:10\n end\n\n errors = zeros(FT, length(cellsrange), length(orderrange))\n for (i, Ne) in enumerate(cellsrange)\n brickrange = (\n range(FT(0); length = Ne + 1, stop = Lˣ),\n range(FT(0); length = Ne + 1, stop = Lʸ),\n )\n topl = BrickTopology(\n mpicomm,\n brickrange,\n periodicity = (false, false),\n boundary = ((1, 1), (1, 1)),\n )\n\n for (j, N) in enumerate(orderrange)\n @info \"running Ne $Ne and N $N with\"\n dt = (Lˣ / c) / Ne / N^2\n @info @sprintf(\"\\n dt = %f\", dt)\n errors[i, j] =\n test_run(mpicomm, topl, ArrayType, N, dt, FT, model, test)\n end\n end\n\n @test errors[end, end] ≈ testval\n\n #=\n msg = \"\"\n for i in length(cellsrange)-1\n rate = log2(errors[i, end] - log2(errors[i+1, end]))\n msg *= @sprintf(\"\\n rate for Ne %d = %e\", cellsrange[i], rate)\n end\n @info msg\n\n msg = \"\"\n for j in length(orderrange)-1\n rate = log2(errors[end, j] - log2(errors[end, j+1]))\n msg *= @sprintf(\"\\n rate for N %d = %e\", orderrange[j], rate)\n end\n @info msg\n =#\n\nend\n" }, { "path": "test/Ocean/ShallowWater/test_2D_spindown.jl", "content": "#!/usr/bin/env julia --project\nusing Test\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Grids: polynomialorders\nusing ClimateMachine.Ocean.ShallowWater\nusing ClimateMachine.Ocean.OceanProblems\n\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VTK\n\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nfunction run_hydrostatic_spindown(; refDat = ())\n mpicomm = MPI.COMM_WORLD\n ArrayType = ClimateMachine.array_type()\n\n ll = uppercase(get(ENV, \"JULIA_LOG_LEVEL\", \"INFO\"))\n loglevel =\n ll == \"DEBUG\" ? Logging.Debug :\n ll == \"WARN\" ? Logging.Warn :\n ll == \"ERROR\" ? Logging.Error : Logging.Info\n logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull\n global_logger(ConsoleLogger(logger_stream, loglevel))\n\n brickrange_2D = (xrange, yrange)\n topl_2D = BrickTopology(\n mpicomm,\n brickrange_2D,\n periodicity = (true, true),\n boundary = ((0, 0), (0, 0)),\n )\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n problem = SimpleBox{FT}(Lˣ, Lʸ, H)\n\n model_2D = ShallowWaterModel{FT}(\n param_set,\n problem,\n ShallowWater.ConstantViscosity{FT}(5e3),\n nothing;\n c = FT(1),\n fₒ = FT(0),\n β = FT(0),\n )\n\n dt_fast = 300\n nout = ceil(Int64, tout / dt_fast)\n dt_fast = tout / nout\n\n dg_2D = DGModel(\n model_2D,\n grid_2D,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_2D = init_ode_state(dg_2D, FT(0); init_on_cpu = true)\n\n lsrk_2D = LSRK54CarpenterKennedy(dg_2D, Q_2D, dt = dt_fast, t0 = 0)\n\n odesolver = lsrk_2D\n\n vtkstep = [0, 0]\n cbvector = make_callbacks(\n vtkpath,\n vtkstep,\n nout,\n mpicomm,\n odesolver,\n dg_2D,\n model_2D,\n Q_2D,\n )\n\n eng0 = norm(Q_2D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 ArrayType\n\n solve!(Q_2D, odesolver; timeend = timeend, callbacks = cbvector)\n\n Qe_2D = init_ode_state(dg_2D, timeend, init_on_cpu = true)\n\n error_2D = euclidean_distance(Q_2D, Qe_2D) / norm(Qe_2D)\n\n println(\"2D error = \", error_2D)\n @test isapprox(error_2D, FT(0.0); atol = 0.005)\n\n ## Check results against reference\n ClimateMachine.StateCheck.scprintref(cbvector[end])\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(cbvector[end], refDat)\n end\n\n return nothing\nend\n\nfunction make_callbacks(\n vtkpath,\n vtkstep,\n nout,\n mpicomm,\n odesolver,\n dg_fast,\n model_fast,\n Q_fast,\n)\n if isdir(vtkpath)\n rm(vtkpath, recursive = true)\n end\n mkpath(vtkpath)\n mkpath(vtkpath * \"/fast\")\n\n function do_output(span, vtkstep, model, dg, Q)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n end\n\n do_output(\"fast\", vtkstep[2], model_fast, dg_fast, Q_fast)\n cbvtk_fast = GenericCallbacks.EveryXSimulationSteps(nout) do (init = false)\n do_output(\"fast\", vtkstep[2], model_fast, dg_fast, Q_fast)\n vtkstep[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_fast)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [(Q_fast, \"2D state\")],\n nout;\n prec = 12,\n )\n\n # don't write vtk during CI testing\n # return (cbvtk_fast, cbinfo, cbcs_dg)\n return (cbinfo, cbcs_dg)\nend\n\n#################\n# RUN THE TESTS #\n#################\nFT = Float64\nvtkpath = abspath(joinpath(\n ClimateMachine.Settings.output_dir,\n \"vtk_shallow_spindown\",\n))\n\nconst timeend = FT(24 * 3600) # s\nconst tout = FT(2 * 3600) # s\n# const timeend = 1200 # s\n# const tout = 600 # s\n\nconst N = 4\nconst Nˣ = 5\nconst Nʸ = 5\nconst Lˣ = 1e6 # m\nconst Lʸ = 1e6 # m\nconst H = 400 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\n\n@testset \"$(@__FILE__)\" begin\n include(\"../refvals/2D_hydrostatic_spindown_refvals.jl\")\n\n run_hydrostatic_spindown(refDat = refVals.explicit) # error = 0.00011327920483879001\nend\n" }, { "path": "test/Ocean/SplitExplicit/hydrostatic_spindown.jl", "content": "include(\"split_explicit.jl\")\n\nfunction SplitConfig(\n name,\n resolution,\n dimensions,\n coupling,\n rotation = Fixed();\n boundary_conditions = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n ),\n solver = SplitExplicitSolver,\n dt_slow = 90 * 60,\n)\n mpicomm = MPI.COMM_WORLD\n ArrayType = ClimateMachine.array_type()\n\n N, Nˣ, Nʸ, Nᶻ = resolution\n Lˣ, Lʸ, H = dimensions\n\n xrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\n yrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\n zrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n\n brickrange_2D = (xrange, yrange)\n topl_2D = BrickTopology(\n mpicomm,\n brickrange_2D,\n periodicity = (true, true),\n boundary = ((0, 0), (0, 0)),\n )\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (true, true, false),\n boundary = ((0, 0), (0, 0), (1, 2)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n problem = SimpleBox{FT}(\n dimensions...;\n BC = boundary_conditions,\n rotation = rotation,\n )\n\n dg_3D, dg_2D = setup_models(\n solver,\n problem,\n grid_3D,\n grid_2D,\n param_set,\n coupling,\n dt_slow,\n )\n\n return SplitConfig(name, dg_3D, dg_2D, solver, mpicomm, ArrayType)\nend\n\nfunction setup_models(\n ::Type{SplitExplicitSolver},\n problem,\n grid_3D,\n grid_2D,\n param_set,\n coupling,\n _,\n)\n\n model_3D = HydrostaticBoussinesqModel{FT}(\n param_set,\n problem;\n coupling = coupling,\n cʰ = FT(1),\n αᵀ = FT(0),\n κʰ = FT(0),\n κᶻ = FT(0),\n )\n\n model_2D = ShallowWaterModel{FT}(\n param_set,\n problem,\n ShallowWater.ConstantViscosity{FT}(model_3D.νʰ),\n nothing;\n coupling = coupling,\n c = FT(1),\n )\n\n N = polynomialorders(grid_3D)\n Nvert = N[end]\n vert_filter = CutoffFilter(grid_3D, Nvert - 1)\n exp_filter = ExponentialFilter(grid_3D, 1, 8)\n\n integral_model = DGModel(\n VerticalIntegralModel(model_3D),\n grid_3D,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n dg_2D = DGModel(\n model_2D,\n grid_2D,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_2D = init_ode_state(dg_2D, FT(0); init_on_cpu = true)\n\n modeldata = (\n dg_2D = dg_2D,\n Q_2D = Q_2D,\n vert_filter = vert_filter,\n exp_filter = exp_filter,\n integral_model = integral_model,\n )\n\n dg_3D = DGModel(\n model_3D,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n modeldata = modeldata,\n )\n\n return dg_3D, dg_2D\n\nend\n\nfunction setup_models(\n ::Type{SplitExplicitLSRK2nSolver},\n problem,\n grid_3D,\n grid_2D,\n param_set,\n _,\n dt_slow,\n)\n add_fast_substeps = 2\n numImplSteps = 5\n numImplSteps > 0 ? ivdc_dt = dt_slow / FT(numImplSteps) : ivdc_dt = dt_slow\n model_3D = OceanModel{FT}(\n param_set,\n problem,\n cʰ = FT(1),\n αᵀ = FT(0),\n κʰ = FT(0),\n κᶻ = FT(0),\n add_fast_substeps = add_fast_substeps,\n numImplSteps = numImplSteps,\n ivdc_dt = ivdc_dt,\n )\n\n model_2D = BarotropicModel(model_3D)\n\n dg_2D = DGModel(\n model_2D,\n grid_2D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_2D = init_ode_state(dg_2D, FT(0); init_on_cpu = true)\n\n dg_3D = OceanDGModel(\n model_3D,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient();\n modeldata = (dg_2D, Q_2D),\n )\n\n return dg_3D, dg_2D\n\nend\n" }, { "path": "test/Ocean/SplitExplicit/simple_box_2dt.jl", "content": "#!/usr/bin/env julia --project\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates: flattenednames\nusing ClimateMachine.Ocean.SplitExplicit01\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.VTK\nusing ClimateMachine.Checkpoint\n\nusing Test\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport ClimateMachine.Ocean.SplitExplicit01:\n ocean_init_aux!,\n ocean_init_state!,\n ocean_boundary_state!,\n CoastlineFreeSlip,\n CoastlineNoSlip,\n OceanFloorFreeSlip,\n OceanFloorNoSlip,\n OceanSurfaceNoStressNoForcing,\n OceanSurfaceStressNoForcing,\n OceanSurfaceNoStressForcing,\n OceanSurfaceStressForcing\nimport ClimateMachine.DGMethods:\n update_auxiliary_state!, update_auxiliary_state_gradient!, VerticalDirection\n# using GPUifyLoops\n\nconst ArrayType = ClimateMachine.array_type()\n\nstruct SimpleBox{T, BC} <: AbstractOceanProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n λʳ::T\n θᴱ::T\n boundary_conditions::BC\nend\n\nfunction ocean_init_state!(p::SimpleBox, Q, A, localgeo, t)\n coords = localgeo.coord\n @inbounds y = coords[2]\n @inbounds z = coords[3]\n @inbounds H = p.H\n\n Q.u = @SVector [-0, -0]\n Q.η = -0\n Q.θ = (5 + 4 * cos(y * π / p.Lʸ)) * (1 + z / H)\n\n return nothing\nend\n\nfunction ocean_init_aux!(m::OceanModel, p::SimpleBox, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # not sure if this is needed but getting weird intialization stuff\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n A.u_d = @SVector [-0, -0]\n A.ΔGu = @SVector [-0, -0]\n\n return nothing\nend\n\n# A is Filled afer the state\nfunction ocean_init_aux!(m::BarotropicModel, P::SimpleBox, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.U_c = @SVector [-0, -0]\n A.η_c = -0\n A.U_s = @SVector [-0, -0]\n A.η_s = -0\n A.Δu = @SVector [-0, -0]\n A.η_diag = -0\n A.Δη = -0\n\n return nothing\nend\n\nfunction main(; restart = 0)\n mpicomm = MPI.COMM_WORLD\n\n ll = uppercase(get(ENV, \"JULIA_LOG_LEVEL\", \"INFO\"))\n loglevel =\n ll == \"DEBUG\" ? Logging.Debug :\n ll == \"WARN\" ? Logging.Warn :\n ll == \"ERROR\" ? Logging.Error : Logging.Info\n logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull\n global_logger(ConsoleLogger(logger_stream, loglevel))\n\n if restart == 0 && MPI.Comm_rank(mpicomm) == 0 && isdir(vtkpath)\n @info @sprintf(\"\"\"Remove old dir: %s and make new one\"\"\", vtkpath)\n rm(vtkpath, recursive = true)\n end\n\n brickrange_2D = (xrange, yrange)\n topl_2D =\n BrickTopology(mpicomm, brickrange_2D, periodicity = (false, false))\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (false, false, false),\n boundary = ((1, 1), (1, 1), (2, 3)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.CoastlineNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceStressForcing(),\n )\n prob = SimpleBox{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, λʳ, θᴱ, BC)\n gravity::FT = grav(param_set)\n\n #- set model time-step:\n dt_fast = 240\n dt_slow = 5400\n # dt_fast = 300\n # dt_slow = 300\n if t_chkp > 0\n n_chkp = ceil(Int64, t_chkp / dt_slow)\n dt_slow = t_chkp / n_chkp\n else\n n_chkp = ceil(Int64, runTime / dt_slow)\n dt_slow = runTime / n_chkp\n n_chkp = 0\n end\n n_outp =\n t_outp > 0 ? floor(Int64, t_outp / dt_slow) :\n ceil(Int64, runTime / dt_slow)\n ivdc_dt = numImplSteps > 0 ? dt_slow / FT(numImplSteps) : dt_slow\n\n model = OceanModel{FT}(\n prob,\n grav = gravity,\n cʰ = cʰ,\n add_fast_substeps = add_fast_substeps,\n )\n # model = OceanModel{FT}(prob, cʰ = cʰ, fₒ = FT(0), β = FT(0) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(1e3), νᶻ = FT(1e-3) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(0), fₒ = FT(0), β = FT(0) )\n\n barotropicmodel = BarotropicModel(model)\n\n minΔx = min_node_distance(grid_3D, HorizontalDirection())\n minΔz = min_node_distance(grid_3D, VerticalDirection())\n #- 2 horiz directions\n gravity_max_dT = 1 / (2 * sqrt(gravity * H) / minΔx)\n # dt_fast = minimum([gravity_max_dT])\n\n #- 2 horiz directions + harmonic visc or diffusion: 2^2 factor in CFL:\n viscous_max_dT = 1 / (2 * model.νʰ / minΔx^2 + model.νᶻ / minΔz^2) / 4\n diffusive_max_dT = 1 / (2 * model.κʰ / minΔx^2 + model.κᶻ / minΔz^2) / 4\n # dt_slow = minimum([diffusive_max_dT, viscous_max_dT])\n\n @info @sprintf(\n \"\"\"Update\n Gravity Max-dT = %.1f\n Timestep = %.1f\"\"\",\n gravity_max_dT,\n dt_fast\n )\n\n @info @sprintf(\n \"\"\"Update\n Viscous Max-dT = %.1f\n Diffusive Max-dT = %.1f\n Timestep = %.1f\"\"\",\n viscous_max_dT,\n diffusive_max_dT,\n dt_slow\n )\n\n if restart > 0\n direction = EveryDirection()\n Q_3D, _, t0 =\n read_checkpoint(vtkpath, \"baroclinic\", ArrayType, mpicomm, restart)\n Q_2D, _, _ =\n read_checkpoint(vtkpath, \"barotropic\", ArrayType, mpicomm, restart)\n\n dg = OceanDGModel(\n model,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = restart_ode_state(dg, Q_3D; init_on_cpu = true)\n Q_2D = restart_ode_state(barotropic_dg, Q_2D; init_on_cpu = true)\n\n else\n t0 = 0\n dg = OceanDGModel(\n model,\n grid_3D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = init_ode_state(dg, FT(0); init_on_cpu = true)\n Q_2D = init_ode_state(barotropic_dg, FT(0); init_on_cpu = true)\n\n end\n timeend = runTime + t0\n\n lsrk_ocean = LSRK54CarpenterKennedy(dg, Q_3D, dt = dt_slow, t0 = t0)\n lsrk_barotropic =\n LSRK54CarpenterKennedy(barotropic_dg, Q_2D, dt = dt_fast, t0 = t0)\n\n odesolver = SplitExplicitLSRK2nSolver(lsrk_ocean, lsrk_barotropic)\n\n #-- Set up State Check call back for config state arrays, called every ntFrq_SC time steps\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [\n (Q_3D, \"oce Q_3D\"),\n (dg.state_auxiliary, \"oce aux\"),\n # (dg.diffstate,\"oce diff\",),\n # (lsrk_ocean.dQ,\"oce_dQ\",),\n # (dg.modeldata.tendency_dg.state_auxiliary,\"tend Int aux\",),\n # (dg.modeldata.conti3d_Q,\"conti3d_Q\",),\n (Q_2D, \"baro Q_2D\"),\n (barotropic_dg.state_auxiliary, \"baro aux\"),\n ],\n ntFrq_SC;\n prec = 12,\n )\n # (barotropic_dg.diffstate,\"baro diff\",),\n # (lsrk_barotropic.dQ,\"baro_dQ\",)\n #--\n\n cb_ntFrq = [n_outp, n_chkp]\n outp_nb = round(Int64, restart * n_chkp / n_outp)\n step = [outp_nb, outp_nb, restart + 1]\n cbvector = make_callbacks(\n vtkpath,\n step,\n cb_ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg,\n model,\n Q_3D,\n barotropic_dg,\n barotropicmodel,\n Q_2D,\n )\n\n eng0 = norm(Q_3D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 ArrayType\n\n # slow fast state tuple\n Qvec = (slow = Q_3D, fast = Q_2D)\n # solve!(Qvec, odesolver; timeend = timeend, callbacks = cbvector)\n cbv = (cbvector..., cbcs_dg)\n solve!(Qvec, odesolver; timeend = timeend, callbacks = cbv)\n\n ## Enable the code block below to print table for use in reference value code\n ## reference value code sits in a file named $(@__FILE__)_refvals.jl. It is hand\n ## edited using code generated by block below when reference values are updated.\n regenRefVals = false\n if regenRefVals\n ## Print state statistics in format for use as reference values\n println(\n \"# SC ========== Test number \",\n 1,\n \" reference values and precision match template. =======\",\n )\n println(\"# SC ========== $(@__FILE__) test reference values ======================================\")\n ClimateMachine.StateCheck.scprintref(cbcs_dg)\n println(\"# SC ====================================================================================\")\n end\n\n ## Check results against reference if present\n checkRefVals = true\n if checkRefVals\n include(\"../refvals/simple_box_2dt_refvals.jl\")\n refDat = (refVals[1], refPrecs[1])\n checkPass = ClimateMachine.StateCheck.scdocheck(cbcs_dg, refDat)\n checkPass ? checkRep = \"Pass\" : checkRep = \"Fail\"\n @info @sprintf(\"\"\"Compare vs RefVals: %s\"\"\", checkRep)\n @test checkPass\n end\n\n return nothing\nend\n\nfunction make_callbacks(\n vtkpath,\n step,\n ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg_slow,\n model_slow,\n Q_slow,\n dg_fast,\n model_fast,\n Q_fast,\n)\n n_outp = ntFrq[1]\n n_chkp = ntFrq[2]\n mkpath(vtkpath)\n mkpath(vtkpath * \"/slow\")\n mkpath(vtkpath * \"/fast\")\n\n function do_output(span, step, model, dg, Q)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n step\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n mycomm = Q.mpicomm\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0 && MPI.Comm_size(mpicomm) > 1\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s/step%04d\", vtkpath, span, step)\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"mpirank%04d_step%04d\", i - 1, step)\n end\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., auxnames...),\n eltype(Q),\n )\n @info \"Done writing VTK: $pvtuprefix\"\n end\n\n end\n\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n cbvtk_slow =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n nothing\n end\n\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n cbvtk_fast =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_slow)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n if n_chkp > 0\n # Note: write zeros instead of Aux vars (not needed to restart); would be\n # better just to write state vars (once write_checkpoint() can handle it)\n cb_checkpoint = GenericCallbacks.EveryXSimulationSteps(n_chkp) do\n write_checkpoint(\n Q_slow,\n zero(Q_slow),\n odesolver,\n vtkpath,\n \"baroclinic\",\n mpicomm,\n step[3],\n )\n\n write_checkpoint(\n Q_fast,\n zero(Q_fast),\n odesolver,\n vtkpath,\n \"barotropic\",\n mpicomm,\n step[3],\n )\n\n step[3] += 1\n nothing\n end\n return (cbvtk_slow, cbvtk_fast, cbinfo, cb_checkpoint)\n else\n return (cbvtk_slow, cbvtk_fast, cbinfo)\n end\n\nend\n\n#################\n# RUN THE TESTS #\n#################\nFT = Float64\nvtkpath = \"vtk_split\"\n\nconst runTime = 5 * 24 * 3600 # s\nconst t_outp = 24 * 3600 # s\nconst t_chkp = runTime # s\n#const runTime = 6 * 3600 # s\n#const t_outp = 6 * 3600 # s\n#const t_chkp = 0\nconst ntFrq_SC = 1 # frequency (in time-step) for State-Check output\n\nconst N = 4\nconst Nˣ = 20\nconst Nʸ = 20\nconst Nᶻ = 20\nconst Lˣ = 4e6 # m\nconst Lʸ = 4e6 # m\nconst H = 1000 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\nzrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n\n#const cʰ = sqrt(gravity * H)\nconst cʰ = 1 # typical of ocean internal-wave speed\nconst cᶻ = 0\n\n#- inverse ratio of additional fast time steps (for weighted average)\n# --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n# e.g., = 1 --> 100% more ; = 2 --> 50% more ; = 3 --> 33% more ...\nadd_fast_substeps = 2\n\n#- number of Implicit vertical-diffusion sub-time-steps within one model full time-step\n# default = 0 : disable implicit vertical diffusion\nnumImplSteps = 0\n\nconst τₒ = 2e-1 # (Pa = N/m^2)\nconst λʳ = 20 // 86400 # m/s\n#- since we are using old BC (with factor of 2), take only half:\n#const τₒ = 1e-1\n#const λʳ = 10 // 86400\nconst θᴱ = 10 # deg.C\n\n@testset \"$(@__FILE__)\" begin\n main(restart = 0)\nend\n" }, { "path": "test/Ocean/SplitExplicit/simple_box_ivd.jl", "content": "#!/usr/bin/env julia --project\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates: flattenednames\nusing ClimateMachine.Ocean.SplitExplicit01\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.VTK\nusing ClimateMachine.Checkpoint\n\nusing Test\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport ClimateMachine.Ocean.SplitExplicit01:\n ocean_init_aux!,\n ocean_init_state!,\n ocean_boundary_state!,\n CoastlineFreeSlip,\n CoastlineNoSlip,\n OceanFloorFreeSlip,\n OceanFloorNoSlip,\n OceanSurfaceNoStressNoForcing,\n OceanSurfaceStressNoForcing,\n OceanSurfaceNoStressForcing,\n OceanSurfaceStressForcing\nimport ClimateMachine.DGMethods:\n update_auxiliary_state!, update_auxiliary_state_gradient!, VerticalDirection\n# using GPUifyLoops\n\nconst ArrayType = ClimateMachine.array_type()\n\nstruct SimpleBox{T, BC} <: AbstractOceanProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n λʳ::T\n θᴱ::T\n boundary_conditions::BC\nend\n\nfunction ocean_init_state!(p::SimpleBox, Q, A, localgeo, t)\n coords = localgeo.coord\n @inbounds y = coords[2]\n @inbounds z = coords[3]\n @inbounds H = p.H\n\n Q.u = @SVector [-0, -0]\n Q.η = -0\n Q.θ = (5 + 4 * cos(y * π / p.Lʸ)) * (1 + z / H)\n\n return nothing\nend\n\nfunction ocean_init_aux!(m::OceanModel, p::SimpleBox, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # not sure if this is needed but getting weird intialization stuff\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n A.u_d = @SVector [-0, -0]\n A.ΔGu = @SVector [-0, -0]\n\n return nothing\nend\n\n# A is Filled afer the state\nfunction ocean_init_aux!(m::BarotropicModel, P::SimpleBox, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.U_c = @SVector [-0, -0]\n A.η_c = -0\n A.U_s = @SVector [-0, -0]\n A.η_s = -0\n A.Δu = @SVector [-0, -0]\n A.η_diag = -0\n A.Δη = -0\n\n return nothing\nend\n\nfunction main(; restart = 0)\n mpicomm = MPI.COMM_WORLD\n\n ll = uppercase(get(ENV, \"JULIA_LOG_LEVEL\", \"INFO\"))\n loglevel =\n ll == \"DEBUG\" ? Logging.Debug :\n ll == \"WARN\" ? Logging.Warn :\n ll == \"ERROR\" ? Logging.Error : Logging.Info\n logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull\n global_logger(ConsoleLogger(logger_stream, loglevel))\n\n if restart == 0 && MPI.Comm_rank(mpicomm) == 0 && isdir(vtkpath)\n @info @sprintf(\"\"\"Remove old dir: %s and make new one\"\"\", vtkpath)\n rm(vtkpath, recursive = true)\n end\n\n brickrange_2D = (xrange, yrange)\n topl_2D =\n BrickTopology(mpicomm, brickrange_2D, periodicity = (false, false))\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (false, false, false),\n boundary = ((1, 1), (1, 1), (2, 3)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.CoastlineNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceStressForcing(),\n )\n prob = SimpleBox{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, λʳ, θᴱ, BC)\n gravity::FT = grav(param_set)\n\n #- set model time-step:\n dt_fast = 240\n dt_slow = 5400\n # dt_fast = 300\n # dt_slow = 300\n if t_chkp > 0\n n_chkp = ceil(Int64, t_chkp / dt_slow)\n dt_slow = t_chkp / n_chkp\n else\n n_chkp = ceil(Int64, runTime / dt_slow)\n dt_slow = runTime / n_chkp\n n_chkp = 0\n end\n n_outp =\n t_outp > 0 ? floor(Int64, t_outp / dt_slow) :\n ceil(Int64, runTime / dt_slow)\n ivdc_dt = numImplSteps > 0 ? dt_slow / FT(numImplSteps) : dt_slow\n\n model = OceanModel{FT}(\n prob,\n grav = gravity,\n cʰ = cʰ,\n add_fast_substeps = add_fast_substeps,\n numImplSteps = numImplSteps,\n ivdc_dt = ivdc_dt,\n κᶜ = FT(0.1),\n )\n # model = OceanModel{FT}(prob, cʰ = cʰ, fₒ = FT(0), β = FT(0) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(1e3), νᶻ = FT(1e-3) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(0), fₒ = FT(0), β = FT(0) )\n\n barotropicmodel = BarotropicModel(model)\n\n minΔx = min_node_distance(grid_3D, HorizontalDirection())\n minΔz = min_node_distance(grid_3D, VerticalDirection())\n #- 2 horiz directions\n gravity_max_dT = 1 / (2 * sqrt(gravity * H) / minΔx)\n # dt_fast = minimum([gravity_max_dT])\n\n #- 2 horiz directions + harmonic visc or diffusion: 2^2 factor in CFL:\n viscous_max_dT = 1 / (2 * model.νʰ / minΔx^2 + model.νᶻ / minΔz^2) / 4\n diffusive_max_dT = 1 / (2 * model.κʰ / minΔx^2 + model.κᶻ / minΔz^2) / 4\n # dt_slow = minimum([diffusive_max_dT, viscous_max_dT])\n\n @info @sprintf(\n \"\"\"Update\n Gravity Max-dT = %.1f\n Timestep = %.1f\"\"\",\n gravity_max_dT,\n dt_fast\n )\n\n @info @sprintf(\n \"\"\"Update\n Viscous Max-dT = %.1f\n Diffusive Max-dT = %.1f\n Timestep = %.1f\"\"\",\n viscous_max_dT,\n diffusive_max_dT,\n dt_slow\n )\n\n if restart > 0\n direction = EveryDirection()\n Q_3D, _, t0 =\n read_checkpoint(vtkpath, \"baroclinic\", ArrayType, mpicomm, restart)\n Q_2D, _, _ =\n read_checkpoint(vtkpath, \"barotropic\", ArrayType, mpicomm, restart)\n\n dg = OceanDGModel(\n model,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = restart_ode_state(dg, Q_3D; init_on_cpu = true)\n Q_2D = restart_ode_state(barotropic_dg, Q_2D; init_on_cpu = true)\n\n else\n t0 = 0\n dg = OceanDGModel(\n model,\n grid_3D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = init_ode_state(dg, FT(0); init_on_cpu = true)\n Q_2D = init_ode_state(barotropic_dg, FT(0); init_on_cpu = true)\n\n end\n timeend = runTime + t0\n\n lsrk_ocean = LSRK54CarpenterKennedy(dg, Q_3D, dt = dt_slow, t0 = t0)\n lsrk_barotropic =\n LSRK54CarpenterKennedy(barotropic_dg, Q_2D, dt = dt_fast, t0 = t0)\n\n odesolver = SplitExplicitLSRK2nSolver(lsrk_ocean, lsrk_barotropic)\n\n #-- Set up State Check call back for config state arrays, called every ntFrq_SC time steps\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [\n (Q_3D, \"oce Q_3D\"),\n (dg.state_auxiliary, \"oce aux\"),\n # (dg.diffstate,\"oce diff\",),\n # (lsrk_ocean.dQ,\"oce_dQ\",),\n # (dg.modeldata.tendency_dg.state_auxiliary,\"tend Int aux\",),\n # (dg.modeldata.conti3d_Q,\"conti3d_Q\",),\n (Q_2D, \"baro Q_2D\"),\n (barotropic_dg.state_auxiliary, \"baro aux\"),\n ],\n ntFrq_SC;\n prec = 12,\n )\n # (barotropic_dg.diffstate,\"baro diff\",),\n # (lsrk_barotropic.dQ,\"baro_dQ\",)\n #--\n\n cb_ntFrq = [n_outp, n_chkp]\n outp_nb = round(Int64, restart * n_chkp / n_outp)\n step = [outp_nb, outp_nb, restart + 1]\n cbvector = make_callbacks(\n vtkpath,\n step,\n cb_ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg,\n model,\n Q_3D,\n barotropic_dg,\n barotropicmodel,\n Q_2D,\n )\n\n eng0 = norm(Q_3D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 ArrayType\n\n # slow fast state tuple\n Qvec = (slow = Q_3D, fast = Q_2D)\n # solve!(Qvec, odesolver; timeend = timeend, callbacks = cbvector)\n cbv = (cbvector..., cbcs_dg)\n solve!(Qvec, odesolver; timeend = timeend, callbacks = cbv)\n\n ## Enable the code block below to print table for use in reference value code\n ## reference value code sits in a file named $(@__FILE__)_refvals.jl. It is hand\n ## edited using code generated by block below when reference values are updated.\n regenRefVals = false\n if regenRefVals\n ## Print state statistics in format for use as reference values\n println(\n \"# SC ========== Test number \",\n 1,\n \" reference values and precision match template. =======\",\n )\n println(\"# SC ========== $(@__FILE__) test reference values ======================================\")\n ClimateMachine.StateCheck.scprintref(cbcs_dg)\n println(\"# SC ====================================================================================\")\n end\n\n ## Check results against reference if present\n checkRefVals = true\n if checkRefVals\n include(\"../refvals/simple_box_ivd_refvals.jl\")\n refDat = (refVals[1], refPrecs[1])\n checkPass = ClimateMachine.StateCheck.scdocheck(cbcs_dg, refDat)\n checkPass ? checkRep = \"Pass\" : checkRep = \"Fail\"\n @info @sprintf(\"\"\"Compare vs RefVals: %s\"\"\", checkRep)\n @test checkPass\n end\n\n return nothing\nend\n\nfunction make_callbacks(\n vtkpath,\n step,\n ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg_slow,\n model_slow,\n Q_slow,\n dg_fast,\n model_fast,\n Q_fast,\n)\n n_outp = ntFrq[1]\n n_chkp = ntFrq[2]\n mkpath(vtkpath)\n mkpath(vtkpath * \"/slow\")\n mkpath(vtkpath * \"/fast\")\n\n function do_output(span, step, model, dg, Q)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n step\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n mycomm = Q.mpicomm\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0 && MPI.Comm_size(mpicomm) > 1\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s/step%04d\", vtkpath, span, step)\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"mpirank%04d_step%04d\", i - 1, step)\n end\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., auxnames...),\n eltype(Q),\n )\n @info \"Done writing VTK: $pvtuprefix\"\n end\n\n end\n\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n cbvtk_slow =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n nothing\n end\n\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n cbvtk_fast =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_slow)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n if n_chkp > 0\n # Note: write zeros instead of Aux vars (not needed to restart); would be\n # better just to write state vars (once write_checkpoint() can handle it)\n cb_checkpoint = GenericCallbacks.EveryXSimulationSteps(n_chkp) do\n write_checkpoint(\n Q_slow,\n zero(Q_slow),\n odesolver,\n vtkpath,\n \"baroclinic\",\n mpicomm,\n step[3],\n )\n\n write_checkpoint(\n Q_fast,\n zero(Q_fast),\n odesolver,\n vtkpath,\n \"barotropic\",\n mpicomm,\n step[3],\n )\n\n step[3] += 1\n nothing\n end\n return (cbvtk_slow, cbvtk_fast, cbinfo, cb_checkpoint)\n # return (cbinfo, cb_checkpoint)\n else\n return (cbvtk_slow, cbvtk_fast, cbinfo)\n # return (cbinfo)\n end\n\nend\n\n#################\n# RUN THE TESTS #\n#################\nFT = Float64\nvtkpath = \"vtk_split\"\n\nconst runTime = 5 * 24 * 3600 # s\nconst t_outp = 24 * 3600 # s\nconst t_chkp = runTime # s\n#const runTime = 6 * 3600 # s\n#const t_outp = 6 * 3600 # s\n#const t_chkp = 0\nconst ntFrq_SC = 1 # frequency (in time-step) for State-Check output\n\nconst N = 4\nconst Nˣ = 20\nconst Nʸ = 20\nconst Nᶻ = 20\nconst Lˣ = 4e6 # m\nconst Lʸ = 4e6 # m\nconst H = 1000 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\nzrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n\n#const cʰ = sqrt(gravity * H)\nconst cʰ = 1 # typical of ocean internal-wave speed\nconst cᶻ = 0\n\n#- inverse ratio of additional fast time steps (for weighted average)\n# --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n# e.g., = 1 --> 100% more ; = 2 --> 50% more ; = 3 --> 33% more ...\nadd_fast_substeps = 2\n\n#- number of Implicit vertical-diffusion sub-time-steps within one model full time-step\n# default = 0 : disable implicit vertical diffusion\nnumImplSteps = 5\n\nconst τₒ = 2e-1 # (Pa = N/m^2)\nconst λʳ = 20 // 86400 # m/s\n#- since we are using old BC (with factor of 2), take only half:\n#const τₒ = 1e-1\n#const λʳ = 10 // 86400\nconst θᴱ = 10 # deg.C\n\n@testset \"$(@__FILE__)\" begin\n main(restart = 0)\nend\n" }, { "path": "test/Ocean/SplitExplicit/simple_box_rk3.jl", "content": "#!/usr/bin/env julia --project\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates: flattenednames\nusing ClimateMachine.Ocean.SplitExplicit01\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.VTK\nusing ClimateMachine.Checkpoint\n\nusing Test\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport ClimateMachine.Ocean.SplitExplicit01:\n ocean_init_aux!,\n ocean_init_state!,\n ocean_boundary_state!,\n CoastlineFreeSlip,\n CoastlineNoSlip,\n OceanFloorFreeSlip,\n OceanFloorNoSlip,\n OceanSurfaceNoStressNoForcing,\n OceanSurfaceStressNoForcing,\n OceanSurfaceNoStressForcing,\n OceanSurfaceStressForcing\nimport ClimateMachine.DGMethods:\n update_auxiliary_state!, update_auxiliary_state_gradient!, VerticalDirection\n# using GPUifyLoops\n\nconst ArrayType = ClimateMachine.array_type()\n\nstruct SimpleBox{T, BC} <: AbstractOceanProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n λʳ::T\n θᴱ::T\n boundary_conditions::BC\nend\n\nfunction ocean_init_state!(p::SimpleBox, Q, A, localgeo, t)\n coords = localgeo.coord\n @inbounds y = coords[2]\n @inbounds z = coords[3]\n @inbounds H = p.H\n\n Q.u = @SVector [-0, -0]\n Q.η = -0\n Q.θ = (5 + 4 * cos(y * π / p.Lʸ)) * (1 + z / H)\n\n return nothing\nend\n\nfunction ocean_init_aux!(m::OceanModel, p::SimpleBox, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # not sure if this is needed but getting weird intialization stuff\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n A.u_d = @SVector [-0, -0]\n A.ΔGu = @SVector [-0, -0]\n\n return nothing\nend\n\n# A is Filled afer the state\nfunction ocean_init_aux!(m::BarotropicModel, P::SimpleBox, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.U_c = @SVector [-0, -0]\n A.η_c = -0\n A.U_s = @SVector [-0, -0]\n A.η_s = -0\n A.Δu = @SVector [-0, -0]\n A.η_diag = -0\n A.Δη = -0\n\n return nothing\nend\n\nfunction main(; restart = 0)\n mpicomm = MPI.COMM_WORLD\n\n ll = uppercase(get(ENV, \"JULIA_LOG_LEVEL\", \"INFO\"))\n loglevel =\n ll == \"DEBUG\" ? Logging.Debug :\n ll == \"WARN\" ? Logging.Warn :\n ll == \"ERROR\" ? Logging.Error : Logging.Info\n logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull\n global_logger(ConsoleLogger(logger_stream, loglevel))\n\n if restart == 0 && MPI.Comm_rank(mpicomm) == 0 && isdir(vtkpath)\n @info @sprintf(\"\"\"Remove old dir: %s and make new one\"\"\", vtkpath)\n rm(vtkpath, recursive = true)\n end\n\n brickrange_2D = (xrange, yrange)\n topl_2D =\n BrickTopology(mpicomm, brickrange_2D, periodicity = (false, false))\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (false, false, false),\n boundary = ((1, 1), (1, 1), (2, 3)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.CoastlineNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceStressForcing(),\n )\n prob = SimpleBox{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, λʳ, θᴱ, BC)\n gravity::FT = grav(param_set)\n\n #- set model time-step:\n dt_fast = 120\n dt_slow = 2400\n # dt_fast = 240\n # dt_slow = 5400\n if t_chkp > 0\n n_chkp = ceil(Int64, t_chkp / dt_slow)\n dt_slow = t_chkp / n_chkp\n else\n n_chkp = ceil(Int64, runTime / dt_slow)\n dt_slow = runTime / n_chkp\n n_chkp = 0\n end\n n_outp =\n t_outp > 0 ? floor(Int64, t_outp / dt_slow) :\n ceil(Int64, runTime / dt_slow)\n ivdc_dt = numImplSteps > 0 ? dt_slow / FT(numImplSteps) : dt_slow\n\n model = OceanModel{FT}(\n prob,\n grav = gravity,\n cʰ = cʰ,\n add_fast_substeps = add_fast_substeps,\n numImplSteps = numImplSteps,\n ivdc_dt = ivdc_dt,\n κᶜ = FT(0.1),\n )\n # model = OceanModel{FT}(prob, cʰ = cʰ, fₒ = FT(0), β = FT(0) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(1e3), νᶻ = FT(1e-3) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(0), fₒ = FT(0), β = FT(0) )\n\n barotropicmodel = BarotropicModel(model)\n\n minΔx = min_node_distance(grid_3D, HorizontalDirection())\n minΔz = min_node_distance(grid_3D, VerticalDirection())\n #- 2 horiz directions\n gravity_max_dT = 1 / (2 * sqrt(gravity * H) / minΔx)\n # dt_fast = minimum([gravity_max_dT])\n\n #- 2 horiz directions + harmonic visc or diffusion: 2^2 factor in CFL:\n viscous_max_dT = 1 / (2 * model.νʰ / minΔx^2 + model.νᶻ / minΔz^2) / 4\n diffusive_max_dT = 1 / (2 * model.κʰ / minΔx^2 + model.κᶻ / minΔz^2) / 4\n # dt_slow = minimum([diffusive_max_dT, viscous_max_dT])\n\n @info @sprintf(\n \"\"\"Update\n Gravity Max-dT = %.1f\n Timestep = %.1f\"\"\",\n gravity_max_dT,\n dt_fast\n )\n\n @info @sprintf(\n \"\"\"Update\n Viscous Max-dT = %.1f\n Diffusive Max-dT = %.1f\n Timestep = %.1f\"\"\",\n viscous_max_dT,\n diffusive_max_dT,\n dt_slow\n )\n\n if restart > 0\n direction = EveryDirection()\n Q_3D, _, t0 =\n read_checkpoint(vtkpath, \"baroclinic\", ArrayType, mpicomm, restart)\n Q_2D, _, _ =\n read_checkpoint(vtkpath, \"barotropic\", ArrayType, mpicomm, restart)\n\n dg = OceanDGModel(\n model,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = restart_ode_state(dg, Q_3D; init_on_cpu = true)\n Q_2D = restart_ode_state(barotropic_dg, Q_2D; init_on_cpu = true)\n\n else\n t0 = 0\n dg = OceanDGModel(\n model,\n grid_3D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = init_ode_state(dg, FT(0); init_on_cpu = true)\n Q_2D = init_ode_state(barotropic_dg, FT(0); init_on_cpu = true)\n\n end\n timeend = runTime + t0\n\n lsrk_ocean = LS3NRK33Heuns(dg, Q_3D, dt = dt_slow, t0 = t0)\n lsrk_barotropic = LS3NRK33Heuns(barotropic_dg, Q_2D, dt = dt_fast, t0 = t0)\n\n odesolver = SplitExplicitLSRK3nSolver(lsrk_ocean, lsrk_barotropic)\n\n #-- Set up State Check call back for config state arrays, called every ntFrq_SC time steps\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [\n (Q_3D, \"oce Q_3D\"),\n (dg.state_auxiliary, \"oce aux\"),\n # (dg.diffstate,\"oce diff\",),\n # (lsrk_ocean.dQ,\"oce_dQ\",),\n # (dg.modeldata.tendency_dg.state_auxiliary,\"tend Int aux\",),\n # (dg.modeldata.conti3d_Q,\"conti3d_Q\",),\n (Q_2D, \"baro Q_2D\"),\n (barotropic_dg.state_auxiliary, \"baro aux\"),\n ],\n ntFrq_SC;\n prec = 12,\n )\n # (barotropic_dg.diffstate,\"baro diff\",),\n # (lsrk_barotropic.dQ,\"baro_dQ\",)\n #--\n\n cb_ntFrq = [n_outp, n_chkp]\n outp_nb = round(Int64, restart * n_chkp / n_outp)\n step = [outp_nb, outp_nb, restart + 1]\n cbvector = make_callbacks(\n vtkpath,\n step,\n cb_ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg,\n model,\n Q_3D,\n barotropic_dg,\n barotropicmodel,\n Q_2D,\n )\n\n eng0 = norm(Q_3D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 ArrayType\n\n # slow fast state tuple\n Qvec = (slow = Q_3D, fast = Q_2D)\n # solve!(Qvec, odesolver; timeend = timeend, callbacks = cbvector)\n cbv = (cbvector..., cbcs_dg)\n solve!(Qvec, odesolver; timeend = timeend, callbacks = cbv)\n\n ## Enable the code block below to print table for use in reference value code\n ## reference value code sits in a file named $(@__FILE__)_refvals.jl. It is hand\n ## edited using code generated by block below when reference values are updated.\n regenRefVals = false\n if regenRefVals\n ## Print state statistics in format for use as reference values\n println(\n \"# SC ========== Test number \",\n 1,\n \" reference values and precision match template. =======\",\n )\n println(\"# SC ========== $(@__FILE__) test reference values ======================================\")\n ClimateMachine.StateCheck.scprintref(cbcs_dg)\n println(\"# SC ====================================================================================\")\n end\n\n ## Check results against reference if present\n checkRefVals = true\n if checkRefVals\n include(\"../refvals/simple_box_rk3_refvals.jl\")\n refDat = (refVals[1], refPrecs[1])\n checkPass = ClimateMachine.StateCheck.scdocheck(cbcs_dg, refDat)\n checkPass ? checkRep = \"Pass\" : checkRep = \"Fail\"\n @info @sprintf(\"\"\"Compare vs RefVals: %s\"\"\", checkRep)\n @test checkPass\n end\n\n return nothing\nend\n\nfunction make_callbacks(\n vtkpath,\n step,\n ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg_slow,\n model_slow,\n Q_slow,\n dg_fast,\n model_fast,\n Q_fast,\n)\n n_outp = ntFrq[1]\n n_chkp = ntFrq[2]\n mkpath(vtkpath)\n mkpath(vtkpath * \"/slow\")\n mkpath(vtkpath * \"/fast\")\n\n function do_output(span, step, model, dg, Q)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n step\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n mycomm = Q.mpicomm\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0 && MPI.Comm_size(mpicomm) > 1\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s/step%04d\", vtkpath, span, step)\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"mpirank%04d_step%04d\", i - 1, step)\n end\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., auxnames...),\n eltype(Q),\n )\n @info \"Done writing VTK: $pvtuprefix\"\n end\n\n end\n\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n cbvtk_slow =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n nothing\n end\n\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n cbvtk_fast =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_slow)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n if n_chkp > 0\n # Note: write zeros instead of Aux vars (not needed to restart); would be\n # better just to write state vars (once write_checkpoint() can handle it)\n cb_checkpoint = GenericCallbacks.EveryXSimulationSteps(n_chkp) do\n write_checkpoint(\n Q_slow,\n zero(Q_slow),\n odesolver,\n vtkpath,\n \"baroclinic\",\n mpicomm,\n step[3],\n )\n\n write_checkpoint(\n Q_fast,\n zero(Q_fast),\n odesolver,\n vtkpath,\n \"barotropic\",\n mpicomm,\n step[3],\n )\n\n step[3] += 1\n nothing\n end\n return (cbvtk_slow, cbvtk_fast, cbinfo, cb_checkpoint)\n else\n return (cbvtk_slow, cbvtk_fast, cbinfo)\n end\n\nend\n\n#################\n# RUN THE TESTS #\n#################\nFT = Float64\nvtkpath = \"vtk_split\"\n\nconst runTime = 3 * 24 * 3600 # s\nconst t_outp = 24 * 3600 # s\nconst t_chkp = runTime # s\n#const runTime = 6 * 3600 # s\n#const t_outp = 6 * 3600 # s\n#const t_chkp = 0\nconst ntFrq_SC = 1 # frequency (in time-step) for State-Check output\n\nconst N = 4\nconst Nˣ = 20\nconst Nʸ = 20\nconst Nᶻ = 20\nconst Lˣ = 4e6 # m\nconst Lʸ = 4e6 # m\nconst H = 1000 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\nzrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n\n#const cʰ = sqrt(gravity * H)\nconst cʰ = 1 # typical of ocean internal-wave speed\nconst cᶻ = 0\n\n#- inverse ratio of additional fast time steps (for weighted average)\n# --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n# e.g., = 1 --> 100% more ; = 2 --> 50% more ; = 3 --> 33% more ...\nadd_fast_substeps = 3\n\n#- number of Implicit vertical-diffusion sub-time-steps within one model full time-step\n# default = 0 : disable implicit vertical diffusion\nnumImplSteps = 5\n\nconst τₒ = 2e-1 # (Pa = N/m^2)\nconst λʳ = 20 // 86400 # m/s\n#- since we are using old BC (with factor of 2), take only half:\n#const τₒ = 1e-1\n#const λʳ = 10 // 86400\nconst θᴱ = 10 # deg.C\n\n@testset \"$(@__FILE__)\" begin\n main(restart = 0)\nend\n" }, { "path": "test/Ocean/SplitExplicit/simple_dbl_gyre.jl", "content": "#!/usr/bin/env julia --project\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates: flattenednames\nusing ClimateMachine.Ocean.SplitExplicit01\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.VTK\nusing ClimateMachine.Checkpoint\n\nusing Test\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nimport ClimateMachine.Ocean.SplitExplicit01:\n ocean_init_aux!,\n ocean_init_state!,\n ocean_boundary_state!,\n CoastlineFreeSlip,\n CoastlineNoSlip,\n OceanFloorFreeSlip,\n OceanFloorNoSlip,\n OceanSurfaceNoStressNoForcing,\n OceanSurfaceStressNoForcing,\n OceanSurfaceNoStressForcing,\n OceanSurfaceStressForcing,\n velocity_flux,\n temperature_flux\nimport ClimateMachine.DGMethods:\n update_auxiliary_state!, update_auxiliary_state_gradient!, VerticalDirection\n# using GPUifyLoops\n\nconst ArrayType = ClimateMachine.array_type()\n\nstruct DoubleGyreBox{T, BC} <: AbstractOceanProblem\n Lˣ::T\n Lʸ::T\n H::T\n τₒ::T\n λʳ::T\n θᴱ::T\n boundary_conditions::BC\nend\n\n@inline velocity_flux(p::DoubleGyreBox, y, ρ) =\n -(p.τₒ / ρ) * cos(2 * π * y / p.Lʸ)\n\n@inline function temperature_flux(p::DoubleGyreBox, y, θ)\n θʳ = p.θᴱ * (1 - y / p.Lʸ)\n return p.λʳ * (θʳ - θ)\nend\n\nfunction ocean_init_state!(p::DoubleGyreBox, Q, A, localgeo, t)\n coords = localgeo.coord\n @inbounds y = coords[2]\n @inbounds z = coords[3]\n @inbounds H = p.H\n\n Q.u = @SVector [-0, -0]\n Q.η = -0\n Q.θ = (12 + 10 * cos(π * y / p.Lʸ)) * (1 + z / H)\n\n return nothing\nend\n\nfunction ocean_init_aux!(m::OceanModel, p::DoubleGyreBox, A, geom)\n FT = eltype(A)\n @inbounds A.y = geom.coord[2]\n\n # not sure if this is needed but getting weird intialization stuff\n A.w = -0\n A.pkin = -0\n A.wz0 = -0\n A.u_d = @SVector [-0, -0]\n A.ΔGu = @SVector [-0, -0]\n\n return nothing\nend\n\n# A is Filled afer the state\nfunction ocean_init_aux!(m::BarotropicModel, P::DoubleGyreBox, A, geom)\n @inbounds A.y = geom.coord[2]\n\n A.Gᵁ = @SVector [-0, -0]\n A.U_c = @SVector [-0, -0]\n A.η_c = -0\n A.U_s = @SVector [-0, -0]\n A.η_s = -0\n A.Δu = @SVector [-0, -0]\n A.η_diag = -0\n A.Δη = -0\n\n return nothing\nend\n\nfunction main(; restart = 0)\n mpicomm = MPI.COMM_WORLD\n\n ll = uppercase(get(ENV, \"JULIA_LOG_LEVEL\", \"INFO\"))\n loglevel =\n ll == \"DEBUG\" ? Logging.Debug :\n ll == \"WARN\" ? Logging.Warn :\n ll == \"ERROR\" ? Logging.Error : Logging.Info\n logger_stream = MPI.Comm_rank(mpicomm) == 0 ? stderr : devnull\n global_logger(ConsoleLogger(logger_stream, loglevel))\n\n if restart == 0 && MPI.Comm_rank(mpicomm) == 0 && isdir(vtkpath)\n @info @sprintf(\"\"\"Remove old dir: %s and make new one\"\"\", vtkpath)\n rm(vtkpath, recursive = true)\n end\n\n brickrange_2D = (xrange, yrange)\n topl_2D =\n BrickTopology(mpicomm, brickrange_2D, periodicity = (false, false))\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (false, false, false),\n boundary = ((1, 1), (1, 1), (2, 3)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.CoastlineNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceStressForcing(),\n )\n prob = DoubleGyreBox{FT, typeof(BC)}(Lˣ, Lʸ, H, τₒ, λʳ, θᴱ, BC)\n gravity::FT = grav(param_set)\n\n #- set model time-step:\n dt_fast = 96\n dt_slow = 3456\n if t_chkp > 0\n n_chkp = ceil(Int64, t_chkp / dt_slow)\n dt_slow = t_chkp / n_chkp\n else\n n_chkp = ceil(Int64, runTime / dt_slow)\n dt_slow = runTime / n_chkp\n n_chkp = 0\n end\n n_outp =\n t_outp > 0 ? floor(Int64, t_outp / dt_slow) :\n ceil(Int64, runTime / dt_slow)\n ivdc_dt = numImplSteps > 0 ? dt_slow / FT(numImplSteps) : dt_slow\n\n model = OceanModel{FT}(\n prob,\n grav = gravity,\n cʰ = cʰ,\n add_fast_substeps = add_fast_substeps,\n numImplSteps = numImplSteps,\n ivdc_dt = ivdc_dt,\n νʰ = FT(15e3),\n νᶻ = FT(5e-3),\n κᶜ = FT(1.0),\n fₒ = FT(3.8e-5),\n β = FT(1.7e-11),\n )\n # model = OceanModel{FT}(prob, cʰ = cʰ, fₒ = FT(0), β = FT(0) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(1e3), νᶻ = FT(1e-3) )\n # model = OceanModel{FT}(prob, cʰ = cʰ, νʰ = FT(0), fₒ = FT(0), β = FT(0) )\n\n barotropicmodel = BarotropicModel(model)\n\n minΔx = min_node_distance(grid_3D, HorizontalDirection())\n minΔz = min_node_distance(grid_3D, VerticalDirection())\n #- 2 horiz directions\n gravity_max_dT = 1 / (2 * sqrt(gravity * H) / minΔx)\n # dt_fast = minimum([gravity_max_dT])\n\n #- 2 horiz directions + harmonic visc or diffusion: 2^2 factor in CFL:\n viscous_max_dT = 1 / (2 * model.νʰ / minΔx^2 + model.νᶻ / minΔz^2) / 4\n diffusive_max_dT = 1 / (2 * model.κʰ / minΔx^2 + model.κᶻ / minΔz^2) / 4\n # dt_slow = minimum([diffusive_max_dT, viscous_max_dT])\n\n @info @sprintf(\n \"\"\"Update\n Gravity Max-dT = %.1f\n Timestep = %.1f\"\"\",\n gravity_max_dT,\n dt_fast\n )\n\n @info @sprintf(\n \"\"\"Update\n Viscous Max-dT = %.1f\n Diffusive Max-dT = %.1f\n Timestep = %.1f\"\"\",\n viscous_max_dT,\n diffusive_max_dT,\n dt_slow\n )\n\n if restart > 0\n direction = EveryDirection()\n Q_3D, _, t0 =\n read_checkpoint(vtkpath, \"baroclinic\", ArrayType, mpicomm, restart)\n Q_2D, _, _ =\n read_checkpoint(vtkpath, \"barotropic\", ArrayType, mpicomm, restart)\n\n dg = OceanDGModel(\n model,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = restart_ode_state(dg, Q_3D; init_on_cpu = true)\n Q_2D = restart_ode_state(barotropic_dg, Q_2D; init_on_cpu = true)\n\n else\n t0 = 0\n dg = OceanDGModel(\n model,\n grid_3D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n barotropic_dg = DGModel(\n barotropicmodel,\n grid_2D,\n # CentralNumericalFluxFirstOrder(),\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = init_ode_state(dg, FT(0); init_on_cpu = true)\n Q_2D = init_ode_state(barotropic_dg, FT(0); init_on_cpu = true)\n\n end\n timeend = runTime + t0\n\n lsrk_ocean = LS3NRK33Heuns(dg, Q_3D, dt = dt_slow, t0 = t0)\n lsrk_barotropic = LS3NRK33Heuns(barotropic_dg, Q_2D, dt = dt_fast, t0 = t0)\n\n odesolver = SplitExplicitLSRK3nSolver(lsrk_ocean, lsrk_barotropic)\n\n #-- Set up State Check call back for config state arrays, called every ntFrq_SC time steps\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [\n (Q_3D, \"oce Q_3D\"),\n (dg.state_auxiliary, \"oce aux\"),\n # (dg.diffstate,\"oce diff\",),\n # (lsrk_ocean.dQ,\"oce_dQ\",),\n # (dg.modeldata.tendency_dg.state_auxiliary,\"tend Int aux\",),\n # (dg.modeldata.conti3d_Q,\"conti3d_Q\",),\n (Q_2D, \"baro Q_2D\"),\n (barotropic_dg.state_auxiliary, \"baro aux\"),\n ],\n ntFrq_SC;\n prec = 12,\n )\n # (barotropic_dg.diffstate,\"baro diff\",),\n # (lsrk_barotropic.dQ,\"baro_dQ\",)\n #--\n\n cb_ntFrq = [n_outp, n_chkp]\n outp_nb = round(Int64, restart * n_chkp / n_outp)\n step = [outp_nb, outp_nb, restart + 1]\n cbvector = make_callbacks(\n vtkpath,\n step,\n cb_ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg,\n model,\n Q_3D,\n barotropic_dg,\n barotropicmodel,\n Q_2D,\n )\n\n eng0 = norm(Q_3D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 ArrayType\n\n # slow fast state tuple\n Qvec = (slow = Q_3D, fast = Q_2D)\n # solve!(Qvec, odesolver; timeend = timeend, callbacks = cbvector)\n cbv = (cbvector..., cbcs_dg)\n solve!(Qvec, odesolver; timeend = timeend, callbacks = cbv)\n\n ## Enable the code block below to print table for use in reference value code\n ## reference value code sits in a file named $(@__FILE__)_refvals.jl. It is hand\n ## edited using code generated by block below when reference values are updated.\n regenRefVals = false\n if regenRefVals\n ## Print state statistics in format for use as reference values\n println(\n \"# SC ========== Test number \",\n 1,\n \" reference values and precision match template. =======\",\n )\n println(\"# SC ========== $(@__FILE__) test reference values ======================================\")\n ClimateMachine.StateCheck.scprintref(cbcs_dg)\n println(\"# SC ====================================================================================\")\n end\n\n ## Check results against reference if present\n checkRefVals = true\n if checkRefVals\n include(\"../refvals/simple_dbl_gyre_refvals.jl\")\n refDat = (refVals[1], refPrecs[1])\n checkPass = ClimateMachine.StateCheck.scdocheck(cbcs_dg, refDat)\n checkPass ? checkRep = \"Pass\" : checkRep = \"Fail\"\n @info @sprintf(\"\"\"Compare vs RefVals: %s\"\"\", checkRep)\n @test checkPass\n end\n\n return nothing\nend\n\nfunction make_callbacks(\n vtkpath,\n step,\n ntFrq,\n timeend,\n mpicomm,\n odesolver,\n dg_slow,\n model_slow,\n Q_slow,\n dg_fast,\n model_fast,\n Q_fast,\n)\n n_outp = ntFrq[1]\n n_chkp = ntFrq[2]\n mkpath(vtkpath)\n mkpath(vtkpath * \"/slow\")\n mkpath(vtkpath * \"/fast\")\n\n function do_output(span, step, model, dg, Q)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n step\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, dg.state_auxiliary, auxnames)\n\n mycomm = Q.mpicomm\n ## Generate the pvtu file for these vtk files\n if MPI.Comm_rank(mpicomm) == 0 && MPI.Comm_size(mpicomm) > 1\n ## name of the pvtu file\n pvtuprefix = @sprintf(\"%s/%s/step%04d\", vtkpath, span, step)\n ## name of each of the ranks vtk files\n prefixes = ntuple(MPI.Comm_size(mpicomm)) do i\n @sprintf(\"mpirank%04d_step%04d\", i - 1, step)\n end\n writepvtu(\n pvtuprefix,\n prefixes,\n (statenames..., auxnames...),\n eltype(Q),\n )\n @info \"Done writing VTK: $pvtuprefix\"\n end\n\n end\n\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n cbvtk_slow =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"slow\", step[1], model_slow, dg_slow, Q_slow)\n step[1] += 1\n nothing\n end\n\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n cbvtk_fast =\n GenericCallbacks.EveryXSimulationSteps(n_outp) do (init = false)\n do_output(\"fast\", step[2], model_fast, dg_fast, Q_fast)\n step[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_slow)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n if n_chkp > 0\n # Note: write zeros instead of Aux vars (not needed to restart); would be\n # better just to write state vars (once write_checkpoint() can handle it)\n cb_checkpoint = GenericCallbacks.EveryXSimulationSteps(n_chkp) do\n write_checkpoint(\n Q_slow,\n zero(Q_slow),\n odesolver,\n vtkpath,\n \"baroclinic\",\n mpicomm,\n step[3],\n )\n\n write_checkpoint(\n Q_fast,\n zero(Q_fast),\n odesolver,\n vtkpath,\n \"barotropic\",\n mpicomm,\n step[3],\n )\n\n step[3] += 1\n nothing\n end\n return (cbvtk_slow, cbvtk_fast, cbinfo, cb_checkpoint)\n else\n return (cbvtk_slow, cbvtk_fast, cbinfo)\n end\n\nend\n\n#################\n# RUN THE TESTS #\n#################\nFT = Float64\nvtkpath = \"vtk_split\"\n\nconst runTime = 3 * 24 * 3600 # s\nconst t_outp = 24 * 3600 # s\nconst t_chkp = runTime # s\n#const runTime = 6 * 3600 # s\n#const t_outp = 6 * 3600 # s\n#const t_chkp = 0\nconst ntFrq_SC = 1 # frequency (in time-step) for State-Check output\n\nconst N = 4\nconst Nˣ = 20\nconst Nʸ = 30\nconst Nᶻ = 15\nconst Lˣ = 4e6 # m\nconst Lʸ = 6e6 # m\nconst H = 3000 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\nzrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n# dz = [576, 540, 433, 339, 266, 208, 162, 128, 99, 75, 58, 43, 31, 22, 20]\n# zrange = zeros(FT, Nᶻ + 1)\n# zrange[2:(Nᶻ + 1)] .= cumsum(dz)\n# zrange .-= H\n\n#const cʰ = sqrt(gravity * H)\nconst cʰ = 1 # typical of ocean internal-wave speed\nconst cᶻ = 0\n\n#- inverse ratio of additional fast time steps (for weighted average)\n# --> do 1/add more time-steps and average from: 1 - 1/add up to: 1 + 1/add\n# e.g., = 1 --> 100% more ; = 2 --> 50% more ; = 3 --> 33% more ...\nadd_fast_substeps = 3\n\n#- number of Implicit vertical-diffusion sub-time-steps within one model full time-step\n# default = 0 : disable implicit vertical diffusion\nnumImplSteps = 5\n\nconst τₒ = 1e-1 # (Pa = N/m^2)\nconst λʳ = 20 // 86400 # m/s\n#- since we are using old BC (with factor of 2), take only half:\n#const τₒ = 5e-2\n#const λʳ = 10 // 86400\nconst θᴱ = 25 # deg.C\n\n@testset \"$(@__FILE__)\" begin\n main(restart = 0)\nend\n" }, { "path": "test/Ocean/SplitExplicit/split_explicit.jl", "content": "#!/usr/bin/env julia --project\nusing Test\nusing ClimateMachine\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Grids: polynomialorders\nusing ClimateMachine.Ocean\nusing ClimateMachine.Ocean.HydrostaticBoussinesq\nusing ClimateMachine.Ocean.ShallowWater\nusing ClimateMachine.Ocean.SplitExplicit: VerticalIntegralModel\nusing ClimateMachine.Ocean.SplitExplicit01\nusing ClimateMachine.Ocean.OceanProblems\n\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.BalanceLaws: vars_state, Prognostic, Auxiliary\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VTK\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.SystemSolvers\n\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nstruct SplitConfig{N, D3, D2, S, M, AT}\n name::N\n dg_3D::D3\n dg_2D::D2\n solver::S\n mpicomm::M\n ArrayType::AT\nend\n\nfunction run_split_explicit(\n config::SplitConfig,\n timespan;\n dt_fast = 300,\n dt_slow = 300,\n refDat = (),\n restart = 0,\n analytic_solution = false,\n)\n Q_3D, Q_2D, t0 = init_states(config, Val(restart))\n\n tout, timeend = timespan\n\n # @show dt_fast = floor(Int, 1 / (2 * sqrt(gravity * H) / minΔx)) # / 4\n # @show dt_slow = floor(Int, minΔx / 15) # / 4\n\n nout = ceil(Int64, tout / dt_slow)\n dt_slow = tout / nout\n\n timeendlocal = timeend + t0\n\n lsrk_3D = LSRK54CarpenterKennedy(config.dg_3D, Q_3D, dt = dt_slow, t0 = t0)\n lsrk_2D = LSRK54CarpenterKennedy(config.dg_2D, Q_2D, dt = dt_fast, t0 = t0)\n\n odesolver = config.solver(lsrk_3D, lsrk_2D;)\n\n vtkstep = [restart, restart, restart + 1, restart + 1]\n cbvector = make_callbacks(\n abspath(joinpath(ClimateMachine.Settings.output_dir, config.name)),\n vtkstep,\n nout,\n config.mpicomm,\n odesolver,\n config.dg_3D,\n config.dg_3D.balance_law,\n Q_3D,\n config.dg_2D,\n config.dg_2D.balance_law,\n Q_2D,\n timeendlocal,\n )\n\n eng0 = norm(Q_3D)\n @info @sprintf \"\"\"Starting\n norm(Q₀) = %.16e\n ArrayType = %s\"\"\" eng0 config.ArrayType\n\n # slow fast state tuple\n Qvec = (slow = Q_3D, fast = Q_2D)\n solve!(Q_3D, odesolver; timeend = timeendlocal, callbacks = cbvector)\n\n if analytic_solution\n Qe_3D = init_ode_state(config.dg_3D, timeendlocal, init_on_cpu = true)\n Qe_2D = init_ode_state(config.dg_2D, timeendlocal, init_on_cpu = true)\n\n error_3D = euclidean_distance(Q_3D, Qe_3D) / norm(Qe_3D)\n error_2D = euclidean_distance(Q_2D, Qe_2D) / norm(Qe_2D)\n\n println(\"3D error = \", error_3D)\n println(\"2D error = \", error_2D)\n @test isapprox(error_3D, FT(0.0); atol = 0.005)\n @test isapprox(error_2D, FT(0.0); atol = 0.005)\n end\n\n ## Check results against reference\n ClimateMachine.StateCheck.scprintref(cbvector[end])\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(cbvector[end], refDat)\n end\n\n return nothing\nend\n\nfunction init_states(config, ::Val{0})\n Q_3D = init_ode_state(config.dg_3D, FT(0); init_on_cpu = true)\n Q_2D = config.dg_3D.modeldata.Q_2D\n\n return Q_3D, Q_2D, 0\nend\n\nfunction init_states(config, ::Val{restart}) where {restart}\n Q_3D, A_3D, t0 = read_checkpoint(\n abspath(joinpath(ClimateMachine.Settings.output_dir, config.name)),\n \"baroclinic\",\n config.ArrayType,\n config.mpicomm,\n restart,\n )\n Q_2D_restart, A_2D, _ = read_checkpoint(\n abspath(joinpath(ClimateMachine.Settings.output_dir, config.name)),\n \"barotropic\",\n config.ArrayType,\n config.mpicomm,\n restart,\n )\n\n direction = EveryDirection()\n A_3D = restart_auxiliary_state(\n config.dg_3D.balance_law,\n config.dg_3D.grid,\n A_3D,\n direction,\n )\n A_2D = restart_auxiliary_state(\n config.dg_2D.balance_law,\n config.dg_2D.grid,\n A_2D,\n direction,\n )\n\n config.dg_3D.state_auxiliary .= A_3D\n config.dg_2D.state_auxiliary .= A_2D\n\n Q_3D = restart_ode_state(config.dg_3D, Q_3D; init_on_cpu = true)\n Q_2D = config.dg_3D.modeldata.Q_2D\n Q_2D .= Q_2D_restart\n\n return Q_3D, Q_2D, t0\nend\n\nfunction make_callbacks(\n vtkpath,\n vtkstep,\n nout,\n mpicomm,\n odesolver,\n dg_slow,\n model_slow,\n Q_slow,\n dg_fast,\n model_fast,\n Q_fast,\n timeend,\n)\n mkpath(vtkpath)\n mkpath(vtkpath * \"/slow\")\n mkpath(vtkpath * \"/fast\")\n\n A_slow = dg_slow.state_auxiliary\n A_fast = dg_fast.state_auxiliary\n\n function do_output(span, vtkstep, model, dg, Q, A)\n outprefix = @sprintf(\n \"%s/%s/mpirank%04d_step%04d\",\n vtkpath,\n span,\n MPI.Comm_rank(mpicomm),\n vtkstep\n )\n @info \"doing VTK output\" outprefix\n statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))\n auxnames = flattenednames(vars_state(model, Auxiliary(), eltype(Q)))\n writevtk(outprefix, Q, dg, statenames, A, auxnames)\n end\n\n do_output(\"slow\", vtkstep[1], model_slow, dg_slow, Q_slow, A_slow)\n vtkstep[1] += 1\n cbvtk_slow = GenericCallbacks.EveryXSimulationSteps(nout) do (init = false)\n do_output(\"slow\", vtkstep[1], model_slow, dg_slow, Q_slow, A_slow)\n vtkstep[1] += 1\n nothing\n end\n\n do_output(\"fast\", vtkstep[2], model_fast, dg_fast, Q_fast, A_fast)\n vtkstep[2] += 1\n cbvtk_fast = GenericCallbacks.EveryXSimulationSteps(nout) do (init = false)\n do_output(\"fast\", vtkstep[2], model_fast, dg_fast, Q_fast, A_fast)\n vtkstep[2] += 1\n nothing\n end\n\n starttime = Ref(now())\n cbinfo = GenericCallbacks.EveryXWallTimeSeconds(60, mpicomm) do (s = false)\n if s\n starttime[] = now()\n else\n energy = norm(Q_slow)\n @info @sprintf(\n \"\"\"Update\n simtime = %8.2f / %8.2f\n runtime = %s\n norm(Q) = %.16e\"\"\",\n ODESolvers.gettime(odesolver),\n timeend,\n Dates.format(\n convert(Dates.DateTime, Dates.now() - starttime[]),\n Dates.dateformat\"HH:MM:SS\",\n ),\n energy\n )\n end\n end\n\n cbcs_dg = ClimateMachine.StateCheck.sccreate(\n [\n (Q_slow, \"3D state\"),\n (A_slow, \"3D aux\"),\n (Q_fast, \"2D state\"),\n (A_fast, \"2D aux\"),\n ],\n nout;\n prec = 12,\n )\n\n cb_checkpoint = GenericCallbacks.EveryXSimulationSteps(nout) do\n write_checkpoint(\n Q_slow,\n A_slow,\n odesolver,\n vtkpath,\n \"baroclinic\",\n mpicomm,\n vtkstep[3],\n )\n\n write_checkpoint(\n Q_fast,\n A_fast,\n odesolver,\n vtkpath,\n \"barotropic\",\n mpicomm,\n vtkstep[4],\n )\n\n rm_checkpoint(vtkpath, \"baroclinic\", mpicomm, vtkstep[3] - 1)\n\n rm_checkpoint(vtkpath, \"barotropic\", mpicomm, vtkstep[4] - 1)\n\n vtkstep[3] += 1\n vtkstep[4] += 1\n nothing\n end\n\n return (cbvtk_slow, cbvtk_fast, cbinfo, cb_checkpoint, cbcs_dg)\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_coriolis.jl", "content": "#!/usr/bin/env julia --project\nusing Test\n\ninclude(\"hydrostatic_spindown.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/hydrostatic_spindown_refvals.jl\")\n\n # simulation time\n timeend = FT(15 * 24 * 3600) # s\n tout = FT(24 * 3600) # s\n timespan = (tout, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(8)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n )\n config = SplitConfig(\n \"rotating_bla\",\n resolution,\n dimensions,\n Coupled(),\n Rotating();\n solver = SplitExplicitSolver,\n boundary_conditions = BC,\n )\n\n #=\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorFreeSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceNoStressNoForcing(),\n )\n\n config = SplitConfig(\n \"rotating_jmc\",\n resolution,\n dimensions,\n Coupled(),\n Rotating();\n solver = SplitExplicitLSRK2nSolver,\n boundary_conditions = BC,\n )\n =#\n\n run_split_explicit(\n config,\n timespan,\n dt_fast = 300,\n dt_slow = 300, # 90 * 60,\n # refDat = refVals.ninety_minutes,\n analytic_solution = true,\n )\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_restart.jl", "content": "#!/usr/bin/env julia --project\nusing Test\n\ninclude(\"hydrostatic_spindown.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/hydrostatic_spindown_refvals.jl\")\n\n # simulation time\n timeend = FT(24 * 3600) # s\n tout = FT(3 * 3600) # s\n timespan = (tout, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(8)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n config = SplitConfig(\"test_restart\", resolution, dimensions, Coupled())\n\n midpoint = timeend / 2\n timespan = (tout, midpoint)\n\n run_split_explicit(config, timespan; dt_slow = 90 * 60)\n\n run_split_explicit(\n config,\n timespan;\n dt_slow = 90 * 60,\n refDat = refVals.ninety_minutes,\n analytic_solution = true,\n restart = Int(midpoint / tout),\n )\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_simple_box.jl", "content": "include(\"../../../experiments/OceanSplitExplicit/simple_box.jl\")\n\nClimateMachine.init()\n\n# Float type\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/simple_box_ivd_refvals.jl\")\n refDat = (refVals[1], refPrecs[1])\n\n # simulation time\n timestart = FT(0) # s\n timeend = FT(5 * 86400) # s\n timespan = (timestart, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(20)\n Nʸ = Int(20)\n Nᶻ = Int(20)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 4e6 # m\n Lʸ = 4e6 # m\n H = 1000 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.CoastlineNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorNoSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceStressForcing(),\n )\n\n config, solver_type = config_simple_box(\n \"test_simple_box\",\n resolution,\n dimensions,\n BC;\n dt_slow = FT(90 * 60),\n dt_fast = FT(240),\n )\n\n run_simple_box(config, timespan, solver_type.dt_slow; refDat = refDat)\n\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_spindown_long.jl", "content": "#!/usr/bin/env julia --project\nusing Test\n\ninclude(\"hydrostatic_spindown.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/hydrostatic_spindown_refvals.jl\")\n\n # simulation time\n timeend = FT(24 * 3600) # s\n tout = FT(3 * 3600) # s\n timespan = (tout, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(8)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n @testset \"Single-Rate\" begin\n @testset \"Not Coupled\" begin\n config =\n SplitConfig(\"uncoupled\", resolution, dimensions, Uncoupled())\n\n run_split_explicit(\n config,\n timespan,\n refDat = refVals.uncoupled,\n analytic_solution = true,\n )\n end\n\n @testset \"Fully Coupled\" begin\n config = SplitConfig(\"coupled\", resolution, dimensions, Coupled())\n\n run_split_explicit(\n config,\n timespan,\n refDat = refVals.coupled,\n analytic_solution = true,\n )\n end\n end\n\n @testset \"Multi-rate\" begin\n @testset \"Δt = 30 mins\" begin\n config = SplitConfig(\"multirate\", resolution, dimensions, Coupled())\n\n run_split_explicit(\n config,\n timespan,\n dt_slow = 30 * 60,\n refDat = refVals.thirty_minutes,\n analytic_solution = true,\n )\n end\n\n @testset \"Δt = 60 mins\" begin\n config = SplitConfig(\"multirate\", resolution, dimensions, Coupled())\n\n run_split_explicit(\n config,\n timespan,\n dt_slow = 60 * 60,\n refDat = refVals.sixty_minutes,\n analytic_solution = true,\n )\n end\n\n @testset \"Δt = 90 mins\" begin\n config = SplitConfig(\"multirate\", resolution, dimensions, Coupled())\n\n run_split_explicit(\n config,\n timespan,\n dt_slow = 90 * 60,\n refDat = refVals.ninety_minutes,\n analytic_solution = true,\n )\n end\n end\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_spindown_short.jl", "content": "#!/usr/bin/env julia --project\n\ninclude(\"hydrostatic_spindown.jl\")\nClimateMachine.init()\n\nconst FT = Float64\n\n#################\n# RUN THE TESTS #\n#################\n@testset \"$(@__FILE__)\" begin\n\n include(\"../refvals/hydrostatic_spindown_refvals.jl\")\n\n # simulation time\n timeend = FT(24 * 3600) # s\n tout = FT(1.5 * 3600) # s\n timespan = (tout, timeend)\n\n # DG polynomial order\n N = Int(4)\n\n # Domain resolution\n Nˣ = Int(5)\n Nʸ = Int(5)\n Nᶻ = Int(8)\n resolution = (N, Nˣ, Nʸ, Nᶻ)\n\n # Domain size\n Lˣ = 1e6 # m\n Lʸ = 1e6 # m\n H = 400 # m\n dimensions = (Lˣ, Lʸ, H)\n\n BC = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n )\n config = SplitConfig(\n \"spindown_bla\",\n resolution,\n dimensions,\n Coupled();\n solver = SplitExplicitSolver,\n boundary_conditions = BC,\n )\n\n #=\n BC = (\n ClimateMachine.Ocean.SplitExplicit01.OceanFloorFreeSlip(),\n ClimateMachine.Ocean.SplitExplicit01.OceanSurfaceNoStressNoForcing(),\n )\n\n config = SplitConfig(\n \"spindown_jmc\",\n resolution,\n dimensions,\n Coupled();\n solver = SplitExplicitLSRK2nSolver,\n boundary_conditions = BC,\n )\n =#\n\n run_split_explicit(\n config,\n timespan;\n dt_fast = 300, # seconds\n dt_slow = 90 * 60, # seconds\n # refDat = refVals.ninety_minutes,\n analytic_solution = true,\n )\nend\n" }, { "path": "test/Ocean/SplitExplicit/test_vertical_integral_model.jl", "content": "#!/usr/bin/env julia --project\nusing Test\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Mesh.Grids: polynomialorders\nusing ClimateMachine.Ocean.HydrostaticBoussinesq\nusing ClimateMachine.Ocean.ShallowWater\nusing ClimateMachine.Ocean.SplitExplicit: VerticalIntegralModel\nusing ClimateMachine.Ocean.OceanProblems\n\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VTK\n\nusing MPI\nusing LinearAlgebra\nusing StaticArrays\nusing Logging, Printf, Dates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nfunction test_vertical_integral_model(::Type{FT}, time; refDat = ()) where {FT}\n mpicomm = MPI.COMM_WORLD\n ArrayType = ClimateMachine.array_type()\n\n\n brickrange_2D = (xrange, yrange)\n topl_2D = BrickTopology(\n mpicomm,\n brickrange_2D,\n periodicity = (true, true),\n boundary = ((0, 0), (0, 0)),\n )\n grid_2D = DiscontinuousSpectralElementGrid(\n topl_2D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n brickrange_3D = (xrange, yrange, zrange)\n topl_3D = StackedBrickTopology(\n mpicomm,\n brickrange_3D;\n periodicity = (true, true, false),\n boundary = ((0, 0), (0, 0), (1, 2)),\n )\n grid_3D = DiscontinuousSpectralElementGrid(\n topl_3D,\n FloatType = FT,\n DeviceArray = ArrayType,\n polynomialorder = N,\n )\n\n problem = SimpleBox{FT}(Lˣ, Lʸ, H)\n\n model_3D = HydrostaticBoussinesqModel{FT}(\n param_set,\n problem;\n cʰ = FT(1),\n αᵀ = FT(0),\n κʰ = FT(0),\n κᶻ = FT(0),\n fₒ = FT(0),\n β = FT(0),\n )\n\n model_2D = ShallowWaterModel{FT}(\n param_set,\n problem,\n ShallowWater.ConstantViscosity{FT}(model_3D.νʰ),\n nothing;\n c = FT(1),\n fₒ = FT(0),\n β = FT(0),\n )\n\n integral_bl = VerticalIntegralModel(model_3D)\n\n integral_model = DGModel(\n integral_bl,\n grid_3D,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n dg_3D = DGModel(\n model_3D,\n grid_3D,\n RusanovNumericalFlux(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n dg_2D = DGModel(\n model_2D,\n grid_2D,\n CentralNumericalFluxFirstOrder(),\n CentralNumericalFluxSecondOrder(),\n CentralNumericalFluxGradient(),\n )\n\n Q_3D = init_ode_state(dg_3D, FT(time); init_on_cpu = true)\n Q_2D = init_ode_state(dg_2D, FT(time); init_on_cpu = true)\n Q_int = integral_model.state_auxiliary\n\n state_check = ClimateMachine.StateCheck.sccreate(\n [\n (Q_int, \"∫u\")\n (Q_2D, \"U\")\n ],\n 1;\n prec = 12,\n )\n\n update_auxiliary_state!(integral_model, integral_bl, Q_3D, time)\n GenericCallbacks.call!(state_check, nothing, nothing, nothing, nothing)\n\n ClimateMachine.StateCheck.scprintref(state_check)\n if length(refDat) > 0\n @test ClimateMachine.StateCheck.scdocheck(state_check, refDat)\n end\n\n return nothing\nend\n\n#################\n# RUN THE TESTS #\n#################\nconst FT = Float64\n\nconst N = 4\nconst Nˣ = 5\nconst Nʸ = 5\nconst Nᶻ = 8\nconst Lˣ = 1e6 # m\nconst Lʸ = 1e6 # m\nconst H = 400 # m\n\nxrange = range(FT(0); length = Nˣ + 1, stop = Lˣ)\nyrange = range(FT(0); length = Nʸ + 1, stop = Lʸ)\nzrange = range(FT(-H); length = Nᶻ + 1, stop = 0)\n\nconst cʰ = 1 # typical of ocean internal-wave speed\nconst cᶻ = 0\n\n@testset \"$(@__FILE__)\" begin\n include(\"../refvals/test_vertical_integral_model_refvals.jl\")\n\n times =\n [0, 86400, 30 * 86400, 365 * 86400, 10 * 365 * 86400, 100 * 365 * 86400]\n for (index, time) in enumerate(times)\n @testset \"$(time)\" begin\n test_vertical_integral_model(FT, time, refDat = refVals[index])\n end\n end\nend\n" }, { "path": "test/Ocean/refvals/2D_hydrostatic_spindown_refvals.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\nparr = [\n [\"2D state\", \"η\", 12, 12, 0, 12],\n [\"2D state\", \"U[1]\", 12, 12, 0, 12],\n [\"2D state\", \"U[2]\", 0, 0, 0, 0],\n]\n\nexplicit = [\n [\n \"2D state\",\n \"η\",\n -8.52722969951589915283e-01,\n 8.52846676313531282254e-01,\n -2.49578135935735214742e-16,\n 6.03454239990563690021e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15431401945821825450e+01,\n 3.15431401945818628008e+01,\n 6.11504145930918957291e-15,\n 2.24273815174625497093e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -7.62224398365580242501e-13,\n 9.72156930292624284356e-13,\n 1.39269607441935025982e-14,\n 1.95606703846656748360e-13,\n ],\n]\n\nrefVals = (explicit = (explicit, parr),)\n" }, { "path": "test/Ocean/refvals/3D_hydrostatic_spindown_refvals.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\nparr = [\n [\"Q\", \"u[1]\", 12, 12, 0, 12],\n [\"Q\", \"u[2]\", 0, 0, 0, 0],\n [\"Q\", \"η\", 12, 12, 0, 12],\n [\"Q\", \"θ\", 15, 15, 15, 15],\n [\"s_aux\", \"y\", 15, 15, 15, 15],\n [\"s_aux\", \"w\", 12, 12, 0, 12],\n [\"s_aux\", \"pkin\", 15, 15, 15, 15],\n [\"s_aux\", \"wz0\", 12, 12, 0, 12],\n [\"aux\", \"uᵈ[1]\", 15, 15, 15, 15],\n [\"aux\", \"uᵈ[2]\", 15, 15, 15, 15],\n [\"aux\", \"ΔGᵘ[1]\", 15, 15, 15, 15],\n [\"aux\", \"ΔGᵘ[2]\", 15, 15, 15, 15],\n]\n\n### fully explicit\nexplicit = [\n [\n \"state\",\n \"u[1]\",\n -9.58544066049463849843e-01,\n 9.58544066049465071089e-01,\n -6.13908923696726568442e-17,\n 4.45400263687296238402e-01,\n ],\n [\n \"state\",\n \"u[2]\",\n -4.33260855197780914020e-14,\n 1.99397359064900580713e-14,\n -1.30343101521973575132e-16,\n 2.95525689323845820606e-15,\n ],\n [\n \"state\",\n \"η\",\n -8.52732886154656810618e-01,\n 8.52845586939211197652e-01,\n 2.20052243093959998331e-14,\n 6.02992088522925295813e-01,\n ],\n [\n \"state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"aux\",\n \"w\",\n -4.04553460063758398447e-04,\n 4.04714358463272711169e-04,\n 4.75730566051879549438e-19,\n 1.63958655681888576441e-04,\n ],\n [\n \"aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"wz0\",\n -2.01164684799271339293e-04,\n 2.01041968159484089494e-04,\n -2.10942374678779754294e-20,\n 1.42228420244455277133e-04,\n ],\n [\n \"aux\",\n \"uᵈ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"uᵈ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"ΔGᵘ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"ΔGᵘ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nimex = [\n [\n \"state\",\n \"u[1]\",\n -9.57705359685807500192e-01,\n 9.57705359685806500991e-01,\n 4.54747350886464140750e-17,\n 4.45323727947873004851e-01,\n ],\n [\n \"state\",\n \"u[2]\",\n -4.07596731389788922865e-14,\n 2.53992382675900717845e-14,\n 9.88087105439151860723e-18,\n 2.63742927063789745855e-15,\n ],\n [\n \"state\",\n \"η\",\n -8.55941716778505390373e-01,\n 8.56014908798837681481e-01,\n 2.16732587432488803528e-14,\n 6.05260814296588400829e-01,\n ],\n [\n \"state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"aux\",\n \"w\",\n -4.03419301824637895407e-04,\n 4.03626959596159180614e-04,\n -3.33066907387546970565e-21,\n 1.63816237485524198066e-04,\n ],\n [\n \"aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"wz0\",\n -1.97016630922629948884e-04,\n 1.97066086934914155761e-04,\n -7.49400541621980656312e-19,\n 1.39393318574472925390e-04,\n ],\n [\n \"aux\",\n \"uᵈ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"uᵈ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"ΔGᵘ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"aux\",\n \"ΔGᵘ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nrefVals = (explicit = (explicit, parr), imex = (imex, parr))\n" }, { "path": "test/Ocean/refvals/hydrostatic_spindown_refvals.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\nparr = [\n [\"3D state\", \"u[1]\", 12, 12, 0, 12],\n [\"3D state\", \"u[2]\", 0, 0, 0, 0],\n [\"3D state\", \"η\", 12, 12, 0, 12],\n [\"3D state\", \"θ\", 15, 15, 15, 15],\n [\"3D aux\", \"y\", 15, 15, 15, 15],\n [\"3D aux\", \"w\", 12, 12, 0, 12],\n [\"3D aux\", \"pkin\", 15, 15, 15, 15],\n [\"3D aux\", \"wz0\", 12, 12, 0, 12],\n [\"3D aux\", \"uᵈ[1]\", 12, 12, 0, 12],\n [\"3D aux\", \"uᵈ[2]\", 0, 0, 0, 0],\n [\"3D aux\", \"ΔGᵘ[1]\", 11, 11, 0, 12],\n [\"3D aux\", \"ΔGᵘ[2]\", 0, 0, 0, 0],\n [\"2D state\", \"η\", 12, 12, 0, 12],\n [\"2D state\", \"U[1]\", 12, 12, 0, 12],\n [\"2D state\", \"U[2]\", 0, 0, 0, 0],\n [\"2D aux\", \"y\", 15, 15, 15, 2],\n [\"2D aux\", \"Gᵁ[1]\", 11, 11, 0, 12],\n [\"2D aux\", \"Gᵁ[2]\", 0, 0, 0, 0],\n [\"2D aux\", \"Δu[1]\", 12, 12, 0, 12],\n [\"2D aux\", \"Δu[2]\", 0, 0, 0, 0],\n]\n\nuncoupled = [\n [\n \"3D state\",\n \"u[1]\",\n -9.58547428246830479637e-01,\n 9.58547428246829258391e-01,\n -6.82121026329696195717e-18,\n 4.45400655325317695876e-01,\n ],\n [\n \"3D state\",\n \"u[2]\",\n -4.13574641907576004779e-14,\n 1.84325494853789287544e-14,\n 1.94567057477529826177e-17,\n 2.12815907156924710074e-15,\n ],\n [\n \"3D state\",\n \"η\",\n -8.52721734395889496838e-01,\n 8.52834259382379111791e-01,\n 2.18869899981655176688e-14,\n 6.02983573168680453414e-01,\n ],\n [\n \"3D state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"3D aux\",\n \"w\",\n -4.04489222828922911912e-04,\n 4.04649979048336849198e-04,\n 7.79376563286859913379e-19,\n 1.63949585575251901345e-04,\n ],\n [\n \"3D aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"wz0\",\n -2.00933099660541601228e-04,\n 2.00809859469744893647e-04,\n 2.62012633811536956698e-19,\n 1.42064682697776552755e-04,\n ],\n [\n \"3D aux\",\n \"uᵈ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"uᵈ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"2D state\",\n \"η\",\n -8.52715894965927923010e-01,\n 8.52861218874714666072e-01,\n -7.70938868299708696148e-16,\n 6.03458382295490314284e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15423945317149438949e+01,\n 3.15423945317147342848e+01,\n 1.42724412149908289268e-14,\n 2.24262219407601044452e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -9.38673083374999972374e-13,\n 1.16534781393478256382e-12,\n 1.93385973223830695204e-14,\n 2.16634164433978071172e-13,\n ],\n [\n \"2D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"2D aux\",\n \"Gᵁ[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"2D aux\",\n \"Gᵁ[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"2D aux\",\n \"Δu[1]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"2D aux\",\n \"Δu[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\ncoupled = [\n [\n \"3D state\",\n \"u[1]\",\n -9.58544854429326798062e-01,\n 9.58544854429326798062e-01,\n -4.09272615797817732838e-17,\n 4.45400401208546903309e-01,\n ],\n [\n \"3D state\",\n \"u[2]\",\n -4.09119364541065205420e-15,\n 3.41954410407034240802e-15,\n 1.57855184364031418231e-19,\n 7.69982545727116047928e-16,\n ],\n [\n \"3D state\",\n \"η\",\n -8.52733075123627615177e-01,\n 8.52843573070954374948e-01,\n 7.27595761418342649852e-17,\n 6.02992194663175995473e-01,\n ],\n [\n \"3D state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"3D aux\",\n \"w\",\n -4.04456725919283483547e-04,\n 4.04616610112968017130e-04,\n 1.33892896769793873666e-18,\n 1.63945200913093809460e-04,\n ],\n [\n \"3D aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"wz0\",\n -2.00813442787548966564e-04,\n 2.00686991710668984502e-04,\n 1.68642877440561282675e-18,\n 1.41979510156587984725e-04,\n ],\n [\n \"3D aux\",\n \"uᵈ[1]\",\n -8.79713597641291755735e-01,\n 8.79713597641291200624e-01,\n -9.26547727431170607429e-17,\n 4.41871673548253185437e-01,\n ],\n [\n \"3D aux\",\n \"uᵈ[2]\",\n -5.88884665747485312740e-28,\n 6.47654803186450692243e-28,\n 1.40221555861782666734e-30,\n 4.22950073030020620100e-29,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[1]\",\n -1.50292385646252865781e-09,\n 1.50292385640701573972e-09,\n -3.03178502810731714103e-21,\n 6.72141424487058536370e-10,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[2]\",\n -2.50330922252795548175e-19,\n 2.79421319642066487612e-19,\n 8.87468518373638336981e-36,\n 5.19465794381941235415e-20,\n ],\n [\n \"2D state\",\n \"η\",\n -8.52733075123627615177e-01,\n 8.52843573070954374948e-01,\n 1.61115565333602722195e-16,\n 6.03463098440266798583e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15402749945522060671e+01,\n 3.15402749945527389741e+01,\n 2.12974776703234167348e-14,\n 2.24262621677375086904e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -1.63647745816416607002e-12,\n 1.36781764162808076476e-12,\n 6.31420737450488849608e-17,\n 3.08233543917742288569e-13,\n ],\n [\n \"2D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"2D aux\",\n \"Gᵁ[1]\",\n -6.01169542562806285963e-07,\n 6.01169542585011466433e-07,\n 1.21270571032004377760e-18,\n 2.69066532005499711718e-07,\n ],\n [\n \"2D aux\",\n \"Gᵁ[2]\",\n -1.11768527856826595815e-16,\n 1.00132368901118217729e-16,\n -3.50057026691823957451e-33,\n 2.07948587451660775232e-17,\n ],\n [\n \"2D aux\",\n \"Δu[1]\",\n -6.53936429129693404076e-04,\n 6.53936429129569046087e-04,\n -2.70378337774435083119e-18,\n 4.64975945389895016328e-04,\n ],\n [\n \"2D aux\",\n \"Δu[2]\",\n -3.76049176138251460297e-17,\n 3.57641761985933683411e-17,\n -4.12443309266433447246e-19,\n 7.09162842149286857838e-18,\n ],\n]\n\nthirty_minutes = [\n [\n \"3D state\",\n \"u[1]\",\n -9.58527737426365766815e-01,\n 9.58527737426364989659e-01,\n -2.95585778076201651428e-17,\n 4.45397585011924779241e-01,\n ],\n [\n \"3D state\",\n \"u[2]\",\n -3.91776258115697290002e-15,\n 3.71573344368902785991e-15,\n 9.99453214450878842160e-18,\n 7.55533641450465537744e-16,\n ],\n [\n \"3D state\",\n \"η\",\n -8.52729582431237531637e-01,\n 8.52830944991449846349e-01,\n -2.54658516496419884307e-16,\n 6.02990141900936249542e-01,\n ],\n [\n \"3D state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"3D aux\",\n \"w\",\n -4.06712181304362465229e-04,\n 4.06869755985816939844e-04,\n 9.07052211118752913371e-19,\n 1.64290385445136678539e-04,\n ],\n [\n \"3D aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"wz0\",\n -2.09003099782708142351e-04,\n 2.08868143992394606845e-04,\n 1.33781874467331362536e-18,\n 1.47768608464728741676e-04,\n ],\n [\n \"3D aux\",\n \"uᵈ[1]\",\n -8.79792191807299950312e-01,\n 8.79792191807298729067e-01,\n -9.54969436861574689412e-17,\n 4.41911149114936674387e-01,\n ],\n [\n \"3D aux\",\n \"uᵈ[2]\",\n -4.49650715975976729085e-28,\n 3.22939933074851707840e-28,\n 1.28559371735991545201e-30,\n 3.12162656228539824003e-29,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[1]\",\n -1.52356619979685413105e-09,\n 1.52356619973800994984e-09,\n -3.00415481336788178489e-21,\n 6.81373145706173027634e-10,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[2]\",\n -2.17115289587171316651e-19,\n 2.99271337844227832879e-19,\n 3.57452597678270976423e-36,\n 5.23982958192948877053e-20,\n ],\n [\n \"2D state\",\n \"η\",\n -8.52729582431237531637e-01,\n 8.52830944991449846349e-01,\n -3.20454773827805192348e-16,\n 6.03461044074932395631e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15407643291233448224e+01,\n 3.15407643291237640426e+01,\n 2.30678864898692384736e-14,\n 2.24265574817805202201e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -1.56710503246274425174e-12,\n 1.48629337747559491236e-12,\n 3.99781285780299048032e-15,\n 3.02449468686897625244e-13,\n ],\n [\n \"2D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"2D aux\",\n \"Gᵁ[1]\",\n -6.09426479895203986555e-07,\n 6.09426479918741655731e-07,\n 1.20165633492970096897e-18,\n 2.72762104279987029506e-07,\n ],\n [\n \"2D aux\",\n \"Gᵁ[2]\",\n -1.19708535137691129300e-16,\n 8.68461158348685235787e-17,\n -1.50869648123518501517e-33,\n 2.09756864038773538444e-17,\n ],\n [\n \"2D aux\",\n \"Δu[1]\",\n -3.89449877638695009935e-03,\n 3.89449877638627485824e-03,\n -2.02396849009311295664e-17,\n 2.76903473119681107703e-03,\n ],\n [\n \"2D aux\",\n \"Δu[2]\",\n -2.13991366909841745553e-16,\n 2.02991354131742782465e-16,\n -1.26781868574011981865e-18,\n 4.34606820130057938426e-17,\n ],\n]\n\nsixty_minutes = [\n [\n \"3D state\",\n \"u[1]\",\n -9.58507705628042327994e-01,\n 9.58507705628042550039e-01,\n -5.22959453519433752618e-17,\n 4.45394100308371843067e-01,\n ],\n [\n \"3D state\",\n \"u[2]\",\n -4.09893498131451878178e-15,\n 3.66595063838075513859e-15,\n 3.52058149555652693147e-18,\n 7.65972478919525466697e-16,\n ],\n [\n \"3D state\",\n \"η\",\n -8.52715753705294066123e-01,\n 8.52819886810596838878e-01,\n 0.00000000000000000000e+00,\n 6.02986351734545844572e-01,\n ],\n [\n \"3D state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"3D aux\",\n \"w\",\n -4.09355343664679748733e-04,\n 4.09526445665209497954e-04,\n 8.01581023779363006131e-19,\n 1.64782726593626366188e-04,\n ],\n [\n \"3D aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"wz0\",\n -2.18597874938533713813e-04,\n 2.18512725277962983460e-04,\n 1.16684439888103955116e-18,\n 1.54581525183602542031e-04,\n ],\n [\n \"3D aux\",\n \"uᵈ[1]\",\n -8.79886113674857694988e-01,\n 8.79886113674856806810e-01,\n -8.81072992342524199517e-17,\n 4.41958323665245178535e-01,\n ],\n [\n \"3D aux\",\n \"uᵈ[2]\",\n -5.04476548888837049561e-28,\n 4.09418809809705127009e-28,\n 1.53142189509630563206e-30,\n 3.27579623353337388378e-29,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[1]\",\n -1.12461707299293835875e-09,\n 1.12461707291060059378e-09,\n -2.89803852573586336546e-21,\n 5.02954103857782471789e-10,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[2]\",\n -2.29220407034248886795e-19,\n 3.23795504633739099693e-19,\n 4.19082355898662551731e-36,\n 5.50021671320874203932e-20,\n ],\n [\n \"2D state\",\n \"η\",\n -8.52715753705294066123e-01,\n 8.52819886810596838878e-01,\n -2.04281036531028799987e-17,\n 6.03457250948630341547e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15415180428826182890e+01,\n 3.15415180428829948767e+01,\n 1.77902803599749859788e-14,\n 2.24265293235027023400e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -1.63957399252576030728e-12,\n 1.46638025535227397199e-12,\n 1.40823259822193190848e-15,\n 3.06628264537953242560e-13,\n ],\n [\n \"2D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"2D aux\",\n \"Gᵁ[1]\",\n -4.49846829164240224277e-07,\n 4.49846829197175353425e-07,\n 1.15921837490966064206e-18,\n 2.01338753352722549629e-07,\n ],\n [\n \"2D aux\",\n \"Gᵁ[2]\",\n -1.29518201853495637566e-16,\n 9.16881628136995516365e-17,\n -2.08062063752041870574e-33,\n 2.20180483211723137251e-17,\n ],\n [\n \"2D aux\",\n \"Δu[1]\",\n -7.71677156339491132631e-03,\n 7.71677156339359640591e-03,\n -2.97775205101297201914e-17,\n 5.48673197909435757247e-03,\n ],\n [\n \"2D aux\",\n \"Δu[2]\",\n -5.01794629524824969143e-16,\n 5.51655281127317078031e-16,\n 3.20748218455203547025e-18,\n 9.36761976343254920594e-17,\n ],\n]\n\nninety_minutes = [\n [\n \"3D state\",\n \"u[1]\",\n -9.58488051495743675900e-01,\n 9.58488051495742121588e-01,\n 2.50111042987555274331e-17,\n 4.45390687454859546257e-01,\n ],\n [\n \"3D state\",\n \"u[2]\",\n -1.23989905776825655281e-15,\n 1.61970487656192169626e-15,\n 3.83419479315953585092e-17,\n 3.55219724018798369983e-16,\n ],\n [\n \"3D state\",\n \"η\",\n -8.52731405452109680887e-01,\n 8.52809353848989926128e-01,\n -2.45563569478690627377e-16,\n 6.02988963894262375298e-01,\n ],\n [\n \"3D state\",\n \"θ\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.92775877460665535182e+05,\n ],\n [\n \"3D aux\",\n \"w\",\n -4.11959288170071617242e-04,\n 4.12126623290091649143e-04,\n -2.00395255944840751533e-19,\n 1.65355949740609838115e-04,\n ],\n [\n \"3D aux\",\n \"pkin\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"3D aux\",\n \"wz0\",\n -2.28048874139989487609e-04,\n 2.27949849600796259491e-04,\n -3.35287353436797287965e-19,\n 1.61271245461141072295e-04,\n ],\n [\n \"3D aux\",\n \"uᵈ[1]\",\n -8.79979607451056078382e-01,\n 8.79979607451055745315e-01,\n -2.84217094304040087969e-18,\n 4.42005283454923181274e-01,\n ],\n [\n \"3D aux\",\n \"uᵈ[2]\",\n -2.89610559829263959062e-28,\n 2.76495747279964637797e-28,\n -2.23087062731398024379e-31,\n 2.03879747098428489165e-29,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[1]\",\n -8.86991425991573654091e-10,\n 8.86991425890194811807e-10,\n -2.95809316169619294608e-22,\n 3.96682558403758137461e-10,\n ],\n [\n \"3D aux\",\n \"ΔGᵘ[2]\",\n -1.46453678699776496163e-19,\n 1.78310815453102811486e-19,\n 3.45126646034192634633e-36,\n 3.36942772885638802169e-20,\n ],\n [\n \"2D state\",\n \"η\",\n -8.52731405452109680887e-01,\n 8.52809353848989926128e-01,\n -2.24886775868071691631e-16,\n 6.03459865148300189652e-01,\n ],\n [\n \"2D state\",\n \"U[1]\",\n -3.15423820746574818941e+01,\n 3.15423820746572900475e+01,\n 1.25987324739451618828e-14,\n 2.24266774102466719398e+01,\n ],\n [\n \"2D state\",\n \"U[2]\",\n -4.95959623107294663097e-13,\n 6.47881950624733081946e-13,\n 1.53367791726381670695e-14,\n 1.42198852443335733848e-13,\n ],\n [\n \"2D aux\",\n \"y\",\n 0.00000000000000000000e+00,\n 1.00000000000000011642e+06,\n 5.00000000000000000000e+05,\n 2.93004519336559635121e+05,\n ],\n [\n \"2D aux\",\n \"Gᵁ[1]\",\n -3.54796570356077906525e-07,\n 3.54796570396629450056e-07,\n 1.18322201808542668922e-19,\n 1.58796938275634264439e-07,\n ],\n [\n \"2D aux\",\n \"Gᵁ[2]\",\n -7.13243261812411242092e-17,\n 5.85814714799106030876e-17,\n -1.27696859032651283100e-33,\n 1.34882362672174716756e-17,\n ],\n [\n \"2D aux\",\n \"Δu[1]\",\n -1.14629123365972158261e-02,\n 1.14629123365972678678e-02,\n 2.41175081207994926868e-18,\n 8.15149713200666835300e-03,\n ],\n [\n \"2D aux\",\n \"Δu[2]\",\n -7.92579521344275918936e-16,\n 6.08374854985514025679e-16,\n -3.41928350044524220733e-18,\n 1.39041335489639128848e-16,\n ],\n]\n\nrefVals = (\n uncoupled = (uncoupled, parr),\n coupled = (coupled, parr),\n thirty_minutes = (thirty_minutes, parr),\n sixty_minutes = (sixty_minutes, parr),\n ninety_minutes = (ninety_minutes, parr),\n)\n" }, { "path": "test/Ocean/refvals/simple_box_2dt_refvals.jl", "content": "refVals = []\nrefPrecs = []\n\n#! format: off\n# SC ========== Test number 1 reference values and precision match template. =======\n# SC ========== /home/jmc/cliMa/cliMa_update/test/Ocean/SplitExplicit/simple_box_2dt.jl test reference values ======================================\n# BEGIN SCPRINT\n# varr - reference values (from reference run)\n# parr - digits match precision (hand edit as needed)\n#\n# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\nvarr = [\n [ \"oce Q_3D\", \"u[1]\", -1.65970172440775332046e-01, 1.61603740424529351838e-01, 4.57235954923597420763e-03, 2.21382850057524928344e-02 ],\n [ \"oce Q_3D\", \"u[2]\", -2.16238066786232502325e-01, 2.44156974607258908661e-01, -2.57322116506966246802e-03, 2.49380331699317753236e-02 ],\n [ \"oce Q_3D\", \"η\", -7.98711302045854054654e-01, 3.47810107484795683064e-01, -2.92595480887335743312e-04, 2.91360229800833869795e-01 ],\n [ \"oce Q_3D\", \"θ\", 3.55385987992295690474e-03, 9.91576780951165659417e+00, 2.49883145933055006438e+00, 2.17846939512263038097e+00 ],\n [ \"oce aux\", \"w\", -1.66960462888799841472e-04, 1.65914321067276468273e-04, 5.22671200400887584530e-07, 1.42481192436579197488e-05 ],\n [ \"oce aux\", \"pkin\", -8.84676906676938301644e+00, 0.00000000000000000000e+00, -3.26734804711110893294e+00, 2.49685702267286213640e+00 ],\n [ \"oce aux\", \"wz0\", -2.12828008736777837579e-05, 3.08227212254399767725e-05, -1.55080158053877423592e-10, 7.82426303256959697840e-06 ],\n [ \"oce aux\", \"u_d[1]\", -1.56747674354107358052e-01, 1.23337664405837849069e-01, -7.85428994855595028748e-05, 1.15205292893568473495e-02 ],\n [ \"oce aux\", \"u_d[2]\", -2.12928714468491042666e-01, 2.30595749843429731474e-01, 2.75559843466661209680e-05, 1.75135782183724539318e-02 ],\n [ \"oce aux\", \"ΔGu[1]\", -2.33711927426000212944e-06, 3.31423592619396043514e-06, -4.04240250953691765080e-08, 3.37695011150041160627e-07 ],\n [ \"oce aux\", \"ΔGu[2]\", -2.38751859804088027546e-06, 2.23171788862484417825e-06, 1.26953194368922374152e-06, 7.64696321260357341382e-07 ],\n [ \"oce aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15573163901915703900e+06 ],\n [ \"baro Q_2D\", \"U[1]\", -1.77749182159153633620e+01, 6.89466489062268976795e+01, 4.64951906527867731000e+00, 1.81233032109943827948e+01 ],\n [ \"baro Q_2D\", \"U[2]\", -3.48869693974309527107e+01, 8.56944551257798678989e+01, -2.60002834745782607229e+00, 1.80715674809898771969e+01 ],\n [ \"baro Q_2D\", \"η\", -7.98718706462308802863e-01, 3.47811046235604659493e-01, -2.92587313874264492736e-04, 2.91375274868600042666e-01 ],\n [ \"baro aux\", \"Gᵁ[1]\", -3.31423592619396051306e-03, 2.33711927426000221075e-03, 4.04240250953691941898e-05, 3.37711728311069743838e-04 ],\n [ \"baro aux\", \"Gᵁ[2]\", -2.23171788862484409693e-03, 2.38751859804088020431e-03, -1.26953194368922394480e-03, 7.64734176576895561574e-04 ],\n [ \"baro aux\", \"U_c[1]\", -1.77747948208394639380e+01, 6.89456621278419703458e+01, 4.64949958798823903550e+00, 1.81231282979409442646e+01 ],\n [ \"baro aux\", \"U_c[2]\", -3.48859937354993405734e+01, 8.56944690523459655651e+01, -2.59996531383233087098e+00, 1.80715328159406638520e+01 ],\n [ \"baro aux\", \"η_c\", -7.98711302045854054654e-01, 3.47810107484795683064e-01, -2.92595480887348482688e-04, 2.91374653217579937525e-01 ],\n [ \"baro aux\", \"U_s[1]\", -1.77749182159153633620e+01, 6.89466489062268976795e+01, 4.64951906527867731000e+00, 1.81233032109943827948e+01 ],\n [ \"baro aux\", \"U_s[2]\", -3.48869693974309527107e+01, 8.56944551257798678989e+01, -2.60002834745782607229e+00, 1.80715674809898771969e+01 ],\n [ \"baro aux\", \"η_s\", -7.98718706462308802863e-01, 3.47811046235604659493e-01, -2.92587313874264492736e-04, 2.91375274868600042666e-01 ],\n [ \"baro aux\", \"Δu[1]\", -4.52063405565574774267e-04, 3.99219316002806577735e-04, -1.09219330063404869891e-05, 1.02341200736279461128e-04 ],\n [ \"baro aux\", \"Δu[2]\", -2.42431689759263984492e-04, 3.83040345117153786629e-04, 5.18877254212526286387e-05, 7.47027376508314802477e-05 ],\n [ \"baro aux\", \"η_diag\", -7.98659935505348972384e-01, 3.47918275378131136577e-01, -2.91845106891554467807e-04, 2.91369719089321743688e-01 ],\n [ \"baro aux\", \"Δη\", -5.74411984127720653959e-04, 7.42725686140199847785e-04, -7.50373995767764799686e-07, 1.62142594044870064609e-04 ],\n [ \"baro aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15578885204060329124e+06 ],\n]\nparr = [\n [ \"oce Q_3D\", \"u[1]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"u[2]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"η\", 12, 12, 8, 12 ],\n [ \"oce Q_3D\", \"θ\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"w\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"pkin\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"wz0\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"u_d[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"u_d[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"y\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[1]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[2]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"η\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Gᵁ[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Gᵁ[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_c\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"U_s[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_s[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_s\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δu[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Δu[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_diag\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δη\", 9, 9, 6, 10 ],\n [ \"baro aux\", \"y\", 12, 12, 12, 12 ],\n]\n# END SCPRINT\n# SC ====================================================================================\n\n append!(refVals ,[ varr ] )\n append!(refPrecs,[ parr ] )\n\n#! format: on\n" }, { "path": "test/Ocean/refvals/simple_box_ivd_refvals.jl", "content": "refVals = []\nrefPrecs = []\n\n#! format: off\n# SC ========== Test number 1 reference values and precision match template. =======\n# SC ========== /home/jmc/cliMa/cliMa_update/test/Ocean/SplitExplicit/simple_box_ivd.jl test reference values ======================================\n# BEGIN SCPRINT\n# varr - reference values (from reference run)\n# parr - digits match precision (hand edit as needed)\n#\n# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\nvarr = [\n [ \"oce Q_3D\", \"u[1]\", -1.66327525125104291881e-01, 1.63223287218068308091e-01, 4.57054199141500999692e-03, 2.21420008866952226778e-02 ],\n [ \"oce Q_3D\", \"u[2]\", -2.16296814964942768489e-01, 2.44105976363460708267e-01, -2.57418981108877616831e-03, 2.49256704872522875938e-02 ],\n [ \"oce Q_3D\", \"η\", -8.06721661019456748321e-01, 3.47844179114868090608e-01, -3.03380420083551127861e-04, 2.92158779168222249023e-01 ],\n [ \"oce Q_3D\", \"θ\", 3.58991077492186536763e-03, 9.91901429055192807027e+00, 2.49592172215209329167e+00, 2.17949780843674956188e+00 ],\n [ \"oce aux\", \"w\", -1.66963788550371410252e-04, 1.64989216196191281925e-04, 5.43118005323253843601e-07, 1.44011197008218777164e-05 ],\n [ \"oce aux\", \"pkin\", -8.84700789178211977060e+00, 0.00000000000000000000e+00, -3.26178716377690980366e+00, 2.50117471654691314598e+00 ],\n [ \"oce aux\", \"wz0\", -2.12922215345145814233e-05, 3.08186566412866553883e-05, -1.53993475195290930982e-10, 7.82327443240865498639e-06 ],\n [ \"oce aux\", \"u_d[1]\", -1.57099271619415420398e-01, 1.22268940698781983234e-01, -7.84250724128123411372e-05, 1.15452721709468873745e-02 ],\n [ \"oce aux\", \"u_d[2]\", -2.12964763257990491452e-01, 2.30578860790304179806e-01, 2.76337873040526178152e-05, 1.75119371638010334902e-02 ],\n [ \"oce aux\", \"ΔGu[1]\", -2.33686060871737704252e-06, 3.34281015324516805469e-06, -4.07662146058456868485e-08, 3.37377005914367152159e-07 ],\n [ \"oce aux\", \"ΔGu[2]\", -2.46095864522509354306e-06, 2.23139964389116549296e-06, 1.29016375934756163675e-06, 7.46140899703298511983e-07 ],\n [ \"oce aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15573163901915703900e+06 ],\n [ \"baro Q_2D\", \"U[1]\", -1.77642579751613922667e+01, 6.90024701711891168543e+01, 4.64758402961474637038e+00, 1.81125164737756669808e+01 ],\n [ \"baro Q_2D\", \"U[2]\", -3.48668457221064898022e+01, 8.56043954389619301537e+01, -2.60107324488959568143e+00, 1.80557567005594421516e+01 ],\n [ \"baro Q_2D\", \"η\", -8.06729084084832237522e-01, 3.47845108442749906263e-01, -3.03371430044830942430e-04, 2.92173862421732766226e-01 ],\n [ \"baro aux\", \"Gᵁ[1]\", -3.34281015324516816989e-03, 2.33686060871737691716e-03, 4.07662146058456834074e-05, 3.37393707332951826861e-04 ],\n [ \"baro aux\", \"Gᵁ[2]\", -2.23139964389116544213e-03, 2.46095864522509347530e-03, -1.29016375934756159609e-03, 7.46177836457347174598e-04 ],\n [ \"baro aux\", \"U_c[1]\", -1.77641349600227265171e+01, 6.90014815427443437557e+01, 4.64756464081321674087e+00, 1.81123414932683921563e+01 ],\n [ \"baro aux\", \"U_c[2]\", -3.48658691554727866446e+01, 8.56044062434442594167e+01, -2.60101037971774573521e+00, 1.80557220201615180599e+01 ],\n [ \"baro aux\", \"η_c\", -8.06721661019456748321e-01, 3.47844179114868090608e-01, -3.03380420083561807253e-04, 2.92173242116136766544e-01 ],\n [ \"baro aux\", \"U_s[1]\", -1.77642579751613922667e+01, 6.90024701711891168543e+01, 4.64758402961474637038e+00, 1.81125164737756669808e+01 ],\n [ \"baro aux\", \"U_s[2]\", -3.48668457221064898022e+01, 8.56043954389619301537e+01, -2.60107324488959568143e+00, 1.80557567005594421516e+01 ],\n [ \"baro aux\", \"η_s\", -8.06729084084832237522e-01, 3.47845108442749906263e-01, -3.03371430044830942430e-04, 2.92173862421732766226e-01 ],\n [ \"baro aux\", \"Δu[1]\", -4.51957131810988945505e-04, 3.99130692959488437947e-04, -1.09363921539353876238e-05, 1.02327771113594037156e-04 ],\n [ \"baro aux\", \"Δu[2]\", -2.42568998424025800446e-04, 3.82991529986384212410e-04, 5.18845385039475763891e-05, 7.47269801857268440920e-05 ],\n [ \"baro aux\", \"η_diag\", -8.06666755113148337131e-01, 3.47953196151883525911e-01, -3.02575448959868274160e-04, 2.92168272060050748795e-01 ],\n [ \"baro aux\", \"Δη\", -5.84450730272745300198e-04, 7.45283859099332701703e-04, -8.04971123704421912570e-07, 1.64194305918131369573e-04 ],\n [ \"baro aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15578885204060329124e+06 ],\n]\nparr = [\n [ \"oce Q_3D\", \"u[1]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"u[2]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"η\", 12, 12, 8, 12 ],\n [ \"oce Q_3D\", \"θ\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"w\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"pkin\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"wz0\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"u_d[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"u_d[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"y\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[1]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[2]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"η\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Gᵁ[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Gᵁ[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_c\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"U_s[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_s[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_s\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δu[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Δu[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_diag\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δη\", 9, 9, 6, 10 ],\n [ \"baro aux\", \"y\", 12, 12, 12, 12 ],\n]\n# END SCPRINT\n# SC ====================================================================================\n\n append!(refVals ,[ varr ] )\n append!(refPrecs,[ parr ] )\n\n#! format: on\n" }, { "path": "test/Ocean/refvals/simple_box_rk3_refvals.jl", "content": "refVals = []\nrefPrecs = []\n\n#! format: off\n# SC ========== Test number 1 reference values and precision match template. =======\n# SC ========== /home/jmc/cliMa/cliMa_new_jmc/test/Ocean/SplitExplicit/simple_box_rk3.jl test reference values ======================================\n# BEGIN SCPRINT\n# varr - reference values (from reference run)\n# parr - digits match precision (hand edit as needed)\n#\n# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\nvarr = [\n [ \"oce Q_3D\", \"u[1]\", -2.16721345244414054232e-01, 1.57840557320998220447e-01, -5.59921465916523013878e-03, 1.73941108901450175450e-02 ],\n [ \"oce Q_3D\", \"u[2]\", -2.19220347308738905401e-01, 2.31784478500299123693e-01, -2.15429258177158092225e-03, 2.12517725248480282563e-02 ],\n [ \"oce Q_3D\", \"η\", -4.69394735091986536890e-01, 2.06288744860751438459e-01, -2.16765734563377927271e-04, 1.47004018586567947180e-01 ],\n [ \"oce Q_3D\", \"θ\", 2.38571451036062890869e-03, 9.91209376647090323331e+00, 2.49760287312251527680e+00, 2.18031198350304133982e+00 ],\n [ \"oce aux\", \"w\", -1.61734977485578068895e-04, 1.50833500332270227301e-04, 5.43794052142437598582e-07, 1.38648560985387203702e-05 ],\n [ \"oce aux\", \"pkin\", -8.85948574009684719499e+00, 0.00000000000000000000e+00, -3.26449934601666758027e+00, 2.50111590039020725840e+00 ],\n [ \"oce aux\", \"wz0\", -3.66555541768163618680e-05, 1.94813350674614766593e-05, -4.61824895762852330113e-10, 9.37524239840742993843e-06 ],\n [ \"oce aux\", \"u_d[1]\", -1.56194066602061892857e-01, 1.21967899279333547025e-01, -2.05350727093651255306e-05, 1.15672828345751727702e-02 ],\n [ \"oce aux\", \"u_d[2]\", -2.09729915133787636616e-01, 2.37375904307950719163e-01, 1.58965003169083721666e-06, 1.81625473749615351515e-02 ],\n [ \"oce aux\", \"ΔGu[1]\", -2.09811375971404995481e-06, 3.23018047166822386688e-06, -4.22417846605267206689e-08, 2.80615396594168210400e-07 ],\n [ \"oce aux\", \"ΔGu[2]\", -2.18036232525426524906e-06, 2.31188192991968871615e-06, 1.27801756611837186519e-06, 7.19811396028193828918e-07 ],\n [ \"oce aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15573163901915703900e+06 ],\n [ \"baro Q_2D\", \"U[1]\", -3.61200008824836729104e+01, 2.32336820840376567787e+01, -5.58234606367367547364e+00, 1.11708003103345792084e+01 ],\n [ \"baro Q_2D\", \"U[2]\", -4.02361402227773581330e+01, 4.41897774443185085147e+01, -2.15561232944511127485e+00, 1.17915802114324606009e+01 ],\n [ \"baro Q_2D\", \"η\", -4.69794642977587939559e-01, 2.06346018349864157582e-01, -2.16567930962150958151e-04, 1.47057680237089760666e-01 ],\n [ \"baro aux\", \"Gᵁ[1]\", -3.23018047166822403968e-03, 2.09811375971404989044e-03, 4.22417846605267028812e-05, 2.80629288101644008488e-04 ],\n [ \"baro aux\", \"Gᵁ[2]\", -2.31188192991968856707e-03, 2.18036232525426528633e-03, -1.27801756611837201427e-03, 7.19847029373728297674e-04 ],\n [ \"baro aux\", \"U_c[1]\", -3.60967490653019851266e+01, 2.32266714099544557826e+01, -5.58248511658658941315e+00, 1.11690207251244704167e+01 ],\n [ \"baro aux\", \"U_c[2]\", -4.02002845377460147347e+01, 4.41870183291104439149e+01, -2.15756908928005142201e+00, 1.17881994981494511165e+01 ],\n [ \"baro aux\", \"η_c\", -4.69394735091986536890e-01, 2.06288744860751438459e-01, -2.16765734563369958385e-04, 1.47011295833105265496e-01 ],\n [ \"baro aux\", \"U_s[1]\", -3.61200008824836729104e+01, 2.32336820840376567787e+01, -5.58234606367367547364e+00, 1.11708003103345792084e+01 ],\n [ \"baro aux\", \"U_s[2]\", -4.02361402227773581330e+01, 4.41897774443185085147e+01, -2.15561232944511127485e+00, 1.17915802114324606009e+01 ],\n [ \"baro aux\", \"η_s\", -4.69794642977587939559e-01, 2.06346018349864157582e-01, -2.16567930962150958151e-04, 1.47057680237089760666e-01 ],\n [ \"baro aux\", \"Δu[1]\", -1.12410191785971484528e-03, 1.98901388437166311979e-03, 1.77992355967350647386e-04, 4.24716110102865956194e-04 ],\n [ \"baro aux\", \"Δu[2]\", -1.40780586763535704373e-03, 1.04939722938436467460e-03, -2.11977336567820129603e-04, 2.73797449735710581031e-04 ],\n [ \"baro aux\", \"η_diag\", -4.69273762873680611030e-01, 2.06327225515284040647e-01, -2.15810591775846216884e-04, 1.47008240438936316208e-01 ],\n [ \"baro aux\", \"Δη\", -2.62922736242607313351e-04, 2.37963172744791451318e-04, -9.55142787533847028187e-07, 6.94341553655788138637e-05 ],\n [ \"baro aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15578885204060329124e+06 ],\n]\nparr = [\n [ \"oce Q_3D\", \"u[1]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"u[2]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"η\", 12, 12, 8, 12 ],\n [ \"oce Q_3D\", \"θ\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"w\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"pkin\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"wz0\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"u_d[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"u_d[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"y\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[1]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[2]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"η\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Gᵁ[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Gᵁ[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_c\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"U_s[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_s[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_s\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δu[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Δu[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_diag\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δη\", 9, 9, 6, 10 ],\n [ \"baro aux\", \"y\", 12, 12, 12, 12 ],\n]\n# END SCPRINT\n# SC ====================================================================================\n\n append!(refVals ,[ varr ] )\n append!(refPrecs,[ parr ] )\n\n#! format: on\n" }, { "path": "test/Ocean/refvals/simple_dbl_gyre_refvals.jl", "content": "refVals = []\nrefPrecs = []\n\n#! format: off\n# SC ========== Test number 1 reference values and precision match template. =======\n# SC ========== /home/jmc/cliMa/cliMa_new_jmc/test/Ocean/SplitExplicit/simple_dbl_gyre.jl test reference values ======================================\n# BEGIN SCPRINT\n# varr - reference values (from reference run)\n# parr - digits match precision (hand edit as needed)\n#\n# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\nvarr = [\n [ \"oce Q_3D\", \"u[1]\", -1.41235234130601461366e-01, 2.44137911410471170059e-01, -1.17638792697003554538e-02, 4.76422973396204429974e-02 ],\n [ \"oce Q_3D\", \"u[2]\", -3.71526973402696303328e-01, 2.63129391123315958811e-01, -3.84758873236049556144e-02, 4.34631279346227028526e-02 ],\n [ \"oce Q_3D\", \"η\", -3.62315954238604787108e+00, 3.25374835544518203889e+00, -1.15530630772209001104e-04, 1.69499420692026547819e+00 ],\n [ \"oce Q_3D\", \"θ\", -1.95989738813091111946e-02, 2.41980669630830433903e+01, 6.00291945213983790808e+00, 5.36502088796993170661e+00 ],\n [ \"oce aux\", \"w\", -1.42436775039925483770e-03, 1.17093883909037276177e-03, -4.27143547438785408186e-07, 1.58700821328606480297e-04 ],\n [ \"oce aux\", \"pkin\", -6.47863971232045656734e+01, 0.00000000000000000000e+00, -2.35470923626465094003e+01, 1.84884052655973327717e+01 ],\n [ \"oce aux\", \"wz0\", -2.03056757204373681007e-04, 2.51217746755789830913e-04, -1.04288746719621225719e-08, 9.09212328801199519022e-05 ],\n [ \"oce aux\", \"u_d[1]\", -1.29380048170506856131e-01, 2.41847160349357936937e-01, 1.50957373152410783811e-04, 4.55063940872466113352e-02 ],\n [ \"oce aux\", \"u_d[2]\", -2.31151064740037659462e-01, 3.17638906635839823878e-01, 7.38946153332456550028e-05, 3.35035482074258969543e-02 ],\n [ \"oce aux\", \"ΔGu[1]\", -2.28025960163159889874e-05, 3.46723882659755934413e-06, -5.56600796147506152338e-07, 2.61685583178905381934e-06 ],\n [ \"oce aux\", \"ΔGu[2]\", -3.00923384704360015996e-06, 1.24960695834485056166e-05, 6.54681130814259427502e-06, 3.13708534178777629812e-06 ],\n [ \"oce aux\", \"y\", 0.00000000000000000000e+00, 6.00000000000000093132e+06, 3.00000000000000000000e+06, 1.73273876310492726043e+06 ],\n [ \"baro Q_2D\", \"U[1]\", -1.37905164694259781299e+02, 1.09296887114357033965e+02, -3.59834173192416955089e+01, 4.24653520384661575804e+01 ],\n [ \"baro Q_2D\", \"U[2]\", -4.93706839547863808093e+02, 5.71157285218092027890e+01, -1.15735478409127139798e+02, 8.23847078816430951065e+01 ],\n [ \"baro Q_2D\", \"η\", -3.62624136658312767878e+00, 3.25919860855558907176e+00, -1.15147341274860083972e-04, 1.69658109382101818241e+00 ],\n [ \"baro aux\", \"Gᵁ[1]\", -1.04017164797926778969e-02, 6.84077880489479678294e-02, 1.66980238844251878058e-03, 7.85082570477715034618e-03 ],\n [ \"baro aux\", \"Gᵁ[2]\", -3.74882087503455169175e-02, 9.02770154113080071367e-03, -1.96404339244277831300e-02, 9.41156556666298931002e-03 ],\n [ \"baro aux\", \"U_c[1]\", -1.36958883839043295438e+02, 1.09074110263089167461e+02, -3.57794851972332565992e+01, 4.23235946471143051895e+01 ],\n [ \"baro aux\", \"U_c[2]\", -4.92966277666792962009e+02, 5.72050394052077919582e+01, -1.15636183420627688179e+02, 8.22648167942200103653e+01 ],\n [ \"baro aux\", \"η_c\", -3.62315954238604787108e+00, 3.25374835544518203889e+00, -1.15530630772203336148e-04, 1.69504995619627174541e+00 ],\n [ \"baro aux\", \"U_s[1]\", -1.37905164694259781299e+02, 1.09296887114357033965e+02, -3.59834173192416955089e+01, 4.24653520384661575804e+01 ],\n [ \"baro aux\", \"U_s[2]\", -4.93706839547863808093e+02, 5.71157285218092027890e+01, -1.15735478409127139798e+02, 8.23847078816430951065e+01 ],\n [ \"baro aux\", \"η_s\", -3.62624136658312767878e+00, 3.25919860855558907176e+00, -1.15147341274860083972e-04, 1.69658109382101818241e+00 ],\n [ \"baro aux\", \"Δu[1]\", -4.66680666769807577128e-04, 1.05054947340137826844e-02, 3.56177517668052924515e-03, 2.34394235403509801352e-03 ],\n [ \"baro aux\", \"Δu[2]\", -2.98072019929583121103e-03, 4.75363215770287662193e-03, 1.22307915357422513150e-03, 1.31189027778619380499e-03 ],\n [ \"baro aux\", \"η_diag\", -3.62078008618114477457e+00, 3.25151939956774960194e+00, -1.15589924206263583179e-04, 1.69498993030730304987e+00 ],\n [ \"baro aux\", \"Δη\", -8.83137470915729139165e-03, 5.97876216292370088468e-03, 5.92934342415881362251e-08, 2.08693123903101827171e-03 ],\n [ \"baro aux\", \"y\", 0.00000000000000000000e+00, 6.00000000000000093132e+06, 3.00000000000000000000e+06, 1.73279575381979602389e+06 ],\n]\nparr = [\n [ \"oce Q_3D\", \"u[1]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"u[2]\", 12, 12, 12, 12 ],\n [ \"oce Q_3D\", \"η\", 12, 12, 8, 12 ],\n [ \"oce Q_3D\", \"θ\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"w\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"pkin\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"wz0\", 12, 12, 8, 12 ],\n [ \"oce aux\", \"u_d[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"u_d[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[1]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"ΔGu[2]\", 12, 12, 12, 12 ],\n [ \"oce aux\", \"y\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[1]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"U[2]\", 12, 12, 12, 12 ],\n [ \"baro Q_2D\", \"η\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Gᵁ[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Gᵁ[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_c[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_c\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"U_s[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"U_s[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_s\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δu[1]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"Δu[2]\", 12, 12, 12, 12 ],\n [ \"baro aux\", \"η_diag\", 12, 12, 8, 12 ],\n [ \"baro aux\", \"Δη\", 9, 9, 6, 10 ],\n [ \"baro aux\", \"y\", 12, 12, 12, 12 ],\n]\n# END SCPRINT\n# SC ====================================================================================\n\n append!(refVals ,[ varr ] )\n append!(refPrecs,[ parr ] )\n\n#! format: on\n" }, { "path": "test/Ocean/refvals/test_ocean_gyre_refvals.jl", "content": "parr = [\n [\"Q\", \"u[1]\", 12, 12, 12, 12],\n [\"Q\", \"u[2]\", 12, 12, 12, 12],\n [\"Q\", \"η\", 12, 12, 11, 12],\n [\"Q\", \"θ\", 12, 12, 12, 12],\n [\"s_aux\", \"y\", 12, 12, 12, 12],\n [\"s_aux\", \"w\", 12, 12, 11, 12],\n [\"s_aux\", \"pkin\", 12, 12, 12, 12],\n [\"s_aux\", \"wz0\", 12, 11, 10, 12],\n [\"s_aux\", \"uᵈ[1]\", 12, 12, 12, 12],\n [\"s_aux\", \"uᵈ[2]\", 12, 12, 12, 12],\n [\"s_aux\", \"ΔGᵘ[1]\", 12, 12, 12, 12],\n [\"s_aux\", \"ΔGᵘ[2]\", 12, 12, 12, 12],\n]\n\n#! format: off\n\nshort = [\n [ \"Q\", \"u[1]\", -1.56752732465427965791e-02, 1.68514505893861757380e-02, -2.29247099512640793717e-03, 3.25160701902235671490e-03 ],\n [ \"Q\", \"u[2]\", -3.79775189267773510826e-02, 6.43003815969073189152e-03, -1.34097283250433057383e-02, 1.20337368194613283934e-02 ],\n [ \"Q\", \"η\", -1.53249855976142324021e-01, 1.57769367725846931805e-01, -6.47997524778634250456e-06, 1.03985821375781786746e-01 ],\n [ \"Q\", \"θ\", 1.07842251824037485168e-05, 9.00370181779731204585e+00, 2.49971752106072075961e+00, 2.19711465681760964586e+00 ],\n [ \"s_aux\", \"y\", 0.00000000000000000000e+00, 1.00000000000000011642e+06, 5.00000000000000000000e+05, 2.92779390974978974555e+05 ],\n [ \"s_aux\", \"w\", -9.71167446044889304474e-05, 8.57892392958965760040e-05, 7.22305481630802806057e-08, 3.80815409312154189983e-05 ],\n [ \"s_aux\", \"pkin\", -8.99913049767501638243e-01, 0.00000000000000000000e+00, -3.32109083459310172604e-01, 2.56226532893150116266e-01 ],\n [ \"s_aux\", \"wz0\", -8.17753668537623428815e-05, 8.25631396299614581233e-05, -1.10719884491561329431e-09, 5.14580958965247714965e-05 ],\n [ \"s_aux\", \"uᵈ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"uᵈ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n]\n\nlong = [\n [ \"Q\", \"u[1]\", -7.88803234115370982549e-02, 6.54328011157656042052e-02, 2.01451525264662988784e-03, 1.24588716783014998718e-02 ],\n [ \"Q\", \"u[2]\", -8.71605274310923855419e-02, 1.47505267452481658719e-01, 5.63158201082450144553e-03, 1.28959719954838351180e-02 ],\n [ \"Q\", \"η\", -4.73491408442001826540e-01, 4.02693687285959778244e-01, -6.49059970677390036392e-05, 2.21689928090663096460e-01 ],\n [ \"Q\", \"θ\", 4.24292935192232840286e-04, 9.24539353455579338004e+00, 2.49938206627401893201e+00, 2.17986626392490689952e+00 ],\n [ \"s_aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15573163901915703900e+06 ],\n [ \"s_aux\", \"w\", -2.22086406767006932887e-04, 2.00575090959222558738e-04, 2.53168866096380013512e-07, 1.66257132341935760973e-05 ],\n [ \"s_aux\", \"pkin\", -9.00869877619916104017e-01, 0.00000000000000000000e+00, -3.33171488779369751043e-01, 2.54740287894525019308e-01 ],\n [ \"s_aux\", \"wz0\", -2.96608015795817572195e-05, 3.66759042312928121082e-05, 3.78116102337067175813e-10, 1.07073256826410757734e-05 ],\n [ \"s_aux\", \"uᵈ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"uᵈ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n]\n\n#! format: on\n\nrefVals = (short = (short, parr), long = (long, parr))\n" }, { "path": "test/Ocean/refvals/test_vertical_integral_model_refvals.jl", "content": "# [\n# [ MPIStateArray Name, Field Name, Maximum, Minimum, Mean, Standard Deviation ],\n# [ : : : : : : ],\n# ]\n\nparr = [\n [\"∫u\", \"∫x[1]\", 12, 12, 0, 12],\n [\"∫u\", \"∫x[2]\", 15, 15, 15, 15],\n [\"U\", \"η\", 12, 12, 0, 12],\n [\"U\", \"U[1]\", 12, 12, 0, 12],\n [\"U\", \"U[2]\", 15, 15, 15, 15],\n]\n\ninitial = [\n [\n \"∫u\",\n \"∫x[1]\",\n -6.36104748927323058183e+01,\n 6.36104748927324763486e+01,\n 3.59432306140661235747e-14,\n 3.16776758596447187699e+01,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -1.00000000000000000000e+00,\n 1.00000000000000000000e+00,\n 3.73034936274052574533e-18,\n 7.07673146340372483110e-01,\n ],\n [\n \"U\",\n \"U[1]\",\n -9.95282537582876658533e-01,\n 9.95282537582876547511e-01,\n -8.82958196287834196438e-18,\n 7.07673146340372150043e-01,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nday = [\n [\n \"∫u\",\n \"∫x[1]\",\n -6.40437964821475844701e+01,\n 6.40437964821476981570e+01,\n 2.21189111471176159638e-14,\n 2.60887556680598677872e+01,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -8.52761706341222169847e-01,\n 8.52761706341222169847e-01,\n -1.86517468137026310378e-17,\n 6.03476559805077306109e-01,\n ],\n [\n \"U\",\n \"U[1]\",\n -3.15397076561132330141e+01,\n 3.15397076561132294614e+01,\n -2.82332564752296143530e-15,\n 2.24255960582435314166e+01,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nmonth = [\n [\n \"∫u\",\n \"∫x[1]\",\n -3.51564130159214158766e+01,\n 3.51564130159213874549e+01,\n -1.16415321826934820032e-14,\n 1.41614248818105643579e+01,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -5.30601964901526002016e-01,\n 5.30601964901526002016e-01,\n -1.24344978758017535116e-17,\n 3.75492761956246756672e-01,\n ],\n [\n \"U\",\n \"U[1]\",\n -3.51564130159208829696e+01,\n 3.51564130159208758641e+01,\n -2.47942593560761348802e-15,\n 2.49971726354604832920e+01,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nyear = [\n [\n \"∫u\",\n \"∫x[1]\",\n -4.25110151351051179791e-01,\n 4.25110151351050846724e-01,\n -1.36424205265939223736e-16,\n 1.74619680058462317662e-01,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -4.39687965160425187072e-02,\n 4.39687965160425187072e-02,\n -6.10622663543836061796e-19,\n 3.11155365713074068268e-02,\n ],\n [\n \"U\",\n \"U[1]\",\n -4.25110151351044573964e-01,\n 4.25110151351044518453e-01,\n -1.99959679826973105365e-17,\n 3.02264961945818255717e-01,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\ndecade = [\n [\n \"∫u\",\n \"∫x[1]\",\n -1.88298126764002322427e-12,\n 1.88298126764002483985e-12,\n 5.62482816536058859320e-28,\n 7.73459738532225749582e-13,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -3.38352596570510824820e-15,\n 3.38352596570510824820e-15,\n 6.31088724176809474572e-34,\n 2.39443046587488032416e-15,\n ],\n [\n \"U\",\n \"U[1]\",\n -1.88298126763999575929e-12,\n 1.88298126763999575929e-12,\n 3.78999175038214983828e-29,\n 1.33885125866565214937e-12,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\ncentury = [\n [\n \"∫u\",\n \"∫x[1]\",\n -3.97994692621354800600e-134,\n 3.97994692621354409568e-134,\n -1.20124985710986044942e-149,\n 1.63481642745144852479e-134,\n ],\n [\n \"∫u\",\n \"∫x[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n [\n \"U\",\n \"η\",\n -2.07116984864437321907e-136,\n 2.07116984864437321907e-136,\n -5.62108290883927252001e-153,\n 1.46571128339547670108e-136,\n ],\n [\n \"U\",\n \"U[1]\",\n -3.97994692621348592969e-134,\n 3.97994692621348544091e-134,\n -4.15669815960076969677e-150,\n 2.82985128060348631090e-134,\n ],\n [\n \"U\",\n \"U[2]\",\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n 0.00000000000000000000e+00,\n ],\n]\n\nrefVals = (\n intitial = (initial, parr),\n day = (day, parr),\n month = (month, parr),\n year = (year, parr),\n decade = (decade, parr),\n century = (century, parr),\n)\n" }, { "path": "test/Ocean/refvals/test_windstress_refvals.jl", "content": "parr = [\n [\"Q\", \"u[1]\", 12, 12, 12, 12],\n [\"Q\", \"u[2]\", 12, 12, 12, 12],\n [\"Q\", \"η\", 12, 12, 11, 12],\n [\"Q\", \"θ\", 12, 12, 12, 0],\n [\"s_aux\", \"y\", 12, 12, 12, 12],\n [\"s_aux\", \"w\", 12, 12, 12, 12],\n [\"s_aux\", \"pkin\", 12, 12, 12, 12],\n [\"s_aux\", \"wz0\", 12, 12, 10, 12],\n [\"s_aux\", \"uᵈ[1]\", 12, 12, 12, 12],\n [\"s_aux\", \"uᵈ[2]\", 12, 12, 12, 12],\n [\"s_aux\", \"ΔGᵘ[1]\", 12, 12, 12, 12],\n [\"s_aux\", \"ΔGᵘ[2]\", 12, 12, 12, 12],\n]\n\n#! format: off\nimex = [\n [ \"Q\", \"u[1]\", -3.79026877309730850230e-02, 3.77520487392776424307e-02, -1.40391503429245286812e-06, 5.33572816364154614566e-03 ],\n [ \"Q\", \"u[2]\", -7.13128520899746973227e-03, 6.54348134460207026680e-03, -5.44654209456798120995e-06, 9.83670484102493409076e-04 ],\n [ \"Q\", \"η\", -5.75156897888350494147e-03, 5.06823909948787842961e-03, -1.61345029906403004787e-05, 1.58599716799621352596e-03 ],\n [ \"Q\", \"θ\", 1.99999999999999928946e+01, 2.00000000000002948752e+01, 2.00000000000000426326e+01, 6.30374859517183302003e-14 ],\n [ \"s_aux\", \"y\", 0.00000000000000000000e+00, 1.00000000000000011642e+06, 5.00000000000000000000e+05, 2.92779390974978974555e+05 ],\n [ \"s_aux\", \"w\", -1.93125254703910637023e-05, 1.89413928194037212366e-05, 6.96287201311761918867e-08, 1.82962329428446588302e-06 ],\n [ \"s_aux\", \"pkin\", -1.60000000000000408562e+00, 0.00000000000000000000e+00, -8.00000000000002042810e-01, 4.68447025559967145103e-01 ],\n [ \"s_aux\", \"wz0\", -1.78460606500689411704e-06, 1.42096532423882541587e-06, -3.15575167929213483841e-09, 7.00511144160156445126e-07 ],\n [ \"s_aux\", \"uᵈ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"uᵈ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n]\n\nexplicit_cpu = [\n [ \"Q\", \"u[1]\", -3.74270752639261211625e-02, 3.72763215301363109999e-02, -1.40392694287316997219e-06, 5.29521849931491837837e-03 ],\n [ \"Q\", \"u[2]\", -7.13776376792300999707e-03, 6.54949226335132476257e-03, -5.45004642311143043000e-06, 9.84283148803805074678e-04 ],\n [ \"Q\", \"η\", -5.75146380759523553894e-03, 5.06819905742867966164e-03, -1.61399463692112299070e-05, 1.58594255118538803723e-03 ],\n [ \"Q\", \"θ\", 1.99999999999996518341e+01, 2.00000000000013891110e+01, 2.00000000000004192202e+01, 2.61373866244148208458e-13 ],\n [ \"s_aux\", \"y\", 0.00000000000000000000e+00, 1.00000000000000011642e+06, 5.00000000000000000000e+05, 2.92779390974978974555e+05 ],\n [ \"s_aux\", \"w\", -1.91903846873268650749e-05, 1.88201368003043060118e-05, 6.88331165208082022114e-08, 1.81657738735220659856e-06 ],\n [ \"s_aux\", \"pkin\", -1.60000000000003339551e+00, 0.00000000000000000000e+00, -8.00000000000017919000e-01, 4.68447025559975749331e-01 ],\n [ \"s_aux\", \"wz0\", -1.79359066322475932405e-06, 1.42978124248321576704e-06, -3.01455602004428728418e-09, 6.97246485535118604352e-07 ],\n [ \"s_aux\", \"uᵈ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"uᵈ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n]\n\nexplicit_gpu = deepcopy(explicit_cpu)\nexplicit_gpu[3][5] = -1.61399463694379806270e-05 # CPU value: -1.61399463692112299070e-05 -> Δ = 2.267507199799762e-16\nexplicit_gpu[6][5] = 6.88331165196929377526e-08 # CPU value: 6.88331165208082022114e-08 -> Δ = -1.1152644588480957e-18\n\nlong = [\n [ \"Q\", \"u[1]\", -1.26047572708648997208e-01, 8.73949752192989676169e-02, 4.77293205335180186562e-05, 7.01620954203814161526e-03 ],\n [ \"Q\", \"u[2]\", -7.74039642996040555545e-02, 1.17079049312627345159e-01, -1.81789027475552530900e-04, 1.05959299008928087282e-02 ],\n [ \"Q\", \"η\", -7.88344390091704622092e-02, 3.98295576513588017731e-02, -9.41579560194082235179e-05, 2.80408708312616314351e-02 ],\n [ \"Q\", \"θ\", 1.99999999999542694695e+01, 2.00000000000501145792e+01, 2.00000000000023661073e+01, 1.20202839114602287074e-12 ],\n [ \"s_aux\", \"y\", 0.00000000000000000000e+00, 4.00000000000000046566e+06, 2.00000000000000000000e+06, 1.15573129229947482236e+06 ],\n [ \"s_aux\", \"w\", -2.39010107630067514077e-05, 6.21800814928817288142e-05, 2.57411032722755278831e-07, 4.26276984128505141513e-06 ],\n [ \"s_aux\", \"pkin\", -1.60000000000220565788e+00, 0.00000000000000000000e+00, -8.00000000000101851860e-01, 4.61946285916540189120e-01 ],\n [ \"s_aux\", \"wz0\", -2.51713525692209742640e-06, 9.22709404863107399641e-07, -8.84750036743458214697e-10, 5.52781925119928323386e-07 ],\n [ \"s_aux\", \"uᵈ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"uᵈ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[1]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n [ \"s_aux\", \"ΔGᵘ[2]\", 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00, 0.00000000000000000000e+00 ],\n]\n\n#! format: on\n\nrefVals = (\n explicit_cpu = (explicit_cpu, parr),\n explicit_gpu = (explicit_gpu, parr),\n imex = (imex, parr),\n long = (long, parr),\n)\n" }, { "path": "test/Ocean/runtests.jl", "content": "using MPI, Test\n\ninclude(\"../testhelpers.jl\")\n\n@testset \"Ocean\" begin\n runmpi(joinpath(@__DIR__, \"HydrostaticBoussinesq/test_ocean_gyre_short.jl\"))\n runmpi(joinpath(@__DIR__, \"SplitExplicit/test_spindown_short.jl\"))\n include(joinpath(\"OceanProblems\", \"test_initial_value_problem.jl\"))\n include(joinpath(\n \"HydrostaticBoussinesqModel\",\n \"test_hydrostatic_boussinesq_model.jl\",\n ))\nend\n" }, { "path": "test/Utilities/SingleStackUtils/runtests.jl", "content": "module TestSingleStackUtils\n\nusing MPI, Test\ninclude(joinpath(\"..\", \"..\", \"testhelpers.jl\"))\n\n@testset \"SingleStackUtils\" begin\n runmpi(joinpath(@__DIR__, \"ssu_tests.jl\"))\nend\n\nend\n" }, { "path": "test/Utilities/SingleStackUtils/ssu_tests.jl", "content": "using OrderedCollections\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.Grids\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.DGMethods: LocalGeometry\n\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Auxiliary, Prognostic, Gradient, GradientFlux\nimport ClimateMachine.BalanceLaws:\n vars_state, nodal_init_state_auxiliary!, init_state_prognostic!\n\nstruct EmptyBalLaw{FT, PS} <: BalanceLaw\n \"Parameters\"\n param_set::PS\n \"Domain height\"\n zmax::FT\n\nend\nEmptyBalLaw(param_set, zmax) =\n EmptyBalLaw{typeof(zmax), typeof(param_set)}(param_set, zmax)\n\nvars_state(::EmptyBalLaw, ::Auxiliary, FT) = @vars(x::FT, y::FT, z::FT)\nvars_state(::EmptyBalLaw, ::Prognostic, FT) = @vars(ρ::FT)\nvars_state(::EmptyBalLaw, ::Gradient, FT) = @vars()\nvars_state(::EmptyBalLaw, ::GradientFlux, FT) = @vars()\n\nfunction nodal_init_state_auxiliary!(\n m::EmptyBalLaw,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.x = geom.coord[1]\n aux.y = geom.coord[2]\n aux.z = geom.coord[3]\nend\n\nfunction init_state_prognostic!(\n m::EmptyBalLaw,\n state::Vars,\n aux::Vars,\n coords,\n t::Real,\n)\n z = aux.z\n x = aux.x\n y = aux.y\n state.ρ = (1 - 4 * (z - m.zmax / 2)^2) * (2 - x - y)\nend\n\nfunction test_hmean(\n grid::DiscontinuousSpectralElementGrid{T, dim, Ns},\n Q::MPIStateArray,\n vars,\n) where {T, dim, Ns}\n state_vars_avg = get_horizontal_mean(grid, Q, vars)\n target = target_meanprof(grid)\n @test state_vars_avg[\"ρ\"] ≈ target\nend\n\nfunction test_hvar(\n grid::DiscontinuousSpectralElementGrid{T, dim, Ns},\n Q::MPIStateArray,\n vars,\n) where {T, dim, Ns}\n state_vars_var = get_horizontal_variance(grid, Q, vars)\n target = target_varprof(grid)\n @test state_vars_var[\"ρ\"] ≈ target\nend\n\nfunction test_horizontally_ave(\n grid::DiscontinuousSpectralElementGrid{T, dim, Ns},\n Q_in::MPIStateArray,\n vars,\n) where {T, dim, Ns}\n Q = deepcopy(Q_in)\n state_vars_var = get_horizontal_variance(grid, Q, vars)\n i_vars = varsindex(vars, :ρ)\n horizontally_average!(grid, Q, i_vars)\n FT = eltype(Q)\n state_vars_var = get_horizontal_variance(grid, Q, vars)\n @test all(isapprox.(state_vars_var[\"ρ\"], 0, atol = 10 * eps(FT)))\nend\n\nfunction target_meanprof(\n grid::DiscontinuousSpectralElementGrid{T, dim, Ns},\n) where {T, dim, Ns}\n Nqs = Ns .+ 1\n Nq_v = Nqs[end]\n Ntot = Nq_v * grid.topology.stacksize\n z = Array(get_z(grid))\n target =\n SVector{Ntot, T}([1.0 - 4.0 * (z_i - z[Ntot] / 2.0)^2 for z_i in z])\n return target\nend\n\nfunction target_varprof(\n grid::DiscontinuousSpectralElementGrid{T, dim, Ns},\n) where {T, dim, Ns}\n Nqs = Ns .+ 1\n Nq_v = Nqs[end]\n nvertelem = grid.topology.stacksize\n Ntot = Nq_v * nvertelem\n z = Array(get_z(grid))\n x = z[1:Nq_v] * nvertelem\n scaled_var = 0.0\n for i in 1:Nq_v\n for j in 1:Nq_v\n scaled_var = scaled_var + (2 - x[i] - x[j]) * (2 - x[i] - x[j])\n end\n end\n target = SVector{Ntot, Float64}([\n (1.0 - 4.0 * (z_i - z[Ntot] / 2.0)^2) *\n (1.0 - 4.0 * (z_i - z[Ntot] / 2.0)^2) *\n (scaled_var / Nq_v / Nq_v - 1) for z_i in z\n ])\n return target\nend\n\nfunction main()\n FT = Float64\n ClimateMachine.init()\n\n m = EmptyBalLaw(param_set, FT(1))\n\n # Prescribe polynomial order of basis functions in finite elements\n N_poly = 5\n # Specify the number of vertical elements\n nelem_vert = 20\n # Specify the domain height\n zmax = m.zmax\n # Initial and final times\n t0 = 0.0\n timeend = 1.0\n dt = 0.1\n # Establish a `ClimateMachine` single stack configuration\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"SingleStackUtilsTest\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n\n # tests\n test_hmean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n test_hvar(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n\n r1, z1 = reduce_nodal_stack(\n max,\n solver_config.dg.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n \"ρ\",\n i = 6,\n j = 6,\n )\n @test r1 ≈ 8.880558532968455e-16 && z1 == 10\n r2, z2 = reduce_nodal_stack(\n +,\n solver_config.dg.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n \"ρ\",\n i = 3,\n j = 3,\n )\n @test r2 ≈ 102.73283921735293 && z2 == 20\n ns = reduce_element_stack(\n +,\n solver_config.dg.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n \"ρ\",\n )\n (r3, z3) = let\n f(a, b) = (a[1] + b[1], b[2])\n reduce(f, ns)\n end\n @test r3 ≈ FT(2877.6) && z3 == 20\n\n test_horizontally_ave(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n\n # Test NodalStack iterator (without interpolation)\n nodal_stack = NodalStack(solver_config; interp = false)\n ρ_iter = [local_states.prog.ρ for local_states in nodal_stack]\n dons = get_vars_from_nodal_stack(\n solver_config.dg.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT);\n interp = false,\n )\n ρ_vfns = dons[\"ρ\"]\n @test all(ρ_iter .≈ ρ_vfns)\n\n # Test NodalStack iterator (with interpolation)\n nodal_stack = NodalStack(solver_config; interp = true)\n ρ_iter = [local_states.prog.ρ for local_states in nodal_stack]\n dons = get_vars_from_nodal_stack(\n solver_config.dg.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT);\n interp = true,\n )\n ρ_vfns_interp = dons[\"ρ\"]\n @test all(ρ_iter .≈ ρ_vfns_interp)\n return nothing\nend\n\n@testset \"Single Stack Utils\" begin\n main()\nend\n" }, { "path": "test/Utilities/TicToc/runtests.jl", "content": "module TestTicToc\n\nusing Test\nusing ClimateMachine.TicToc\n\nfunction foo()\n @tic foo\n sleep(0.25)\n @toc foo\nend\n\nfunction bar()\n @tic bar\n sleep(1)\n @toc bar\nend\n\n\nif TicToc.tictoc_enabled\n @testset \"TicToc\" begin\n @test tictoc() >= 2\n foo_i = findfirst(s -> s == :tictoc__foo, TicToc.timing_info_names)\n bar_i = findfirst(s -> s == :tictoc__bar, TicToc.timing_info_names)\n @test foo_i != nothing\n @test bar_i != nothing\n foo()\n foo()\n @test TicToc.timing_infos[foo_i].ncalls == 2\n @test TicToc.timing_infos[bar_i].ncalls == 0\n @test TicToc.timing_infos[foo_i].time >= 5e8\n bar()\n @test TicToc.timing_infos[bar_i].ncalls == 1\n @test TicToc.timing_infos[bar_i].time >= 1e9\n buf = IOBuffer()\n old_stdout = stdout\n try\n rd, = redirect_stdout()\n TicToc.print_timing_info()\n Libc.flush_cstdio()\n flush(stdout)\n write(buf, readavailable(rd))\n finally\n redirect_stdout(old_stdout)\n end\n str = String(take!(buf))\n @test findnext(\"tictoc__foo\", str, 1) != nothing\n @test findnext(\"tictoc__bar\", str, 1) != nothing\n end\nend\n\nend # module TestTicToc\n" }, { "path": "test/Utilities/VariableTemplates/complex_models.jl", "content": "using Test\nusing StaticArrays\n\nabstract type AbstractModel end\nabstract type OneLayerModel <: AbstractModel end\n\nstruct EmptyModel <: OneLayerModel end\nstate(m::EmptyModel, T) = @vars()\n\nstruct ScalarModel <: OneLayerModel end\nstate(m::ScalarModel, T) = @vars(x::T)\n\nstruct VectorModel{N} <: OneLayerModel end\nstate(m::VectorModel{N}, T) where {N} = @vars(x::SVector{N, T})\n\nstruct MatrixModel{N, M} <: OneLayerModel end\nstate(m::MatrixModel{N, M}, T) where {N, M} =\n @vars(x::SHermitianCompact{N, T, M})\n\nabstract type TwoLayerModel <: AbstractModel end\n\nstruct CompositModel{EM, SM, VM, MM} <: TwoLayerModel\n empty_model::EM\n scalar_model::SM\n vector_model::VM\n matrix_model::MM\nend\nfunction CompositModel(\n Nv,\n N,\n M;\n empty_model = EmptyModel(),\n scalar_model = ScalarModel(),\n vector_model = VectorModel{Nv}(),\n matrix_model = MatrixModel{N, M}(),\n)\n args = (empty_model, scalar_model, vector_model, matrix_model)\n return CompositModel{typeof.(args)...}(args...)\nend\n\n\nfunction state(m::CompositModel, T)\n @vars begin\n empty_model::state(m.empty_model, T)\n scalar_model::state(m.scalar_model, T)\n vector_model::state(m.vector_model, T)\n matrix_model::state(m.matrix_model, T)\n end\nend\n\nBase.@kwdef struct NTupleModel{S, V, Nv} <: OneLayerModel\n scalar_model::S = ScalarModel()\n vector_model::V = VectorModel{Nv}()\nend\n\nfunction NTupleModel(\n Nv;\n scalar_model::S = ScalarModel(),\n vector_model::V = VectorModel{Nv}(),\n) where {S, V}\n return NTupleModel{S, V, Nv}(scalar_model, vector_model)\nend\n\nfunction state(m::NTupleModel, T)\n @vars begin\n scalar_model::state(m.scalar_model, T)\n vector_model::state(m.vector_model, T)\n end\nend\n\nstate(m::NTuple{N, NTupleModel}, FT) where {N} =\n Tuple{ntuple(i -> state(m[i], FT), N)...}\n\nstruct NTupleContainingModel{N, NTM, VM, SM} <: TwoLayerModel\n ntuple_model::NTM\n vector_model::VM\n scalar_model::SM\nend\nfunction NTupleContainingModel(\n N,\n Nv;\n ntuple_model = ntuple(i -> NTupleModel(Nv), N),\n vector_model = VectorModel{Nv}(),\n scalar_model = ScalarModel(),\n)\n args = (ntuple_model, vector_model, scalar_model)\n return NTupleContainingModel{N, typeof.(args)...}(args...)\nend\n\nfunction state(m::NTupleContainingModel, T)\n @vars begin\n ntuple_model::state(m.ntuple_model, T)\n vector_model::state(m.vector_model, T)\n scalar_model::state(m.scalar_model, T)\n end\nend\n" }, { "path": "test/Utilities/VariableTemplates/runtests.jl", "content": "module TestVariableTemplates\n\nusing Test\nusing StaticArrays\nusing ClimateMachine.VariableTemplates\n\n@testset \"VariableTemplates\" begin\n include(\"test_base_functionality.jl\")\n include(\"varsindex.jl\")\n include(\"test_complex_models.jl\")\nend\n\nend\n" }, { "path": "test/Utilities/VariableTemplates/runtests_gpu.jl", "content": "module TestVariableTemplatesGPU\n\nusing Test\nusing StaticArrays\nusing ClimateMachine.VariableTemplates\n\n@testset \"VariableTemplates - GPU\" begin\n include(\"test_complex_models_gpu.jl\")\nend\n\nend\n" }, { "path": "test/Utilities/VariableTemplates/test_base_functionality.jl", "content": "using Test\nusing StaticArrays\nusing ClimateMachine.VariableTemplates\n\nstruct TestModel{A, B, C}\n a::A\n b::B\n c::C\nend\n\nstruct SubModelA end\nstruct SubModelB end\nstruct SubModelC{N} end\n\nfunction state(m::TestModel, T)\n @vars begin\n ρ::T\n ρu::SVector{3, T}\n ρe::T\n a::state(m.a, T)\n b::state(m.b, T)\n c::state(m.c, T)\n S::SHermitianCompact{3, T, 6}\n end\nend\n\nstate(m::SubModelA, T) = @vars()\nstate(m::SubModelB, T) = @vars(ρqt::T)\nstate(m::SubModelC{N}, T) where {N} = @vars(ρk::SVector{N, T})\n\nmodel = TestModel(SubModelA(), SubModelB(), SubModelC{5}())\n\nst = state(model, Float64)\n\n@test varsize(st) == 17\n@test varsize(typeof(())) == 0\n\nv = Vars{st}(zeros(MVector{varsize(st), Float64}))\ng = Grad{st}(zeros(MMatrix{3, varsize(st), Float64}))\n\n@test v.ρ === 0.0\n@test v.ρu === SVector(0.0, 0.0, 0.0)\nv.ρu = SVector(1, 2, 3)\n@test v.ρu === SVector(1.0, 2.0, 3.0)\n\n@test v.b.ρqt === 0.0\nv.b.ρqt = 12.0\n@test v.b.ρqt === 12.0\n\n@test v.S === zeros(SHermitianCompact{3, Float64, 6})\nv.S = SHermitianCompact{3, Float64, 6}(1, 2, 3, 2, 3, 4, 3, 4, 5)\n@test v.S[1, 1] === 1.0\n@test v.S[1, 3] === 3.0\n@test v.S[3, 1] === 3.0\n@test v.S[3, 3] === 5.0\n\nv.S = ones(SMatrix{3, 3, Int64})\n@test v.S[1, 1] === 1.0\n@test v.S[1, 3] === 1.0\n@test v.S[3, 1] === 1.0\n@test v.S[3, 3] === 1.0\n\n@test propertynames(v.a) == ()\n@test propertynames(g.a) == ()\n\n@test g.ρu == zeros(SMatrix{3, 3, Float64})\ng.ρu = SMatrix{3, 3}(1:9)\n@test g.ρu == SMatrix{3, 3, Float64}(1:9)\n\n@test size(v.c.ρk) == (5,)\n@test size(g.c.ρk) == (3, 5)\n\nsv = similar(v)\n@test typeof(sv) == typeof(v)\n@test size(parent(sv)) == size(parent(v))\n\nsg = similar(g)\n@test typeof(sg) == typeof(g)\n@test size(parent(sg)) == size(parent(g))\n\n@test flattenednames(st) == [\n \"ρ\",\n \"ρu[1]\",\n \"ρu[2]\",\n \"ρu[3]\",\n \"ρe\",\n \"b.ρqt\",\n \"c.ρk[1]\",\n \"c.ρk[2]\",\n \"c.ρk[3]\",\n \"c.ρk[4]\",\n \"c.ρk[5]\",\n \"S[1,1]\",\n \"S[2,1]\",\n \"S[3,1]\",\n \"S[2,2]\",\n \"S[3,2]\",\n \"S[3,3]\",\n]\n" }, { "path": "test/Utilities/VariableTemplates/test_complex_models.jl", "content": "using Test\nusing StaticArrays\nusing Printf\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.VariableTemplates: wrap_val\nimport ClimateMachine.VariableTemplates\nVT = VariableTemplates\n\n@testset \"Test complex models\" begin\n include(\"complex_models.jl\")\n\n FT = Float32\n\n # test getproperty\n m = ScalarModel()\n st = state(m, FT)\n vs = varsize(st)\n a_global = collect(1:vs)\n v = Vars{st}(a_global)\n @test v.x == FT(1)\n\n Nv = 4\n m = VectorModel{Nv}()\n st = state(m, FT)\n vs = varsize(st)\n a_global = collect(1:vs)\n v = Vars{st}(a_global)\n @test v.x == SVector{Nv, FT}(1:Nv)\n\n N = 3\n M = 6\n m = MatrixModel{N, M}()\n st = state(m, FT)\n vs = varsize(st)\n a_global = collect(1:vs)\n v = Vars{st}(a_global)\n @test v.x == SHermitianCompact{N, FT, M}(collect(1:(1 + M - 1)))\n\n Nv = 3\n N = 3\n M = 6\n m = CompositModel(Nv, N, M)\n st = state(m, FT)\n vs = varsize(st)\n a_global = collect(1:vs)\n v = Vars{st}(a_global)\n\n scalar_model = v.scalar_model\n @test v.scalar_model.x == FT(1)\n\n vector_model = v.vector_model\n @test v.vector_model.x == SVector{Nv, FT}([2, 3, 4])\n\n matrix_model = v.matrix_model\n @test v.matrix_model.x ==\n SHermitianCompact{N, FT, M}(collect(5:(5 + M - 1)))\n\n Nv = 3\n N = 5\n m = NTupleContainingModel(N, Nv)\n st = state(m, FT)\n vs = varsize(st)\n a_global = collect(1:vs)\n v = Vars{st}(a_global)\n\n offset = (Nv + 1) * N\n @test v.vector_model.x == SVector{Nv, FT}(collect(1:Nv) .+ offset)\n @test v.scalar_model.x == FT(1 + Nv) + offset\n\n # Make sure we bounds error for NTupleModel's:\n @test_throws BoundsError m.ntuple_model[0]\n @test_throws BoundsError m.ntuple_model[N + 1]\n\n unval(::Val{i}) where {i} = i\n @unroll_map(N) do i\n @test m.ntuple_model[i] isa NTupleModel\n @test m.ntuple_model[i].scalar_model isa ScalarModel\n @test v.ntuple_model[i].scalar_model.x ==\n FT(unval(i)) + (Nv) * (unval(i) - 1)\n @test v.vector_model.x == SVector{Nv, FT}(1:Nv) .+ offset\n @test v.scalar_model.x == FT(Nv + 1) + offset\n end\n\n @test vuntuple(x -> x, 5) == ntuple(i -> Val(i), Val(5))\n\n # test flattenednames\n fn = flattenednames(st)\n j = 1\n for i in 1:N\n @test fn[j] === \"ntuple_model[$i].scalar_model.x\"\n j += 1\n for k in 1:Nv\n @test fn[j] === \"ntuple_model[$i].vector_model.x[$k]\"\n j += 1\n end\n end\n for k in 1:Nv\n @test fn[j] === \"vector_model.x[$k]\"\n j += 1\n end\n @test fn[j] === \"scalar_model.x\"\n\n # flattened_tup_chain - empty/generic cases\n struct Foo end\n @test flattened_tup_chain(NamedTuple{(), Tuple{}}) == ()\n @test flattened_tup_chain(Foo, RetainArr()) == ((Symbol(),),)\n @test flattened_tup_chain(Foo, FlattenArr()) == ((Symbol(),),)\n\n # flattened_tup_chain - SHermitianCompact\n Nv, M = 3, 6\n A = SHermitianCompact{Nv, FT, M}(collect(1:(1 + M - 1)))\n\n ftc = flattened_tup_chain(typeof(A), FlattenArr())\n @test ftc == ntuple(i -> (Symbol(), i), M)\n\n ftc = flattened_tup_chain(typeof(A), RetainArr())\n @test ftc == ((Symbol(),),)\n\n # flattened_tup_chain - Retain arrays\n\n ftc = flattened_tup_chain(st, RetainArr())\n j = 1\n for i in 1:N\n @test ftc[j] === (:ntuple_model, i, :scalar_model, :x)\n j += 1\n @test ftc[j] === (:ntuple_model, i, :vector_model, :x)\n j += 1\n end\n @test ftc[j] === (:vector_model, :x)\n j += 1\n @test ftc[j] === (:scalar_model, :x)\n\n # test varsindex\n ntuple(N) do i\n i_val = Val(i)\n i_sm = varsindex(st, :ntuple_model, i_val, :scalar_model, :x)\n i_vm = varsindex(st, :ntuple_model, i_val, :vector_model, :x)\n nt_offset = (Nv + 1) - 1\n\n i_sm_correct = (i + nt_offset * (i - 1)):(i + nt_offset * (i - 1))\n @test i_sm == i_sm_correct\n\n offset = 1\n i_start = i + nt_offset * (i - 1) + offset\n i_vm_correct = (i_start):(i_start + Nv - 1)\n @test i_vm == i_vm_correct\n end\n\n # test that getproperty matches varsindex\n ntuple(N) do i\n i_ϕ = varsindex(st, wrap_val.(ftc[i])...)\n ϕ = getproperty(v, wrap_val.(ftc[i]))\n @test all(parent(v)[i_ϕ] .≈ ϕ)\n end\n\n # test getproperty with tup-chain\n @unroll_map(N) do i\n @test v.scalar_model.x == getproperty(v, (:scalar_model, :x))\n @test v.vector_model.x == getproperty(v, (:vector_model, :x))\n @test v.ntuple_model[i] == getproperty(v, (:ntuple_model, i))\n @test v.ntuple_model[i].scalar_model ==\n getproperty(v, (:ntuple_model, i, :scalar_model))\n @test v.ntuple_model[i].scalar_model.x ==\n getproperty(v, (:ntuple_model, i, :scalar_model, :x))\n end\n\n # Test converting to flattened NamedTuple\n fnt = flattened_named_tuple(v, RetainArr())\n @test fnt.ntuple_model_1_scalar_model_x == 1.0f0\n @test fnt.ntuple_model_1_vector_model_x == Float32[2.0, 3.0, 4.0]\n @test fnt.ntuple_model_2_scalar_model_x == 5.0f0\n @test fnt.ntuple_model_2_vector_model_x == Float32[6.0, 7.0, 8.0]\n @test fnt.ntuple_model_3_scalar_model_x == 9.0f0\n @test fnt.ntuple_model_3_vector_model_x == Float32[10.0, 11.0, 12.0]\n @test fnt.ntuple_model_4_scalar_model_x == 13.0f0\n @test fnt.ntuple_model_4_vector_model_x == Float32[14.0, 15.0, 16.0]\n @test fnt.ntuple_model_5_scalar_model_x == 17.0f0\n @test fnt.ntuple_model_5_vector_model_x == Float32[18.0, 19.0, 20.0]\n @test fnt.vector_model_x == Float32[21.0, 22.0, 23.0]\n @test fnt.scalar_model_x == 24.0f0\n\n # flattened_tup_chain - Flatten arrays\n\n ftc = flattened_tup_chain(st, FlattenArr())\n j = 1\n for i in 1:N\n @test ftc[j] === (:ntuple_model, i, :scalar_model, :x)\n j += 1\n for k in 1:Nv\n @test ftc[j] === (:ntuple_model, i, :vector_model, :x, k)\n j += 1\n end\n end\n for i in 1:Nv\n @test ftc[j] === (:vector_model, :x, i)\n j += 1\n end\n @test ftc[j] === (:scalar_model, :x)\n\n # test varsindex (flatten arrays)\n ntuple(N) do i\n i_val = Val(i)\n i_sm = varsindex(st, :ntuple_model, i_val, :scalar_model, :x)\n nt_offset = (Nv + 1) - 1\n\n i_sm_correct = (i + nt_offset * (i - 1)):(i + nt_offset * (i - 1))\n @test i_sm == i_sm_correct\n\n for j in 1:Nv\n i_vm =\n varsindex(st, :ntuple_model, i_val, :vector_model, :x, Val(j))\n offset = 1\n i_start = i + nt_offset * (i - 1) + offset\n i_vm_correct = i_start + j - 1\n @test i_vm == i_vm_correct:i_vm_correct\n end\n end\n\n # test that getproperty matches varsindex\n ntuple(N) do i\n i_ϕ = varsindex(st, wrap_val.(ftc[i])...)\n ϕ = getproperty(v, wrap_val.(ftc[i]))\n @test all(parent(v)[i_ϕ] .≈ ϕ)\n end\n\n # test getproperty with tup-chain\n for k in 1:Nv\n @test v.vector_model.x[k] == getproperty(v, (:vector_model, :x, Val(k)))\n end\n\n # test getindex with Val\n @test getindex((1, 2), Val(1)) == 1\n @test getindex((1, 2), Val(2)) == 2\n @test getindex(SVector(1, 2), Val(1)) == 1\n @test getindex(SVector(1, 2), Val(2)) == 2\n\n nt = (; a = ((; x = 1), (; x = 2)))\n fnt = VT.flattened_tuple(FlattenArr(), nt)\n vg = Grad{typeof(nt)}(zeros(MMatrix{3, length(fnt), FT}))\n parent(vg)[1, :] .= fnt\n parent(vg)[2, :] .= fnt\n parent(vg)[3, :] .= fnt\n for i in 1:2\n @test getindex(vg.a, Val(i)).x[1] == i\n @test getindex(vg.a, Val(i)).x[2] == i\n @test getindex(vg.a, Val(i)).x[3] == i\n end\n\n # getpropertyorindex\n @test VT.getpropertyorindex((1, 2), Val(1)) == 1\n @test VT.getpropertyorindex((1, 2), Val(2)) == 2\n @test VT.getpropertyorindex([1, 2], Val(1)) == 1\n @test VT.getpropertyorindex([1, 2], Val(2)) == 2\n @test VT.getpropertyorindex(v, :scalar_model) == v.scalar_model\n for i in 1:N\n @test VT.getpropertyorindex(v.ntuple_model, Val(i)) ==\n v.ntuple_model[Val(i)]\n @test VT.getpropertyorindex(v.ntuple_model, (Val(i),)) ==\n v.ntuple_model[Val(i)]\n @test getindex(v.ntuple_model, (Val(i),)) ==\n VT.getpropertyorindex(v.ntuple_model, (Val(i),))\n end\n\n # Test converting to flattened NamedTuple\n fnt = flattened_named_tuple(v, FlattenArr())\n @test fnt.ntuple_model_1_scalar_model_x == 1.0f0\n @test fnt.ntuple_model_1_vector_model_x_1 == 2.0\n @test fnt.ntuple_model_1_vector_model_x_2 == 3.0\n @test fnt.ntuple_model_1_vector_model_x_3 == 4.0\n @test fnt.ntuple_model_2_scalar_model_x == 5.0f0\n @test fnt.ntuple_model_2_vector_model_x_1 == 6.0\n @test fnt.ntuple_model_2_vector_model_x_2 == 7.0\n @test fnt.ntuple_model_2_vector_model_x_3 == 8.0\n @test fnt.ntuple_model_3_scalar_model_x == 9.0f0\n @test fnt.ntuple_model_3_vector_model_x_1 == 10.0\n @test fnt.ntuple_model_3_vector_model_x_2 == 11.0\n @test fnt.ntuple_model_3_vector_model_x_3 == 12.0\n @test fnt.ntuple_model_4_scalar_model_x == 13.0f0\n @test fnt.ntuple_model_4_vector_model_x_1 == 14.0\n @test fnt.ntuple_model_4_vector_model_x_2 == 15.0\n @test fnt.ntuple_model_4_vector_model_x_3 == 16.0\n @test fnt.ntuple_model_5_scalar_model_x == 17.0f0\n @test fnt.ntuple_model_5_vector_model_x_1 == 18.0\n @test fnt.ntuple_model_5_vector_model_x_2 == 19.0\n @test fnt.ntuple_model_5_vector_model_x_3 == 20.0\n @test fnt.vector_model_x_1 == 21.0\n @test fnt.vector_model_x_2 == 22.0\n @test fnt.vector_model_x_3 == 23.0\n @test fnt.scalar_model_x == 24.0f0\n\n struct Foo end\n nt = (;\n nest = (;\n v = SVector(1, 2, 3),\n nt = (;\n shc = SHermitianCompact{3, FT, 6}(collect(1:6)),\n f = FT(1.0),\n ),\n d = SDiagonal(collect(1:3)...),\n tt = (Foo(), Foo()),\n t = Foo(),\n ),\n )\n # Test flattened_tuple:\n\n @test VT.flattened_tuple(RetainArr(), NamedTuple()) == ()\n @test VT.flattened_tuple(FlattenArr(), NamedTuple()) == ()\n @test VT.flattened_tuple(RetainArr(), Tuple(NamedTuple())) == ()\n @test VT.flattened_tuple(FlattenArr(), Tuple(NamedTuple())) == ()\n\n ft = FlattenArr()\n @test VT.flattened_tuple(ft, nt.nest.nt.f) == (1.0f0,)\n @test VT.flattened_tuple(ft, nt.nest.nt) ==\n (1.0f0, 2.0f0, 3.0f0, 4.0f0, 5.0f0, 6.0f0, 1.0f0)\n @test VT.flattened_tuple(ft, nt.nest.d) == (1, 2, 3)\n @test VT.flattened_tuple(ft, nt.nest.t) == (Foo(),)\n @test VT.flattened_tuple(ft, nt.nest.tt) == (Foo(), Foo())\n\n ft = RetainArr()\n @test VT.flattened_tuple(ft, nt.nest.nt.f) == (1.0f0,)\n @test VT.flattened_tuple(ft, nt.nest.nt)[1] == nt.nest.nt.shc.lowertriangle\n @test VT.flattened_tuple(ft, nt.nest.nt)[2] == 1.0f0\n @test VT.flattened_tuple(ft, nt.nest.d) == (nt.nest.d.diag,)\n @test VT.flattened_tuple(ft, nt.nest.t) == (Foo(),)\n @test VT.flattened_tuple(ft, nt.nest.tt) == (Foo(), Foo())\n\n # Test flattened_named_tuple for NamedTuples\n fnt = flattened_named_tuple(nt, FlattenArr())\n @test fnt.nest_v_1 == 1\n @test fnt.nest_v_2 == 2\n @test fnt.nest_v_3 == 3\n @test fnt.nest_nt_shc_1 == 1.0\n @test fnt.nest_nt_shc_2 == 2.0\n @test fnt.nest_nt_shc_3 == 3.0\n @test fnt.nest_nt_shc_4 == 4.0\n @test fnt.nest_nt_shc_5 == 5.0\n @test fnt.nest_nt_shc_6 == 6.0\n @test fnt.nest_nt_f == 1.0\n @test fnt.nest_tt_1 == Foo()\n @test fnt.nest_tt_2 == Foo()\n @test fnt.nest_t == Foo()\n\n fnt = flattened_named_tuple(nt, RetainArr())\n @test fnt.nest_v == SVector(1, 2, 3)\n @test fnt.nest_nt_shc == nt.nest.nt.shc.lowertriangle\n @test fnt.nest_nt_f == 1.0\n @test fnt.nest_tt_1 == Foo()\n @test fnt.nest_tt_2 == Foo()\n @test fnt.nest_t == Foo()\n\n # Test that show doesn't break\n sprint(show, v)\n sprint(show, vg)\n\nend\n" }, { "path": "test/Utilities/VariableTemplates/test_complex_models_gpu.jl", "content": "using Test\nusing StaticArrays\nusing ClimateMachine.VariableTemplates\nusing CUDA\nusing KernelAbstractions\nusing KernelAbstractions.Extras: @unroll\n\ninclude(\"complex_models.jl\")\n\nrun_gpu = CUDA.has_cuda_gpu()\nif run_gpu\n CUDA.allowscalar(false)\nelse\n CUDA.allowscalar(true)\nend\n\nget_device() = run_gpu ? CPU() : CUDADevice()\ndevice_array(a, ::CUDADevice) = CuArray(a)\ndevice_array(a, ::CPU) = Array(a)\ndevice_rand(::CUDADevice, args...) = CUDA.rand(args...)\ndevice_rand(::CPU, args...) = rand(args...)\nnumber_states(m) = varsize(state(m, Int))\n\n@kernel function mem_copy_kernel!(\n m::AbstractModel,\n dst::AbstractArray{FT, N},\n src::AbstractArray{FT, N},\n) where {FT, N}\n @uniform begin\n ns = number_states(m)\n vs = state(m, FT)\n local_src = MArray{Tuple{ns}, FT}(undef)\n local_dst = MArray{Tuple{ns}, FT}(undef)\n end\n i = @index(Group, Linear)\n @inbounds begin\n @unroll for s in 1:ns\n local_src[s] = src[s, i]\n end\n mem_copy!(m, Vars{vs}(local_dst), Vars{vs}(local_src))\n @unroll for s in 1:ns\n dst[s, i] = local_dst[s]\n end\n end\nend\n\nfunction mem_copy!(m::ScalarModel, dst::Vars, src::Vars)\n dst.x = src.x\nend\n\nfunction mem_copy!(m::NTupleContainingModel{N}, dst::Vars, src::Vars) where {N}\n dst.vector_model.x = src.vector_model.x\n dst.scalar_model.x = src.scalar_model.x\n\n up = vuntuple(i -> src.ntuple_model[i].scalar_model.x, N)\n up_v = vuntuple(i -> src.ntuple_model[i].vector_model.x, N)\n up_sv = SVector(up...)\n\n @unroll_map(N) do i\n dst.ntuple_model[i].scalar_model.x = up[i] # index into tuple\n dst.ntuple_model[i].scalar_model.x = up_sv[i] # index into SArray\n dst.ntuple_model[i].vector_model.x = up_v[i]\n end\nend\n\n@testset \"ScalarModel\" begin\n FT = Float32\n device = get_device()\n n_elem = 10\n m = ScalarModel()\n ns = number_states(m)\n a_src = Array{FT}(undef, ns, n_elem)\n a_dst = Array{FT}(undef, ns, n_elem)\n d_src = device_array(a_src, device)\n d_dst = device_array(a_dst, device)\n d_src .= device_rand(device, FT, ns, n_elem)\n fill!(d_dst, 0)\n\n work_groups = (1,)\n ndrange = (n_elem,)\n kernel! = mem_copy_kernel!(device, work_groups)\n event = kernel!(m, d_dst, d_src, ndrange = ndrange)\n wait(device, event)\n\n a_src = Array(d_src)\n a_dst = Array(d_dst)\n @test a_src == a_dst\nend\n\n@testset \"NTupleContainingModel\" begin\n FT = Float32\n device = get_device()\n n_elem = 10\n Nv = 4\n N = 3\n m = NTupleContainingModel(N, Nv)\n ns = number_states(m)\n a_src = Array{FT}(undef, ns, n_elem)\n a_dst = Array{FT}(undef, ns, n_elem)\n d_src = device_array(a_src, device)\n d_dst = device_array(a_dst, device)\n d_src .= device_rand(device, FT, ns, n_elem)\n fill!(d_dst, 0)\n\n work_groups = (1,)\n ndrange = (n_elem,)\n kernel! = mem_copy_kernel!(device, work_groups)\n event = kernel!(m, d_dst, d_src, ndrange = ndrange)\n wait(device, event)\n\n a_src = Array(d_src)\n a_dst = Array(d_dst)\n @test a_src == a_dst\nend\n" }, { "path": "test/Utilities/VariableTemplates/varsindex.jl", "content": "using Test, StaticArrays\nusing ClimateMachine.VariableTemplates: varsindex, @vars, varsindices\nusing StaticArrays\n\n@testset \"varsindex\" begin\n struct TestMoistureModel{FT} end\n struct TestAtmosModel{FT}\n moisture::TestMoistureModel{FT}\n end\n\n function vars_state(::TestMoistureModel, FT)\n @vars begin\n ρq_tot::FT\n ρq_x::SVector{5, FT}\n ρq_liq::FT\n ρq_vap::FT\n end\n end\n function vars_state(m::TestAtmosModel, FT)\n @vars begin\n ρ::FT\n ρu::SVector{3, FT}\n ρe::FT\n moisture::vars_state(m.moisture, FT)\n S::SHermitianCompact{3, FT, 6}\n end\n end\n FT = Float64\n m = TestAtmosModel(TestMoistureModel{FT}())\n\n @test 1:1 === varsindex(vars_state(m, FT), :ρ)\n @test 2:4 === varsindex(vars_state(m, FT), :ρu)\n @test 5:5 === varsindex(vars_state(m, FT), :ρe)\n\n # Since moisture is defined recusively this will get all the fields\n moist = varsindex(vars_state(m, FT), :moisture)\n @test 6:13 === moist\n\n # To get the specific ones we can do\n @test 6:6 === varsindex(vars_state(m, FT), :moisture, :ρq_tot)\n @test 7:11 === varsindex(vars_state(m, FT), :moisture, :ρq_x)\n @test 12:12 === varsindex(vars_state(m, FT), :moisture, :ρq_liq)\n @test 13:13 === varsindex(vars_state(m, FT), :moisture, :ρq_vap)\n # or\n @test 6:6 === moist[varsindex(vars_state(m.moisture, FT), :ρq_tot)]\n @test 7:11 === moist[varsindex(vars_state(m.moisture, FT), :ρq_x)]\n @test 12:12 === moist[varsindex(vars_state(m.moisture, FT), :ρq_liq)]\n @test 13:13 === moist[varsindex(vars_state(m.moisture, FT), :ρq_vap)]\n\n @test 14:19 == varsindex(vars_state(m, FT), :S)\n\n @test (1,) === varsindices(vars_state(m, FT), :ρ)\n @test (2, 3, 4) === varsindices(vars_state(m, FT), :ρu)\n @test (1, 5) === varsindices(vars_state(m, FT), :ρ, :ρe)\n @test (12,) === varsindices(vars_state(m, FT), :(moisture.ρq_liq))\n let\n vars = (\"ρe\", \"moisture.ρq_x\", \"moisture.ρq_vap\")\n @test (5, 7, 8, 9, 10, 11, 13) === varsindices(vars_state(m, FT), vars)\n end\nend\n" }, { "path": "test/Utilities/runtests.jl", "content": "using Test, Pkg\n\n@testset \"Utilities\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n for submodule in [\"TicToc\", \"VariableTemplates\", \"SingleStackUtils\"]\n if all_tests ||\n \"$submodule\" in ARGS ||\n \"Utilities/$submodule\" in ARGS ||\n \"Utilities\" in ARGS\n include_test(submodule)\n end\n end\n\nend\n" }, { "path": "test/runtests.jl", "content": "using Test, Pkg\n\nENV[\"DATADEPS_ALWAYS_ACCEPT\"] = true\nENV[\"JULIA_LOG_LEVEL\"] = \"WARN\"\n\nfunction include_test(_module)\n println(\"Starting tests for $_module\")\n t = @elapsed include(joinpath(_module, \"runtests.jl\"))\n println(\"Completed tests for $_module, $(round(Int, t)) seconds elapsed\")\n return nothing\nend\n\n\n@testset \"ClimateMachine\" begin\n all_tests = isempty(ARGS) || \"all\" in ARGS ? true : false\n\n function has_submodule(sm)\n any(ARGS) do a\n a == sm && return true\n first(split(a, '/')) == sm && return true\n return false\n end\n end\n\n for submodule in [\n \"InputOutput\",\n \"Utilities\",\n \"Common\",\n \"Arrays\",\n \"BalanceLaws\",\n \"Atmos\",\n \"Land\",\n \"Numerics\",\n \"Diagnostics\",\n \"Ocean\",\n \"Driver\",\n ]\n if all_tests || has_submodule(submodule) || \"ClimateMachine\" in ARGS\n include_test(submodule)\n end\n end\nend\n" }, { "path": "test/runtests_gpu.jl", "content": "using Test, Pkg, CUDA\n\nENV[\"JULIA_LOG_LEVEL\"] = \"WARN\"\n\n@test CUDA.functional()\n\nfor submodule in [\n \"Arrays\",\n #\"Numerics/Mesh\",\n #\"Numerics/DGMethods\",\n \"Numerics/ODESolvers\",\n]\n println(\"Starting tests for $submodule\")\n t = @elapsed include(joinpath(submodule, \"runtests.jl\"))\n println(\"Completed tests for $submodule, $(round(Int, t)) seconds elapsed\")\nend\n" }, { "path": "test/testhelpers.jl", "content": "using MPI\n\nfunction runmpi(file; ntasks = 1, localhost = false)\n localhostenv = try\n parse(Bool, get(ENV, \"CLIMATEMACHINE_TEST_RUNMPI_LOCALHOST\", \"false\"))\n catch\n false\n end\n # Force mpiexec to exec on the localhost node\n # TODO: MicrosoftMPI mpiexec has issues if mpiexec.exe is in a different \n # folder as the MPI script to run so ignore for now.\n if (localhost || localhostenv) && (MPI.MPI_LIBRARY != MPI.MicrosoftMPI)\n localhostonly = `-host localhost`\n else\n localhostonly = ``\n end\n # by default some mpi runtimes will\n # complain if more resources (processes)\n # are requested than available on the node\n if MPI.MPI_LIBRARY == MPI.OpenMPI\n oversubscribe = `--oversubscribe`\n else\n oversubscribe = ``\n end\n @info \"Running MPI test...\" file ntasks\n # Running this way prevents:\n # Balance Law Solver | No tests\n # since external tests are not returned as passed/fail\n @time @test MPI.mpiexec() do cmd\n Base.run(\n `$cmd $localhostonly $oversubscribe -np $ntasks $(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $file`;\n wait = true,\n )\n true\n end\nend\n" }, { "path": "tutorials/Atmos/agnesi_hs_lin.jl", "content": "# # [Linear HS mountain waves (Topography)](@id EX-LIN_HS-docs)\n#\n# ## Description of experiment\n# 1) Dry linear Hydrostatic Mountain Waves\n# The atmosphere is dry and the flow impinges against a witch of Agnesi mountain of heigh $h_{m}=1$m\n# and base parameter $a=10000m$ and centered on $x_{c} = 120km$ in a 2D domain\n# $\\Omega = 240km \\times 50km$. The mountain is defined as\n#\n# ```math\n# z = \\frac{h_m}{1 + \\frac{x - x_c}{a}}\n# ```\n#\n# The 2D problem is setup in 3D by using 1 element in the y direction.\n# To damp the upward moving gravity waves, a Reyleigh absorbing layer is added at $z = 15000 m$.\n#\n# The initial atmosphere is defined such that it has a stability frequency $N=g/\\sqrt{c_p T_0}$, where\n#\n# $T_0 = 250 K$\n# so that\n#\n# ```math\n# \\theta = \\theta_0 = T_0\n# ```\n#\n# ```math\n# \\pi = 1 + \\frac{g^2}{c_p \\theta_0 N^2}\\left(\\exp\\left(\\frac{-N^2 z}{g} \\right)\\right)\n# ```\n#\n# where $\\theta_0 = T_0 K$.\n#\n# so that\n#\n# ```math\n# ρ = \\frac{p_{sfc}}{R_{gas}\\theta}\\pi^{c_v/R_{gas}}\n# ```\n# and\n# ```math\n# T = \\theta \\pi\n# ```\n#\n# 2) Boundaries\n# - `Impenetrable(FreeSlip())` - Top and bottom: no momentum flux, no mass flux through\n# walls.\n# - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through\n# walls.\n# - Agnesi topography built via meshwarp.\n# - Laterally periodic\n# 3) Domain - 240,000 m (horizontal) x 4000 m (horizontal) x 30,000m (vertical)\n# 4) Resolution - 1000m X 240 m effective resolution\n# 5) Total simulation time - 15,000 s\n# 6) Overrides defaults for\n# - CPU Initialisation\n# - Time integrator\n# - Sources\n#\n#md # !!! note\n#md # This experiment setup assumes that you have installed the\n#md # `ClimateMachine` according to the instructions on the landing page.\n#md # We assume the users' familiarity with the conservative form of the\n#md # equations of motion for a compressible fluid (see the\n#md # [AtmosModel](@ref AtmosModel-docs) page).\n#md #\n#md # The following topics are covered in this example\n#md # - Defining the initial conditions\n#md # - Applying source terms\n#md # - Add an idealized topography defined by a warping function\n#\n# ## Boilerplate (Using Modules)\n#\n# The setup of this problem is taken from Case 6 of [giraldoRestelli2008a](@cite)\n#\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true);\nnothing\n\n# Setting `parse_clargs=true` allows the use of command-line arguments (see API > Driver docs)\n# to control simulation update and output intervals.\n\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nusing CLIMAParameters.Atmos.SubgridScale: C_smag\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# ## [Initial Conditions](@id init)\n#md # !!! note\n#md # The following variables are assigned in the initial condition\n#md # - `state.ρ` = Scalar quantity for initial density profile\n#md # - `state.ρu`= 3-component vector for initial momentum profile\n#md # - `state.energy.ρe`= Scalar quantity for initial total-energy profile\n#md # humidity\nfunction init_agnesi_hs_lin!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n ## Problem float-type\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n ## Unpack constant parameters\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n γ::FT = c_p / c_v\n\n c::FT = c_v / R_gas\n c2::FT = R_gas / c_p\n\n Tiso::FT = 250.0\n θ0::FT = Tiso\n\n ## Calculate the Brunt-Vaisaila frequency for an isothermal field\n Brunt::FT = _grav / sqrt(c_p * Tiso)\n Brunt2::FT = Brunt * Brunt\n g2::FT = _grav * _grav\n\n π_exner::FT = exp(-_grav * z / (c_p * Tiso))\n θ::FT = θ0 * exp(Brunt2 * z / _grav)\n ρ::FT = p0 / (R_gas * θ) * (π_exner)^c\n\n ## Compute perturbed thermodynamic state:\n T = θ * π_exner\n e_int = internal_energy(param_set, T)\n ts = PhaseDry(param_set, e_int, ρ)\n\n ## initial velocity\n u = FT(20.0)\n\n ## State (prognostic) variable assignment\n e_kin = FT(0) # kinetic energy\n e_pot = gravitational_potential(bl.orientation, aux)# potential energy\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts) # total energy\n\n state.ρ = ρ\n state.ρu = SVector{3, FT}(ρ * u, 0, 0)\n state.energy.ρe = ρe_tot\nend\n\n# Define a `setmax` method\nfunction setmax(f, xmax, ymax, zmax)\n function setmaxima(xin, yin, zin)\n return f(xin, yin, zin; xmax = xmax, ymax = ymax, zmax = zmax)\n end\n return setmaxima\nend\n\n# Define a warping function to build an analytic topography:\nfunction warp_agnesi(xin, yin, zin; xmax = 1000.0, ymax = 1000.0, zmax = 1000.0)\n\n FT = eltype(xin)\n\n ac = FT(10000)\n hm = FT(1)\n xc = FT(0.5) * xmax\n zdiff = hm / (FT(1) + ((xin - xc) / ac)^2)\n\n ## Linear relaxation towards domain maximum height\n x, y, z = xin, yin, zin + zdiff * (zmax - zin) / zmax\n return x, y, z\nend\n\n# ## [Model Configuration](@id config-helper)\n# We define a configuration function to assist in prescribing the physical\n# model. The purpose of this is to populate the\n# `AtmosLESConfiguration` with arguments\n# appropriate to the problem being considered.\nfunction config_agnesi_hs_lin(\n ::Type{FT},\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n) where {FT}\n ##\n ## Explicit Rayleigh damping:\n ##\n ## ``\n ## \\tau_s = \\alpha * \\sin\\left(0.5\\pi \\frac{z - z_s}{zmax - z_s} \\right)^2,\n ## ``\n ## where\n ## ``sponge_ampz`` is the wave damping coefficient (1/s)\n ## ``z_s`` is the level where the Rayleigh sponge starts\n ## ``zmax`` is the domain top\n ##\n ## Setup the parameters for the gravity wave absorbing layer\n ## at the top of the domain\n ##\n ## u_relaxation(xvelo, vvelo, wvelo) contains the background velocity values to which\n ## the sponge relaxes the vertically moving wave\n u_relaxation = SVector(FT(20), FT(0), FT(0))\n\n ## Wave damping coefficient (1/s)\n sponge_ampz = FT(0.5)\n\n ## Vertical level where the absorbing layer starts\n z_s = FT(25000.0)\n\n ## Pass the sponge parameters to the sponge calculator\n rayleigh_sponge =\n RayleighSponge{FT}(zmax, z_s, sponge_ampz, u_relaxation, 2)\n\n ## Setup the source terms for this problem:\n source = (Gravity(), rayleigh_sponge)\n\n ## Define the reference state:\n T_virt = FT(250)\n temp_profile_ref = IsothermalProfile(param_set, T_virt)\n ref_state = HydrostaticState(temp_profile_ref)\n nothing # hide\n\n _C_smag = FT(0.21)\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = Vreman(_C_smag),\n moisture = DryModel(),\n tracers = NoTracers(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = init_agnesi_hs_lin!,\n source = source,\n )\n\n config = ClimateMachine.AtmosLESConfiguration(\n \"Agnesi_HS_LINEAR\", # Problem title [String]\n N, # Polynomial order [Int]\n resolution, # (Δx, Δy, Δz) effective resolution [m]\n xmax, # Domain maximum size [m]\n ymax, # Domain maximum size [m]\n zmax, # Domain maximum size [m]\n param_set, # Parameter set.\n init_agnesi_hs_lin!, # Function specifying initial condition\n model = model, # Model type\n meshwarp = setmax(warp_agnesi, xmax, ymax, zmax),\n )\n\n return config\nend\n\n# Define a `main` method (entry point)\nfunction main()\n\n FT = Float64\n\n ## Define the polynomial order and effective grid spacings:\n N = 4\n\n ## Define the domain size and spatial resolution\n Nx = 20\n Ny = 20\n Nz = 20\n xmax = FT(244000)\n ymax = FT(4000)\n zmax = FT(50000)\n Δx = xmax / FT(Nx)\n Δy = ymax / FT(Ny)\n Δz = zmax / FT(Nz)\n resolution = (Δx, Δy, Δz)\n\n t0 = FT(0)\n timeend = FT(150) #FT(hrs * 60 * 60)\n\n ## Define the max Courant for the time time integrator (ode_solver).\n ## The default value is 1.7 for LSRK144:\n CFL = FT(1.5)\n\n ## Assign configurations so they can be passed to the `invoke!` function\n driver_config = config_agnesi_hs_lin(FT, N, resolution, xmax, ymax, zmax)\n\n ## Define the time integrator:\n ## We chose an explicit single-rate LSRK144 for this problem\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFL,\n )\n\n ## Set up the spectral filter to remove the solutions spurious modes\n ## Define the order of the exponential filter: use 32 or 64 for this problem.\n ## The larger the value, the less dissipation you get:\n filterorder = 64\n filter = ExponentialFilter(solver_config.dg.grid, 0, filterorder)\n cbfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n AtmosFilterPerturbations(driver_config.bl),\n solver_config.dg.grid,\n filter,\n state_auxiliary = solver_config.dg.state_auxiliary,\n )\n nothing\n end\n ## End exponential filter\n\n ## Invoke solver (calls `solve!` function for time-integrator),\n ## pass the driver, solver and diagnostic config information.\n result = ClimateMachine.invoke!(\n solver_config;\n user_callbacks = (cbfilter,),\n check_euclidean_distance = true,\n )\n\n ## Check that the solution norm is reasonable.\n @test isapprox(result, FT(1); atol = 1.5e-3)\nend\n\n# Call `main`\nmain()\n" }, { "path": "tutorials/Atmos/agnesi_nh_lin.jl", "content": "# # [Linear NH mountain waves (Topography)](@id EX-LIN_NH-docs)\n#\n# ## Description of experiment\n# 1) Dry linear Non-hydrostatic Mountain Waves\n# This example of a non-linear hydrostatic mountain wave can be classified as an initial value\n# problem.\n#\n# The atmosphere is dry and the flow impinges against a witch of Agnesi mountain of heigh $h_m=1m$\n# and base parameter $a=1000$m and centered in $x_c = 72km$ in a 2D domain\n# $\\Omega = 144km \\times 30 km$. The mountain is defined as\n#\n# ```math\n# z = \\frac{h_m}{1 + \\frac{x - x_c}{a}}\n# ```\n# The 2D problem is setup in 3D by using 1 element in the y direction.\n# To damp the upward moving gravity waves, a Reyleigh absorbing layer is added at $z = 10,000m$.\n#\n# The initial atmosphere is defined such that it has a stability frequency $N=0.01 s^{-1}$, where\n#\n# ```math\n# N^2 = g\\frac{\\rm d \\ln \\theta}{ \\rm dz}\n# ```\n# so that\n#\n# ```math\n# \\theta = \\theta_0 \\exp\\left(\\frac{N^2 z}{g} \\right),\n# ```\n# ```math\n# \\pi = 1 + \\frac{g^2}{c_p \\theta_0 N^2}\\left(\\exp\\left(\\frac{-N^2 z}{g} \\right)\\right)\n# ```\n#\n# where $\\theta_0 = 280K$.\n#\n# so that\n#\n# $ρ = \\frac{p_{sfc}}{R_{gas}\\theta}pi^{c_v/R_{gas}}$\n# and $T = \\theta \\pi$\n#\n# 2) Boundaries\n# - `Impenetrable(FreeSlip())` - Top and bottom: no momentum flux, no mass flux through\n# walls.\n# - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through\n# walls.\n# - Agnesi topography built via meshwarp.\n# - Laterally periodic\n# 3) Domain - 144,000 m (horizontal) x 1360 m (horizontal) x 30,000m (vertical) (infinite domain in y)\n# 4) Resolution - 340 m X 200 m effective resolution\n# 5) Total simulation time - 18,000 s\n# 6) Overrides defaults for\n# - CPU Initialisation\n# - Time integrator\n# - Sources\n#\n#md # !!! note\n#md # This experiment setup assumes that you have installed the\n#md # `ClimateMachine` according to the instructions on the landing page.\n#md # We assume the users' familiarity with the conservative form of the\n#md # equations of motion for a compressible fluid (see the\n#md # [AtmosModel](@ref AtmosModel-docs) page).\n#md #\n#md # The following topics are covered in this example\n#md # - Defining the initial conditions\n#md # - Applying source terms\n#md # - Add an idealized topography defined by a warping function\n#\n# ## Boilerplate (Using Modules)\n#\n# The setup of this problem is taken from Case 6 of [giraldoRestelli2008a](@cite)\n#\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing StaticArrays\nusing Test\n\nusing CLIMAParameters\nusing CLIMAParameters.Atmos.SubgridScale: C_smag\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# ## [Initial Conditions](@id init)\n#md # !!! note\n#md # The following variables are assigned in the initial condition\n#md # - `state.ρ` = Scalar quantity for initial density profile\n#md # - `state.ρu`= 3-component vector for initial momentum profile\n#md # - `state.energy.ρe`= Scalar quantity for initial total-energy profile\n#md # humidity\nfunction init_agnesi_hs_lin!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n ## Problem float-type\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n ## Unpack constant parameters\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n γ::FT = c_p / c_v\n\n c::FT = c_v / R_gas\n c2::FT = R_gas / c_p\n\n ## Define initial thermal field as isothermal\n Tiso::FT = 250.0\n θ0::FT = Tiso\n\n ## Assign a value to the Brunt-Vaisala frquencey:\n Brunt::FT = 0.01\n Brunt2::FT = Brunt * Brunt\n g2::FT = _grav * _grav\n\n π_exner::FT = exp(-_grav * z / (c_p * Tiso))\n θ::FT = θ0 * exp(Brunt2 * z / _grav)\n ρ::FT = p0 / (R_gas * θ) * (π_exner)^c\n\n ## Compute perturbed thermodynamic state:\n T = θ * π_exner\n e_int = internal_energy(param_set, T)\n ts = PhaseDry(param_set, e_int, ρ)\n\n ## initial velocity\n u = FT(10.0)\n\n ## State (prognostic) variable assignment\n e_kin = FT(0) # kinetic energy\n e_pot = gravitational_potential(bl.orientation, aux)# potential energy\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts) # total energy\n\n state.ρ = ρ\n state.ρu = SVector{3, FT}(ρ * u, 0, 0)\n state.energy.ρe = ρe_tot\nend\n\nfunction setmax(f, xmax, ymax, zmax)\n function setmaxima(xin, yin, zin)\n return f(xin, yin, zin; xmax = xmax, ymax = ymax, zmax = zmax)\n end\n return setmaxima\nend\n\n# Define a warping function to build an analytic topography:\nfunction warp_agnesi(xin, yin, zin; xmax = 1000.0, ymax = 1000.0, zmax = 1000.0)\n\n FT = eltype(xin)\n\n ac = FT(1000)\n hm = FT(1)\n xc = FT(0.5) * xmax\n zdiff = hm / (FT(1) + ((xin - xc) / ac)^2)\n\n ## Linear relaxation towards domain maximum height\n x, y, z = xin, yin, zin + zdiff * (zmax - zin) / zmax\n return x, y, z\nend\n\n# ## [Model Configuration](@id config-helper)\n# We define a configuration function to assist in prescribing the physical\n# model. The purpose of this is to populate the\n# `AtmosLESConfiguration` with arguments\n# appropriate to the problem being considered.\nfunction config_agnesi_hs_lin(\n ::Type{FT},\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n) where {FT}\n ##\n ## Explicit Rayleigh damping:\n ##\n ## ``\n ## \\tau_s = \\alpha * \\sin\\left(0.5\\pi \\frac{z - z_s}{zmax - z_s} \\right)^2,\n ## ``\n ## where\n ## ``sponge_ampz`` is the wave damping coefficient (1/s)\n ## ``z_s`` is the level where the Rayleigh sponge starts\n ## ``zmax`` is the domain top\n ##\n ## Setup the parameters for the gravity wave absorbing layer\n ## at the top of the domain\n ##\n ## u_relaxation(xvelo, vvelo, wvelo) contains the background velocity values to which\n ## the sponge relaxes the vertically moving wave\n u_relaxation = SVector(FT(10), FT(0), FT(0))\n\n ## Wave damping coefficient (1/s)\n sponge_ampz = FT(0.5)\n\n ## Vertical level where the absorbing layer starts\n z_s = FT(10000.0)\n\n ## Pass the sponge parameters to the sponge calculator\n rayleigh_sponge =\n RayleighSponge{FT}(zmax, z_s, sponge_ampz, u_relaxation, 2)\n\n ## Setup the source terms for this problem:\n source = (Gravity(), rayleigh_sponge)\n\n ## Define the reference state:\n T_virt = FT(280)\n temp_profile_ref = IsothermalProfile(param_set, T_virt)\n ref_state = HydrostaticState(temp_profile_ref)\n nothing # hide\n\n _C_smag = FT(0.0)\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = Vreman(_C_smag),\n moisture = DryModel(),\n tracers = NoTracers(),\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n init_state_prognostic = init_agnesi_hs_lin!,\n source = source,\n )\n\n config = ClimateMachine.AtmosLESConfiguration(\n \"Agnesi_NH_LINEAR\", # Problem title [String]\n N, # Polynomial order [Int]\n resolution, # (Δx, Δy, Δz) effective resolution [m]\n xmax, # Domain maximum size [m]\n ymax, # Domain maximum size [m]\n zmax, # Domain maximum size [m]\n param_set, # Parameter set.\n init_agnesi_hs_lin!, # Function specifying initial condition\n model = model, # Model type\n meshwarp = setmax(warp_agnesi, xmax, ymax, zmax),\n )\n\n return config\nend\n\n# Define a `main` method (entry point)\nfunction main()\n\n FT = Float64\n\n ## Define the polynomial order and effective grid spacings:\n N = 4\n\n ## Define the domain size and spatial resolution\n Nx = 20\n Ny = 20\n Nz = 20\n xmax = FT(144000)\n ymax = FT(4000)\n zmax = FT(30000)\n Δx = xmax / FT(Nx)\n Δy = ymax / FT(Ny)\n Δz = zmax / FT(Nz)\n resolution = (Δx, Δy, Δz)\n\n t0 = FT(0)\n timeend = FT(100)\n\n ## Define the max Courant for the time time integrator (ode_solver).\n ## The default value is 1.7 for LSRK144:\n CFL = FT(1.5)\n\n ## Assign configurations so they can be passed to the `invoke!` function\n driver_config = config_agnesi_hs_lin(FT, N, resolution, xmax, ymax, zmax)\n\n ## Define the time integrator:\n ## We chose an explicit single-rate LSRK144 for this problem\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFL,\n )\n\n ## Set up the spectral filter to remove the solutions spurious modes\n ## Define the order of the exponential filter: use 32 or 64 for this problem.\n ## The larger the value, the less dissipation you get:\n filterorder = 64\n filter = ExponentialFilter(solver_config.dg.grid, 0, filterorder)\n cbfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n AtmosFilterPerturbations(driver_config.bl),\n solver_config.dg.grid,\n filter,\n state_auxiliary = solver_config.dg.state_auxiliary,\n )\n nothing\n end\n ## End exponential filter\n\n ## Invoke solver (calls `solve!` function for time-integrator),\n ## pass the driver, solver and diagnostic config information.\n result = ClimateMachine.invoke!(\n solver_config;\n user_callbacks = (cbfilter,),\n check_euclidean_distance = true,\n )\n\n ## Check that the solution norm is reasonable.\n @test FT(0.9) < result < FT(1)\n ## Some energy is lost from the sponge, so we cannot\n ## expect perfect energy conservation.\nend\n\n# Call `main`\nmain();\n" }, { "path": "tutorials/Atmos/burgers_single_stack.jl", "content": "# # Single stack tutorial based on the 3D Burgers + tracer equations\n\n# This tutorial implements the Burgers equations with a tracer field\n# in a single element stack. The flow is initialized with a horizontally\n# uniform profile of horizontal velocity and uniform initial temperature. The fluid\n# is heated from the bottom surface. Gaussian noise is imposed to the horizontal\n# velocity field at each node at the start of the simulation. The tutorial demonstrates how to\n#\n# * Initialize a [`BalanceLaw`](@ref ClimateMachine.BalanceLaws.BalanceLaw) in a single stack configuration;\n# * Return the horizontal velocity field to a given profile (e.g., large-scale advection);\n# * Remove any horizontal inhomogeneities or noise from the flow.\n#\n# The second and third bullet points are demonstrated imposing Rayleigh friction, horizontal\n# diffusion and 2D divergence damping to the horizontal momentum prognostic equation.\n\n# Equations solved in balance law form:\n\n# ```math\n# \\begin{align}\n# \\frac{∂ ρ}{∂ t} =& - ∇ ⋅ (ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρ\\mathbf{u}}{∂ t} =& - ∇ ⋅ (-μ ∇\\mathbf{u}) - ∇ ⋅ (ρ\\mathbf{u} \\mathbf{u}') - γ[ (ρ\\mathbf{u}-ρ̄\\mathbf{ū}) - (ρ\\mathbf{u}-ρ̄\\mathbf{ū})⋅ẑ ẑ] - ν_d ∇_h (∇_h ⋅ ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρcT}{∂ t} =& - ∇ ⋅ (-α ∇ρcT) - ∇ ⋅ (\\mathbf{u} ρcT)\n# \\end{align}\n# ```\n\n# Boundary conditions:\n# ```math\n# \\begin{align}\n# z_{\\mathrm{min}}: & ρ = 1 \\\\\n# z_{\\mathrm{min}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{min}}: & ρcT = ρc T_{\\mathrm{fixed}} \\\\\n# z_{\\mathrm{max}}: & ρ = 1 \\\\\n# z_{\\mathrm{max}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{max}}: & -α∇ρcT = 0\n# \\end{align}\n# ```\n\n# where\n# - ``t`` is time\n# - ``ρ`` is the density\n# - ``\\mathbf{u}`` is the velocity (vector)\n# - ``\\mathbf{ū}`` is the horizontally averaged velocity (vector)\n# - ``μ`` is the dynamic viscosity tensor\n# - ``γ`` is the Rayleigh friction frequency\n# - ``ν_d`` is the horizontal divergence damping coefficient\n# - ``T`` is the temperature\n# - ``α`` is the thermal diffusivity tensor\n# - ``c`` is the heat capacity\n# - ``ρcT`` is the thermal energy\n\n# Solving these equations is broken down into the following steps:\n# 1) Preliminary configuration\n# 2) PDEs\n# 3) Space discretization\n# 4) Time discretization\n# 5) Solver hooks / callbacks\n# 6) Solve\n# 7) Post-processing\n\n# # Preliminary configuration\n\n# ## [Loading code](@id Loading-code-burgers)\n\n# First, we'll load our pre-requisites\n# - load external packages:\nusing MPI\nusing Distributions\nusing OrderedCollections\nusing Plots\nusing StaticArrays\nusing LinearAlgebra: Diagonal, tr\n\n# - load CLIMAParameters and set up to use it:\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# - load necessary ClimateMachine modules:\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Writers\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, parameter_set\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\n\n# - import necessary ClimateMachine modules: (`import`ing enables us to\n# provide implementations of these structs/methods)\nusing ClimateMachine.Orientations:\n Orientation,\n FlatOrientation,\n init_aux!,\n vertical_unit_vector,\n projection_tangential\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n source!,\n flux_second_order!,\n flux_first_order!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n init_state_auxiliary!,\n init_state_prognostic!,\n BoundaryCondition,\n boundary_conditions,\n boundary_state!\n\n# ## Initialization\n\n# Define the float type (`Float64` or `Float32`)\nconst FT = Float64;\n# Initialize ClimateMachine for CPU.\nClimateMachine.init(; disable_gpu = true);\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\n# Load some helper functions for plotting\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Define the set of Partial Differential Equations (PDEs)\n\n# ## Define the model\n\n# Model parameters can be stored in the particular [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw), in this case, the `BurgersEquation`:\n\nBase.@kwdef struct BurgersEquation{FT, APS, O} <: BalanceLaw\n \"Parameters\"\n param_set::APS\n \"Orientation model\"\n orientation::O\n \"Heat capacity\"\n c::FT = 1\n \"Vertical dynamic viscosity\"\n μv::FT = 1e-4\n \"Horizontal dynamic viscosity\"\n μh::FT = 1\n \"Vertical thermal diffusivity\"\n αv::FT = 1e-2\n \"Horizontal thermal diffusivity\"\n αh::FT = 1\n \"IC Gaussian noise standard deviation\"\n σ::FT = 5e-2\n \"Rayleigh damping\"\n γ::FT = 5\n \"Domain height\"\n zmax::FT = 1\n \"Initial conditions for temperature\"\n initialT::FT = 295.15\n \"Bottom boundary value for temperature (Dirichlet boundary conditions)\"\n T_bottom::FT = 300.0\n \"Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)\"\n flux_top::FT = 0.0\n \"Divergence damping coefficient (horizontal)\"\n νd::FT = 1\nend\n\n# Create an instance of the `BurgersEquation`:\norientation = FlatOrientation()\n\nm = BurgersEquation{FT, typeof(param_set), typeof(orientation)}(\n param_set = param_set,\n orientation = orientation,\n);\n\n# This model dictates the flow control, using [Dynamic Multiple\n# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which\n# kernels are executed.\n\n# ## Define the variables\n\n# All of the methods defined in this section were `import`ed in\n# [Loading code](@ref Loading-code-burgers) to let us provide\n# implementations for our `BurgersEquation` as they will be used\n# by the solver.\n\n# Specify auxiliary variables for `BurgersEquation`\nfunction vars_state(m::BurgersEquation, st::Auxiliary, FT)\n @vars begin\n coord::SVector{3, FT}\n orientation::vars_state(m.orientation, st, FT)\n end\nend\n\n# Specify prognostic variables, the variables solved for in the PDEs, for\n# `BurgersEquation`\nvars_state(::BurgersEquation, ::Prognostic, FT) =\n @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);\n\n# Specify state variables whose gradients are needed for `BurgersEquation`\nvars_state(::BurgersEquation, ::Gradient, FT) =\n @vars(u::SVector{3, FT}, ρcT::FT, ρu::SVector{3, FT});\n\n# Specify gradient variables for `BurgersEquation`\nvars_state(::BurgersEquation, ::GradientFlux, FT) = @vars(\n μ∇u::SMatrix{3, 3, FT, 9},\n α∇ρcT::SVector{3, FT},\n νd∇D::SMatrix{3, 3, FT, 9}\n);\n\n# ## Define the compute kernels\n\n# Specify the initial values in `aux::Vars`, which are available in\n# `init_state_prognostic!`. Note that\n# - this method is only called at `t=0`.\n# - `aux.coord` is available here because we've specified `coord` in `vars_state(m, aux, FT)`.\nfunction nodal_init_state_auxiliary!(\n m::BurgersEquation,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.coord = geom.coord\nend;\n\n# `init_aux!` initializes the auxiliary gravitational potential field needed for vertical projections\nfunction init_state_auxiliary!(\n m::BurgersEquation,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_aux!(m, m.orientation, state_auxiliary, grid, direction)\n\n init_state_auxiliary!(\n m,\n nodal_init_state_auxiliary!,\n state_auxiliary,\n grid,\n direction,\n )\nend;\n\n# Specify the initial values in `state::Vars`. Note that\n# - this method is only called at `t=0`.\n# - `state.ρ`, `state.ρu` and`state.ρcT` are available here because we've specified `ρ`, `ρu` and `ρcT` in `vars_state(m, state, FT)`.\nfunction init_state_prognostic!(\n m::BurgersEquation,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n z = aux.coord[3]\n ε1 = rand(Normal(0, m.σ))\n ε2 = rand(Normal(0, m.σ))\n state.ρ = 1\n ρu = 1 - 4 * (z - m.zmax / 2)^2 + ε1\n ρv = 1 - 4 * (z - m.zmax / 2)^2 + ε2\n ρw = 0\n state.ρu = SVector(ρu, ρv, ρw)\n\n state.ρcT = state.ρ * m.c * m.initialT\nend;\n\n# The remaining methods, defined in this section, are called at every\n# time-step in the solver by the [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw) framework.\n\n# Since we have second-order fluxes, we must tell `ClimateMachine` to compute\n# the gradient of `ρcT`, `u` and `ρu`. Here, we specify how `ρcT`, `u` and `ρu` are computed. Note that\n# e.g. `transform.ρcT` is available here because we've specified `ρcT` in `vars_state(m, ::Gradient, FT)`.\nfunction compute_gradient_argument!(\n m::BurgersEquation,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρcT = state.ρcT\n transform.u = state.ρu / state.ρ\n transform.ρu = state.ρu\nend;\n\n# Specify where in `diffusive::Vars` to store the computed gradient from\n# `compute_gradient_argument!`. Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - `∇transform.u` is available here because we've specified `u` in `vars_state(m, ::Gradient, FT)`.\n# - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor.\n# - The `divergence` may be computed from the trace of tensor `∇ρu`.\nfunction compute_gradient_flux!(\n m::BurgersEquation{FT},\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n param_set = parameter_set(m)\n ∇ρu = ∇transform.ρu\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n divergence = tr(∇ρu) - ẑ' * ∇ρu * ẑ\n diffusive.α∇ρcT = Diagonal(SVector(m.αh, m.αh, m.αv)) * ∇transform.ρcT\n diffusive.μ∇u = Diagonal(SVector(m.μh, m.μh, m.μv)) * ∇transform.u\n diffusive.νd∇D =\n Diagonal(SVector(m.νd, m.νd, FT(0))) *\n Diagonal(SVector(divergence, divergence, FT(0)))\nend;\n\n# Introduce Rayleigh friction towards a target profile as a source.\n# Note that:\n# - Rayleigh damping is only applied in the horizontal using the `projection_tangential` method.\nfunction source!(\n m::BurgersEquation{FT},\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n args...,\n) where {FT}\n param_set = parameter_set(m)\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n z = aux.coord[3]\n ρ̄ū =\n state.ρ * SVector{3, FT}(\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.0,\n )\n ρu_p = state.ρu - ρ̄ū\n source.ρu -=\n m.γ * projection_tangential(m.orientation, param_set, aux, ρu_p)\nend;\n\n# Compute advective flux.\n# Note that:\n# - `state.ρu` is available here because we've specified `ρu` in `vars_state(m, state, FT)`.\nfunction flux_first_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n flux.ρ = state.ρu\n\n u = state.ρu / state.ρ\n flux.ρu = state.ρu * u'\n flux.ρcT = u * state.ρcT\nend;\n\n# Compute diffusive flux (e.g. ``F(μ, \\mathbf{u}, t) = -μ∇\\mathbf{u}`` in the original PDE).\n# Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - The divergence gradient can be written as a diffusive flux using a divergence diagonal tensor.\nfunction flux_second_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρcT -= diffusive.α∇ρcT\n flux.ρu -= diffusive.μ∇u\n flux.ρu -= diffusive.νd∇D\nend;\n\n# ### Boundary conditions\n\n# Second-order terms in our equations, ``∇⋅(G)`` where ``G = μ∇\\mathbf{u}``, are\n# internally reformulated to first-order unknowns.\n# Boundary conditions must be specified for all unknowns, both first-order and\n# second-order unknowns which have been reformulated.\n\nstruct TopBC <: BoundaryCondition end;\nstruct BottomBC <: BoundaryCondition end;\nboundary_conditions(::BurgersEquation) = (BottomBC(), TopBC());\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` (first order unknowns)\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\nend;\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` are specified here for\n# second-order unknowns\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n diff⁺.α∇ρcT = -n⁻ * m.flux_top\nend;\n\n# # Spatial discretization\n\n# Prescribe polynomial order of basis functions in finite elements\nN_poly = 5;\n\n# Specify the number of vertical elements\nnelem_vert = 10;\n\n# Specify the domain height\nzmax = m.zmax;\n\n# Establish a `ClimateMachine` single stack configuration\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"BurgersEquation\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n# # Time discretization\n\n# Specify simulation time (SI units)\nt0 = FT(0);\ntimeend = FT(1);\n\n# We'll define the time-step based on the Fourier\n# number and the [Courant number](https://en.wikipedia.org/wiki/Courant–Friedrichs–Lewy_condition)\n# of the flow\nΔ = min_node_distance(driver_config.grid)\n\ngiven_Fourier = FT(0.5);\nFourier_bound = given_Fourier * Δ^2 / max(m.αh, m.μh, m.νd);\nCourant_bound = FT(0.5) * Δ;\ndt = min(Fourier_bound, Courant_bound)\n\n# # Configure a `ClimateMachine` solver.\n\n# This initializes the state vector and allocates memory for the solution in\n# space (`dg` has the model `m`, which describes the PDEs as well as the\n# function used for initialization). This additionally initializes the ODE\n# solver, by default an explicit Low-Storage\n# [Runge-Kutta](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods)\n# method.\n\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\n# ## Inspect the initial conditions for a single nodal stack\n\n# Let's export plots of the initial state\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nz_scale = 100 # convert from meters to cm\nz_key = \"z\"\nz_label = \"z [cm]\"\nz = get_z(driver_config.grid; z_scale = z_scale)\nstate_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\n# Create an array to store the solution:\nstate_data = Dict[state_vars] # store initial condition at ``t=0``\ntime_data = FT[0] # store time data\n\n# Generate plots of initial conditions for the southwest nodal stack\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρcT\",),\n joinpath(output_dir, \"initial_condition_T_nodal.png\");\n xlabel = \"ρcT at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_u_nodal.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[2]\",),\n joinpath(output_dir, \"initial_condition_v_nodal.png\");\n xlabel = \"ρv at southwest node\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_T_nodal.png)\n# ![](initial_condition_u_nodal.png)\n\n# ## Inspect the initial conditions for the horizontal averages\n\n# Horizontal statistics of variables\n\nstate_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\nstate_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\ndata_avg = Dict[state_vars_avg]\ndata_var = Dict[state_vars_var]\n\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_avg_u.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_variance_u.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_avg_u.png)\n# ![](initial_condition_variance_u.png)\n\n# # Solver hooks / callbacks\n\n# Define the number of outputs from `t0` to `timeend`\nconst n_outputs = 5;\nconst every_x_simulation_time = timeend / n_outputs;\n\n# Create a dictionary for `z` coordinate (and convert to cm) NCDatasets IO:\ndims = OrderedDict(z_key => collect(z));\n\n# Create dictionaries to store outputs:\ndata_var = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_var[1] = state_vars_var\n\ndata_avg = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_avg[1] = state_vars_avg\n\ndata_nodal = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_nodal[1] = state_vars\n\n# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which\n# allow users to inject code to be executed at specified intervals. In this\n# callback, the state variables are collected, combined into a single\n# `OrderedDict` and written to a NetCDF file (for each output step).\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n state_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n i = 1,\n j = 1,\n )\n push!(data_var, state_vars_var)\n push!(data_avg, state_vars_avg)\n push!(data_nodal, state_vars)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Solve\n\n# This is the main `ClimateMachine` solver invocation. While users do not have\n# access to the time-stepping loop, code may be injected via `user_callbacks`,\n# which is a `Tuple` of [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n# # Post-processing\n\n# Our solution has now been calculated and exported to NetCDF files in\n# `output_dir`.\n\n# Let's plot the horizontal statistics of `ρu` and `ρcT`, as well as the evolution of\n# `ρu` for the southwest nodal stack:\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_avg.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\"),\n joinpath(output_dir, \"variance_vs_time_u.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρcT\"),\n joinpath(output_dir, \"solution_vs_time_T_avg.png\");\n xlabel = \"Horizontal mean of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρcT\"),\n joinpath(output_dir, \"variance_vs_time_T.png\");\n xlabel = \"Horizontal variance of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_nodal,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_nodal.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\n# ![](solution_vs_time_u_avg.png)\n# ![](variance_vs_time_u.png)\n# ![](solution_vs_time_T_avg.png)\n# ![](variance_vs_time_T.png)\n# ![](solution_vs_time_u_nodal.png)\n\n# Rayleigh friction returns the horizontal velocity to the objective\n# profile on the timescale of the simulation (1 second), since `γ`∼1. The horizontal viscosity\n# and 2D divergence damping act to reduce the horizontal variance over the same timescale.\n# The initial Gaussian noise is propagated to the temperature field through advection.\n# The horizontal diffusivity acts to reduce this `ρcT` variance in time, although in a longer\n# timescale.\n\n# To run this file, and\n# inspect the solution, include this tutorial in the Julia REPL\n# with:\n\n# ```julia\n# include(joinpath(\"tutorials\", \"Atmos\", \"burgers_single_stack.jl\"))\n# ```\n" }, { "path": "tutorials/Atmos/burgers_single_stack_bjfnk.jl", "content": "# # Single stack with HEVI solver tutorial based on the 3D Burgers + tracer equations\n\n# This tutorial implements the Burgers equations with a tracer field\n# in a single element stack. The flow is initialized with a horizontally\n# uniform profile of horizontal velocity and uniform initial temperature. The fluid\n# is heated from the bottom surface. Gaussian noise is imposed to the horizontal\n# velocity field at each node at the start of the simulation. The tutorial demonstrates how to\n#\n# * Initialize a [`BalanceLaw`](@ref ClimateMachine.BalanceLaws.BalanceLaw) in a single stack configuration;\n# * Return the horizontal velocity field to a given profile (e.g., large-scale advection);\n# * Remove any horizontal inhomogeneities or noise from the flow.\n# * Use horizontal explicit vertical implicit (HEVI) solver in the Single stack setup\n#\n# The second and third bullet points are demonstrated imposing Rayleigh friction, horizontal\n# diffusion and 2D divergence damping to the horizontal momentum prognostic equation.\n\n# Equations solved in balance law form:\n\n# ```math\n# \\begin{align}\n# \\frac{∂ ρ}{∂ t} =& - ∇ ⋅ (ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρ\\mathbf{u}}{∂ t} =& - ∇ ⋅ (-μ ∇\\mathbf{u}) - ∇ ⋅ (ρ\\mathbf{u} \\mathbf{u}') - γ[ (ρ\\mathbf{u}-ρ̄\\mathbf{ū}) - (ρ\\mathbf{u}-ρ̄\\mathbf{ū})⋅ẑ ẑ] - ν_d ∇_h (∇_h ⋅ ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρcT}{∂ t} =& - ∇ ⋅ (-α ∇ρcT) - ∇ ⋅ (\\mathbf{u} ρcT)\n# \\end{align}\n# ```\n\n# Boundary conditions:\n# ```math\n# \\begin{align}\n# z_{\\mathrm{min}}: & ρ = 1 \\\\\n# z_{\\mathrm{min}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{min}}: & ρcT = ρc T_{\\mathrm{fixed}} \\\\\n# z_{\\mathrm{max}}: & ρ = 1 \\\\\n# z_{\\mathrm{max}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{max}}: & -α∇ρcT = 0\n# \\end{align}\n# ```\n\n# where\n# - ``t`` is time\n# - ``ρ`` is the density\n# - ``\\mathbf{u}`` is the velocity (vector)\n# - ``\\mathbf{ū}`` is the horizontally averaged velocity (vector)\n# - ``μ`` is the dynamic viscosity tensor\n# - ``γ`` is the Rayleigh friction frequency\n# - ``ν_d`` is the horizontal divergence damping coefficient\n# - ``T`` is the temperature\n# - ``α`` is the thermal diffusivity tensor\n# - ``c`` is the heat capacity\n# - ``ρcT`` is the thermal energy\n\n# Solving these equations is broken down into the following steps:\n# 1) Preliminary configuration\n# 2) PDEs\n# 3) Space discretization\n# 4) Time discretization\n# 5) Solver hooks / callbacks\n# 6) Solve\n# 7) Post-processing\n\n# # Preliminary configuration\n\n# ## [Loading code](@id Loading-code-burgers-bjfnk)\n\n# First, we'll load our pre-requisites\n# - load external packages:\nusing MPI\nusing Distributions\nusing OrderedCollections\nusing Plots\nusing StaticArrays\nusing LinearAlgebra: Diagonal, tr\n\n# - load CLIMAParameters and set up to use it:\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# - load necessary ClimateMachine modules:\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Writers\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, parameter_set\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\n\n# - import necessary ClimateMachine modules: (`import`ing enables us to\n# provide implementations of these structs/methods)\nusing ClimateMachine.Orientations:\n Orientation,\n FlatOrientation,\n init_aux!,\n vertical_unit_vector,\n projection_tangential\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n source!,\n flux_second_order!,\n flux_first_order!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n init_state_auxiliary!,\n init_state_prognostic!,\n BoundaryCondition,\n boundary_conditions,\n boundary_state!\n\n# ## Initialization\n\n# Define the float type (`Float64` or `Float32`)\nconst FT = Float64;\n# Initialize ClimateMachine for CPU.\nClimateMachine.init(; disable_gpu = true);\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\n# Load some helper functions for plotting\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Define the set of Partial Differential Equations (PDEs)\n\n# ## Define the model\n\n# Model parameters can be stored in the particular [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw), in this case, the `BurgersEquation`:\n\nBase.@kwdef struct BurgersEquation{FT, APS, O} <: BalanceLaw\n \"Parameters\"\n param_set::APS\n \"Orientation model\"\n orientation::O\n \"Heat capacity\"\n c::FT = 1\n \"Vertical dynamic viscosity\"\n μv::FT = 1e-4\n \"Horizontal dynamic viscosity\"\n μh::FT = 1\n \"Vertical thermal diffusivity\"\n αv::FT = 1e-2\n \"Horizontal thermal diffusivity\"\n αh::FT = 1\n \"IC Gaussian noise standard deviation\"\n σ::FT = 5e-2\n \"Rayleigh damping\"\n γ::FT = 5\n \"Domain height\"\n zmax::FT = 1\n \"Initial conditions for temperature\"\n initialT::FT = 295.15\n \"Bottom boundary value for temperature (Dirichlet boundary conditions)\"\n T_bottom::FT = 300.0\n \"Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)\"\n flux_top::FT = 0.0\n \"Divergence damping coefficient (horizontal)\"\n νd::FT = 1\nend\n\n# Create an instance of the `BurgersEquation`:\norientation = FlatOrientation()\n\nm = BurgersEquation{FT, typeof(param_set), typeof(orientation)}(\n param_set = param_set,\n orientation = orientation,\n);\n\n# This model dictates the flow control, using [Dynamic Multiple\n# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which\n# kernels are executed.\n\n# ## Define the variables\n\n# All of the methods defined in this section were `import`ed in\n# [Loading code](@ref Loading-code-burgers) to let us provide\n# implementations for our `BurgersEquation` as they will be used\n# by the solver.\n\n# Specify auxiliary variables for `BurgersEquation`\nfunction vars_state(m::BurgersEquation, st::Auxiliary, FT)\n @vars begin\n coord::SVector{3, FT}\n orientation::vars_state(m.orientation, st, FT)\n end\nend\n\n# Specify prognostic variables, the variables solved for in the PDEs, for\n# `BurgersEquation`\nvars_state(::BurgersEquation, ::Prognostic, FT) =\n @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);\n\n# Specify state variables whose gradients are needed for `BurgersEquation`\nvars_state(::BurgersEquation, ::Gradient, FT) =\n @vars(u::SVector{3, FT}, ρcT::FT, ρu::SVector{3, FT});\n\n# Specify gradient variables for `BurgersEquation`\nvars_state(::BurgersEquation, ::GradientFlux, FT) = @vars(\n μ∇u::SMatrix{3, 3, FT, 9},\n α∇ρcT::SVector{3, FT},\n νd∇D::SMatrix{3, 3, FT, 9}\n);\n\n# ## Define the compute kernels\n\n# Specify the initial values in `aux::Vars`, which are available in\n# `init_state_prognostic!`. Note that\n# - this method is only called at `t=0`.\n# - `aux.coord` is available here because we've specified `coord` in `vars_state(m, aux, FT)`.\nfunction nodal_init_state_auxiliary!(\n m::BurgersEquation,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.coord = geom.coord\nend;\n\n# `init_aux!` initializes the auxiliary gravitational potential field needed for vertical projections\nfunction init_state_auxiliary!(\n m::BurgersEquation,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_aux!(m, m.orientation, state_auxiliary, grid, direction)\n\n init_state_auxiliary!(\n m,\n nodal_init_state_auxiliary!,\n state_auxiliary,\n grid,\n direction,\n )\nend;\n\n# Specify the initial values in `state::Vars`. Note that\n# - this method is only called at `t=0`.\n# - `state.ρ`, `state.ρu` and`state.ρcT` are available here because we've specified `ρ`, `ρu` and `ρcT` in `vars_state(m, state, FT)`.\nfunction init_state_prognostic!(\n m::BurgersEquation,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n z = aux.coord[3]\n ε1 = rand(Normal(0, m.σ))\n ε2 = rand(Normal(0, m.σ))\n state.ρ = 1\n ρu = 1 - 4 * (z - m.zmax / 2)^2 + ε1\n ρv = 1 - 4 * (z - m.zmax / 2)^2 + ε2\n ρw = 0\n state.ρu = SVector(ρu, ρv, ρw)\n\n state.ρcT = state.ρ * m.c * m.initialT\nend;\n\n# The remaining methods, defined in this section, are called at every\n# time-step in the solver by the [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw) framework.\n\n# Since we have second-order fluxes, we must tell `ClimateMachine` to compute\n# the gradient of `ρcT`, `u` and `ρu`. Here, we specify how `ρcT`, `u` and `ρu` are computed. Note that\n# e.g. `transform.ρcT` is available here because we've specified `ρcT` in `vars_state(m, ::Gradient, FT)`.\nfunction compute_gradient_argument!(\n m::BurgersEquation,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρcT = state.ρcT\n transform.u = state.ρu / state.ρ\n transform.ρu = state.ρu\nend;\n\n# Specify where in `diffusive::Vars` to store the computed gradient from\n# `compute_gradient_argument!`. Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - `∇transform.u` is available here because we've specified `u` in `vars_state(m, ::Gradient, FT)`.\n# - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor.\n# - The `divergence` may be computed from the trace of tensor `∇ρu`.\nfunction compute_gradient_flux!(\n m::BurgersEquation{FT},\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n param_set = parameter_set(m)\n ∇ρu = ∇transform.ρu\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n divergence = tr(∇ρu) - ẑ' * ∇ρu * ẑ\n diffusive.α∇ρcT = Diagonal(SVector(m.αh, m.αh, m.αv)) * ∇transform.ρcT\n diffusive.μ∇u = Diagonal(SVector(m.μh, m.μh, m.μv)) * ∇transform.u\n diffusive.νd∇D =\n Diagonal(SVector(m.νd, m.νd, FT(0))) *\n Diagonal(SVector(divergence, divergence, FT(0)))\nend;\n\n# Introduce Rayleigh friction towards a target profile as a source.\n# Note that:\n# - Rayleigh damping is only applied in the horizontal using the `projection_tangential` method.\nfunction source!(\n m::BurgersEquation{FT},\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n args...,\n) where {FT}\n param_set = parameter_set(m)\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n z = aux.coord[3]\n ρ̄ū =\n state.ρ * SVector{3, FT}(\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.0,\n )\n ρu_p = state.ρu - ρ̄ū\n source.ρu -=\n m.γ * projection_tangential(m.orientation, param_set, aux, ρu_p)\nend;\n\n# Compute advective flux.\n# Note that:\n# - `state.ρu` is available here because we've specified `ρu` in `vars_state(m, state, FT)`.\nfunction flux_first_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n flux.ρ = state.ρu\n\n u = state.ρu / state.ρ\n flux.ρu = state.ρu * u'\n flux.ρcT = u * state.ρcT\nend;\n\n# Compute diffusive flux (e.g. ``F(μ, \\mathbf{u}, t) = -μ∇\\mathbf{u}`` in the original PDE).\n# Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - The divergence gradient can be written as a diffusive flux using a divergence diagonal tensor.\nfunction flux_second_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρcT -= diffusive.α∇ρcT\n flux.ρu -= diffusive.μ∇u\n flux.ρu -= diffusive.νd∇D\nend;\n\n# ### Boundary conditions\n\n# Second-order terms in our equations, ``∇⋅(G)`` where ``G = μ∇\\mathbf{u}``, are\n# internally reformulated to first-order unknowns.\n# Boundary conditions must be specified for all unknowns, both first-order and\n# second-order unknowns which have been reformulated.\n\nstruct TopBC <: BoundaryCondition end;\nstruct BottomBC <: BoundaryCondition end;\nboundary_conditions(::BurgersEquation) = (BottomBC(), TopBC());\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` (first order unknowns)\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\nend;\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` are specified here for\n# second-order unknowns\n\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\n\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n diff⁺.α∇ρcT = -n⁻ * m.flux_top\nend;\n\n# # Spatial discretization\n\n# Prescribe polynomial order of basis functions in finite elements\nN_poly = 5;\n\n# Specify the number of vertical elements\nnelem_vert = 10;\n\n# Specify the domain height\nzmax = m.zmax;\n\n# # Temporal discretization\n\n# This initializes the ODE\n# solver, the horizontal explicit vertical implicit scheme\n# with ARK2GiraldoKellyConstantinescu method.\n\node_solver_type = ClimateMachine.HEVISolverType(\n FT;\n solver_method = ARK2GiraldoKellyConstantinescu,\n linear_max_subspace_size = Int(30),\n linear_atol = FT(-1.0),\n linear_rtol = FT(1e-5),\n nonlinear_max_iterations = Int(10),\n nonlinear_rtol = FT(1e-4),\n nonlinear_ϵ = FT(1.e-10),\n preconditioner_update_freq = Int(50),\n)\n\n\n# Establish a `ClimateMachine` single stack configuration\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"BurgersEquation\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n# # Time discretization\n\n# Specify simulation time (SI units)\nt0 = FT(0);\ntimeend = FT(1);\n\n# We'll define the time-step based on the Fourier\n# number and the [Courant number](https://en.wikipedia.org/wiki/Courant–Friedrichs–Lewy_condition)\n# of the flow\nΔ = min_node_distance(driver_config.grid)\n\ngiven_Fourier = FT(0.5);\nFourier_bound = given_Fourier * Δ^2 / max(m.αh, m.μh, m.νd);\nCourant_bound = FT(0.5) * Δ;\n\n# We define the time step 50 times larger than that defined by CFL law\ndt = FT(50.0) * min(Fourier_bound, Courant_bound)\n\n# # Configure a `ClimateMachine` solver.\n\n# This initializes the state vector and allocates memory for the solution in\n# space (`dg` has the model `m`, which describes the PDEs as well as the\n# function used for initialization). \n\n\nsolver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n ode_solver_type = ode_solver_type,\n);\n\n\n# ## Inspect the initial conditions for a single nodal stack\n\n# Let's export plots of the initial state\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nz_scale = 100 # convert from meters to cm\nz_key = \"z\"\nz_label = \"z [cm]\"\nz = get_z(driver_config.grid; z_scale = z_scale)\nstate_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\n# Create an array to store the solution:\nstate_data = Dict[state_vars] # store initial condition at ``t=0``\ntime_data = FT[0] # store time data\n\n# Generate plots of initial conditions for the southwest nodal stack\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρcT\",),\n joinpath(output_dir, \"initial_condition_T_nodal_bjfnk.png\");\n xlabel = \"ρcT at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_u_nodal_bjfnk.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[2]\",),\n joinpath(output_dir, \"initial_condition_v_nodal_bjfnk.png\");\n xlabel = \"ρv at southwest node\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_T_nodal_bjfnk.png)\n# ![](initial_condition_u_nodal_bjfnk.png)\n\n# ## Inspect the initial conditions for the horizontal averages\n\n# Horizontal statistics of variables\n\nstate_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\nstate_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\ndata_avg = Dict[state_vars_avg]\ndata_var = Dict[state_vars_var]\n\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_avg_u_bjfnk.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_variance_u_bjfnk.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_avg_u_bjfnk.png)\n# ![](initial_condition_variance_u_bjfnk.png)\n\n# # Solver hooks / callbacks\n\n# Define the number of outputs from `t0` to `timeend`\nconst n_outputs = 5;\nconst every_x_simulation_time = timeend / n_outputs;\n\n# Create a dictionary for `z` coordinate (and convert to cm) NCDatasets IO:\ndims = OrderedDict(z_key => collect(z));\n\n# Create dictionaries to store outputs:\ndata_var = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_var[1] = state_vars_var\n\ndata_avg = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_avg[1] = state_vars_avg\n\ndata_nodal = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_nodal[1] = state_vars\n\n# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which\n# allow users to inject code to be executed at specified intervals. In this\n# callback, the state variables are collected, combined into a single\n# `OrderedDict` and written to a NetCDF file (for each output step `step`).\nstep = [0];\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n state_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n i = 1,\n j = 1,\n )\n step[1] += 1\n push!(data_var, state_vars_var)\n push!(data_avg, state_vars_avg)\n push!(data_nodal, state_vars)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Solve\n\n# This is the main `ClimateMachine` solver invocation. While users do not have\n# access to the time-stepping loop, code may be injected via `user_callbacks`,\n# which is a `Tuple` of [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n# # Post-processing\n\n# Our solution has now been calculated and exported to NetCDF files in\n# `output_dir`.\n\n# Let's plot the horizontal statistics of `ρu` and `ρcT`, as well as the evolution of\n# `ρu` for the southwest nodal stack:\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_avg_bjfnk.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\"),\n joinpath(output_dir, \"variance_vs_time_u_bjfnk.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρcT\"),\n joinpath(output_dir, \"solution_vs_time_T_avg_bjfnk.png\");\n xlabel = \"Horizontal mean of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρcT\"),\n joinpath(output_dir, \"variance_vs_time_T_bjfnk.png\");\n xlabel = \"Horizontal variance of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_nodal,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_nodal_bjfnk.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\n# ![](solution_vs_time_u_avg_bjfnk.png)\n# ![](variance_vs_time_u_bjfnk.png)\n# ![](solution_vs_time_T_avg_bjfnk.png)\n# ![](variance_vs_time_T_bjfnk.png)\n# ![](solution_vs_time_u_nodal_bjfnk.png)\n\n# Rayleigh friction returns the horizontal velocity to the objective\n# profile on the timescale of the simulation (1 second), since `γ`∼1. The horizontal viscosity\n# and 2D divergence damping act to reduce the horizontal variance over the same timescale.\n# The initial Gaussian noise is propagated to the temperature field through advection.\n# The horizontal diffusivity acts to reduce this `ρcT` variance in time, although in a longer\n# timescale.\n\n# To run this file, and\n# inspect the solution, include this tutorial in the Julia REPL\n# with:\n\n# ```julia\n# include(joinpath(\"tutorials\", \"Atmos\", \"burgers_single_stack.jl\"))\n# ```\n" }, { "path": "tutorials/Atmos/burgers_single_stack_fvm.jl", "content": "# # Finite volume single stack tutorial based on the 3D Burgers + tracer equations\n\n# This tutorial implements the Burgers equations with a tracer field\n# in a single element stack. The flow is initialized with a horizontally\n# uniform profile of horizontal velocity and uniform initial temperature. The fluid\n# is heated from the bottom surface. Gaussian noise is imposed to the horizontal\n# velocity field at each node at the start of the simulation. The tutorial demonstrates how to\n#\n# * Initialize a [`BalanceLaw`](@ref ClimateMachine.BalanceLaws.BalanceLaw) in a single stack configuration;\n# * Return the horizontal velocity field to a given profile (e.g., large-scale advection);\n# * Remove any horizontal inhomogeneities or noise from the flow.\n#\n# The second and third bullet points are demonstrated imposing Rayleigh friction, horizontal\n# diffusion and 2D divergence damping to the horizontal momentum prognostic equation.\n\n# Equations solved in balance law form:\n\n# ```math\n# \\begin{align}\n# \\frac{∂ ρ}{∂ t} =& - ∇ ⋅ (ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρ\\mathbf{u}}{∂ t} =& - ∇ ⋅ (-μ ∇\\mathbf{u}) - ∇ ⋅ (ρ\\mathbf{u} \\mathbf{u}') - γ[ (ρ\\mathbf{u}-ρ̄\\mathbf{ū}) - (ρ\\mathbf{u}-ρ̄\\mathbf{ū})⋅ẑ ẑ] - ν_d ∇_h (∇_h ⋅ ρ\\mathbf{u}) \\\\\n# \\frac{∂ ρcT}{∂ t} =& - ∇ ⋅ (-α ∇ρcT) - ∇ ⋅ (\\mathbf{u} ρcT)\n# \\end{align}\n# ```\n\n# Boundary conditions:\n# ```math\n# \\begin{align}\n# z_{\\mathrm{min}}: & ρ = 1 \\\\\n# z_{\\mathrm{min}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{min}}: & ρcT = ρc T_{\\mathrm{fixed}} \\\\\n# z_{\\mathrm{max}}: & ρ = 1 \\\\\n# z_{\\mathrm{max}}: & ρ\\mathbf{u} = \\mathbf{0} \\\\\n# z_{\\mathrm{max}}: & -α∇ρcT = 0\n# \\end{align}\n# ```\n\n# where\n# - ``t`` is time\n# - ``ρ`` is the density\n# - ``\\mathbf{u}`` is the velocity (vector)\n# - ``\\mathbf{ū}`` is the horizontally averaged velocity (vector)\n# - ``μ`` is the dynamic viscosity tensor\n# - ``γ`` is the Rayleigh friction frequency\n# - ``ν_d`` is the horizontal divergence damping coefficient\n# - ``T`` is the temperature\n# - ``α`` is the thermal diffusivity tensor\n# - ``c`` is the heat capacity\n# - ``ρcT`` is the thermal energy\n\n# Solving these equations is broken down into the following steps:\n# 1) Preliminary configuration\n# 2) PDEs\n# 3) Space discretization\n# 4) Time discretization\n# 5) Solver hooks / callbacks\n# 6) Solve\n# 7) Post-processing\n\n# # Preliminary configuration\n\n# ## [Loading code](@id Loading-code-burgers-fvm)\n\n# First, we'll load our pre-requisites\n# - load external packages:\nusing MPI\nusing Distributions\nusing OrderedCollections\nusing Plots\nusing StaticArrays\nusing LinearAlgebra: Diagonal, tr\n\n# - load CLIMAParameters and set up to use it:\nusing CLIMAParameters\nusing CLIMAParameters.Planet: grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# - load necessary ClimateMachine modules:\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Writers\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, parameter_set\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nimport ClimateMachine.DGMethods.FVReconstructions: FVLinear\n\n# - import necessary ClimateMachine modules: (`import`ing enables us to\n# provide implementations of these structs/methods)\nusing ClimateMachine.Orientations:\n Orientation,\n NoOrientation,\n FlatOrientation,\n init_aux!,\n vertical_unit_vector,\n projection_tangential\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n source!,\n flux_second_order!,\n flux_first_order!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n init_state_auxiliary!,\n init_state_prognostic!,\n construct_face_auxiliary_state!,\n BoundaryCondition,\n boundary_conditions,\n boundary_state!\n\n# ## Initialization\n\n# Define the float type (`Float64` or `Float32`)\nconst FT = Float64;\n# Initialize ClimateMachine for CPU.\nClimateMachine.init(; disable_gpu = true);\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\n# Load some helper functions for plotting\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Define the set of Partial Differential Equations (PDEs)\n\n# ## Define the model\n\n# Model parameters can be stored in the particular [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw), in this case, the `BurgersEquation`:\n\nBase.@kwdef struct BurgersEquation{FT, APS, O} <: BalanceLaw\n \"Parameters\"\n param_set::APS\n \"Orientation model\"\n orientation::O\n \"Heat capacity\"\n c::FT = 1\n \"Vertical dynamic viscosity\"\n μv::FT = 1e-4\n \"Horizontal dynamic viscosity\"\n μh::FT = 1\n \"Vertical thermal diffusivity\"\n αv::FT = 1e-2\n \"Horizontal thermal diffusivity\"\n αh::FT = 1\n \"IC Gaussian noise standard deviation\"\n σ::FT = 5e-2\n \"Rayleigh damping\"\n γ::FT = 5\n \"Domain height\"\n zmax::FT = 1\n \"Initial conditions for temperature\"\n initialT::FT = 295.15\n \"Bottom boundary value for temperature (Dirichlet boundary conditions)\"\n T_bottom::FT = 300.0\n \"Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)\"\n flux_top::FT = 0.0\n \"Divergence damping coefficient (horizontal)\"\n νd::FT = 1\nend\n\n# Create an instance of the `BurgersEquation`:\norientation = FlatOrientation()\n\nm = BurgersEquation{FT, typeof(param_set), typeof(orientation)}(\n param_set = param_set,\n orientation = orientation,\n);\n\n# This model dictates the flow control, using [Dynamic Multiple\n# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which\n# kernels are executed.\n\n# ## Define the variables\n\n# All of the methods defined in this section were `import`ed in\n# [Loading code](@ref Loading-code-burgers) to let us provide\n# implementations for our `BurgersEquation` as they will be used\n# by the solver.\n\n# Specify auxiliary variables for `BurgersEquation`\nfunction vars_state(m::BurgersEquation, st::Auxiliary, FT)\n @vars begin\n coord::SVector{3, FT}\n orientation::vars_state(m.orientation, st, FT)\n end\nend\n\n# Specify prognostic variables, the variables solved for in the PDEs, for\n# `BurgersEquation`\nvars_state(::BurgersEquation, ::Prognostic, FT) =\n @vars(ρ::FT, ρu::SVector{3, FT}, ρcT::FT);\n\n# Specify state variables whose gradients are needed for `BurgersEquation`\nvars_state(::BurgersEquation, ::Gradient, FT) =\n @vars(u::SVector{3, FT}, ρcT::FT, ρu::SVector{3, FT});\n\n# Specify gradient variables for `BurgersEquation`\nvars_state(::BurgersEquation, ::GradientFlux, FT) = @vars(\n μ∇u::SMatrix{3, 3, FT, 9},\n α∇ρcT::SVector{3, FT},\n νd∇D::SMatrix{3, 3, FT, 9}\n);\n\n# ## Define the compute kernels\n\n# Specify the initial values in `aux::Vars`, which are available in\n# `init_state_prognostic!`. Note that\n# - this method is only called at `t=0`.\n# - `aux.coord` is available here because we've specified `coord` in `vars_state(m, aux, FT)`.\nfunction nodal_init_state_auxiliary!(\n m::BurgersEquation,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.coord = geom.coord\nend;\n\n# `init_aux!` initializes the auxiliary gravitational potential field needed for vertical projections\nfunction init_state_auxiliary!(\n m::BurgersEquation,\n state_auxiliary::MPIStateArray,\n grid,\n direction,\n)\n init_aux!(m, m.orientation, state_auxiliary, grid, direction)\n\n init_state_auxiliary!(\n m,\n nodal_init_state_auxiliary!,\n state_auxiliary,\n grid,\n direction,\n )\nend;\n\n# Specify the initial values in `state::Vars`. Note that\n# - this method is only called at `t=0`.\n# - `state.ρ`, `state.ρu` and`state.ρcT` are available here because we've specified `ρ`, `ρu` and `ρcT` in `vars_state(m, state, FT)`.\nfunction init_state_prognostic!(\n m::BurgersEquation,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n z = aux.coord[3]\n ε1 = rand(Normal(0, m.σ))\n ε2 = rand(Normal(0, m.σ))\n state.ρ = 1\n ρu = 1 - 4 * (z - m.zmax / 2)^2 + ε1\n ρv = 1 - 4 * (z - m.zmax / 2)^2 + ε2\n ρw = 0\n state.ρu = SVector(ρu, ρv, ρw)\n\n state.ρcT = state.ρ * m.c * m.initialT\nend;\n\nfunction construct_face_auxiliary_state!(\n bl::BurgersEquation,\n aux_face::AbstractArray,\n aux_cell::AbstractArray,\n Δz::FT,\n) where {FT <: Real}\n param_set = parameter_set(bl)\n _grav = FT(grav(param_set))\n var_aux = Vars{vars_state(bl, Auxiliary(), FT)}\n aux_face .= aux_cell\n\n if !(bl.orientation isa NoOrientation)\n var_aux(aux_face).orientation.Φ =\n var_aux(aux_cell).orientation.Φ + _grav * Δz / 2\n end\nend\n\n\n# The remaining methods, defined in this section, are called at every\n# time-step in the solver by the [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw) framework.\n\n# Since we have second-order fluxes, we must tell `ClimateMachine` to compute\n# the gradient of `ρcT`, `u` and `ρu`. Here, we specify how `ρcT`, `u` and `ρu` are computed. Note that\n# e.g. `transform.ρcT` is available here because we've specified `ρcT` in `vars_state(m, ::Gradient, FT)`.\nfunction compute_gradient_argument!(\n m::BurgersEquation,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρcT = state.ρcT\n transform.u = state.ρu / state.ρ\n transform.ρu = state.ρu\nend;\n\n# Specify where in `diffusive::Vars` to store the computed gradient from\n# `compute_gradient_argument!`. Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - `∇transform.u` is available here because we've specified `u` in `vars_state(m, ::Gradient, FT)`.\n# - `diffusive.μ∇u` is built using an anisotropic diffusivity tensor.\n# - The `divergence` may be computed from the trace of tensor `∇ρu`.\nfunction compute_gradient_flux!(\n m::BurgersEquation{FT},\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n) where {FT}\n param_set = parameter_set(m)\n ∇ρu = ∇transform.ρu\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n divergence = tr(∇ρu) - ẑ' * ∇ρu * ẑ\n diffusive.α∇ρcT = Diagonal(SVector(m.αh, m.αh, m.αv)) * ∇transform.ρcT\n diffusive.μ∇u = Diagonal(SVector(m.μh, m.μh, m.μv)) * ∇transform.u\n diffusive.νd∇D =\n Diagonal(SVector(m.νd, m.νd, FT(0))) *\n Diagonal(SVector(divergence, divergence, FT(0)))\nend;\n\n# Introduce Rayleigh friction towards a target profile as a source.\n# Note that:\n# - Rayleigh damping is only applied in the horizontal using the `projection_tangential` method.\nfunction source!(\n m::BurgersEquation{FT},\n source::Vars,\n state::Vars,\n diffusive::Vars,\n aux::Vars,\n args...,\n) where {FT}\n param_set = parameter_set(m)\n ẑ = vertical_unit_vector(m.orientation, param_set, aux)\n z = aux.coord[3]\n ρ̄ū =\n state.ρ * SVector{3, FT}(\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.5 - 2 * (z - m.zmax / 2)^2,\n 0.0,\n )\n ρu_p = state.ρu - ρ̄ū\n source.ρu -=\n m.γ * projection_tangential(m.orientation, param_set, aux, ρu_p)\nend;\n\n# Compute advective flux.\n# Note that:\n# - `state.ρu` is available here because we've specified `ρu` in `vars_state(m, state, FT)`.\nfunction flux_first_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n flux.ρ = state.ρu\n\n u = state.ρu / state.ρ\n flux.ρu = state.ρu * u'\n flux.ρcT = u * state.ρcT\nend;\n\n# Compute diffusive flux (e.g. ``F(μ, \\mathbf{u}, t) = -μ∇\\mathbf{u}`` in the original PDE).\n# Note that:\n# - `diffusive.μ∇u` is available here because we've specified `μ∇u` in `vars_state(m, ::GradientFlux, FT)`.\n# - The divergence gradient can be written as a diffusive flux using a divergence diagonal tensor.\nfunction flux_second_order!(\n m::BurgersEquation,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρcT -= diffusive.α∇ρcT\n flux.ρu -= diffusive.μ∇u\n flux.ρu -= diffusive.νd∇D\nend;\n\n# ### Boundary conditions\n\n# Second-order terms in our equations, ``∇⋅(G)`` where ``G = μ∇\\mathbf{u}``, are\n# internally reformulated to first-order unknowns.\n# Boundary conditions must be specified for all unknowns, both first-order and\n# second-order unknowns which have been reformulated.\n\nstruct TopBC <: BoundaryCondition end;\nstruct BottomBC <: BoundaryCondition end;\nboundary_conditions(::BurgersEquation) = (BottomBC(), TopBC());\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` (first order unknowns)\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\nend;\n\n# The boundary conditions for `ρ`, `ρu` and `ρcT` are specified here for\n# second-order unknowns\nfunction boundary_state!(\n nf,\n bc::BottomBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n state⁺.ρcT = state⁺.ρ * m.c * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::TopBC,\n m::BurgersEquation,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n _...,\n)\n state⁺.ρ = 1\n state⁺.ρu = SVector(0, 0, 0)\n diff⁺.α∇ρcT = -n⁻ * m.flux_top\nend;\n\n# # Spatial discretization\n\n# Prescribe polynomial order of basis functions in finite elements\n# The second index 0 indicates that finite volume method is \n# applied in the vertical direction\nN_poly = (1, 0);\n\n# Specify the number of vertical elements\nnelem_vert = 50;\n\n# Specify the domain height\nzmax = m.zmax;\n\n# Establish a `ClimateMachine` single stack configuration\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"BurgersEquation\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n fv_reconstruction = FVLinear(),\n);\n\n# # Time discretization\n\n# Specify simulation time (SI units)\nt0 = FT(0);\ntimeend = FT(1);\n\n# We'll define the time-step based on the Fourier\n# number and the [Courant number](https://en.wikipedia.org/wiki/Courant–Friedrichs–Lewy_condition)\n# of the flow\nΔ = min_node_distance(driver_config.grid)\n\ngiven_Fourier = FT(0.5);\nFourier_bound = given_Fourier * Δ^2 / max(m.αh, m.μh, m.νd);\nCourant_bound = FT(0.5) * Δ;\ndt = min(Fourier_bound, Courant_bound)\n\n# # Configure a `ClimateMachine` solver.\n\n# This initializes the state vector and allocates memory for the solution in\n# space (`dg` has the model `m`, which describes the PDEs as well as the\n# function used for initialization). This additionally initializes the ODE\n# solver, by default an explicit Low-Storage\n# [Runge-Kutta](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods)\n# method.\n\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\n# ## Inspect the initial conditions for a single nodal stack\n\n# Let's export plots of the initial state\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nz_scale = 100 # convert from meters to cm\nz_key = \"z\"\nz_label = \"z [cm]\"\nz = get_z(driver_config.grid; z_scale = z_scale)\nstate_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\n# Create an array to store the solution:\nstate_data = Dict[state_vars] # store initial condition at ``t=0``\ntime_data = FT[0] # store time data\n\n# Generate plots of initial conditions for the southwest nodal stack\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρcT\",),\n joinpath(output_dir, \"initial_condition_T_nodal_fvm.png\");\n xlabel = \"ρcT at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_u_nodal_fvm.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n state_data,\n (\"ρu[2]\",),\n joinpath(output_dir, \"initial_condition_v_nodal_fvm.png\");\n xlabel = \"ρv at southwest node\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_T_nodal_fvm.png)\n# ![](initial_condition_u_nodal_fvm.png)\n\n# ## Inspect the initial conditions for the horizontal averages\n\n# Horizontal statistics of variables\n\nstate_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\nstate_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n);\n\ndata_avg = Dict[state_vars_avg]\ndata_var = Dict[state_vars_var]\n\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_avg_u_fvm.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\",),\n joinpath(output_dir, \"initial_condition_variance_u_fvm.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\n\n# ![](initial_condition_avg_u_fvm.png)\n# ![](initial_condition_variance_u_fvm.png)\n\n# # Solver hooks / callbacks\n\n# Define the number of outputs from `t0` to `timeend`\nconst n_outputs = 5;\nconst every_x_simulation_time = timeend / n_outputs;\n\n# Create a dictionary for `z` coordinate (and convert to cm) NCDatasets IO:\ndims = OrderedDict(z_key => collect(z));\n\n# Create dictionaries to store outputs:\ndata_var = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_var[1] = state_vars_var\n\ndata_avg = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_avg[1] = state_vars_avg\n\ndata_nodal = Dict[Dict([k => Dict() for k in 0:n_outputs]...),]\ndata_nodal[1] = state_vars\n\n# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which\n# allow users to inject code to be executed at specified intervals. In this\n# callback, the state variables are collected, combined into a single\n# `OrderedDict` and written to a NetCDF file (for each output step).\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n state_vars_var = get_horizontal_variance(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars_avg = get_horizontal_mean(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n )\n state_vars = get_vars_from_nodal_stack(\n driver_config.grid,\n solver_config.Q,\n vars_state(m, Prognostic(), FT),\n i = 1,\n j = 1,\n )\n push!(data_var, state_vars_var)\n push!(data_avg, state_vars_avg)\n push!(data_nodal, state_vars)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Solve\n\n# This is the main `ClimateMachine` solver invocation. While users do not have\n# access to the time-stepping loop, code may be injected via `user_callbacks`,\n# which is a `Tuple` of [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,))\n\n# # Post-processing\n\n# Our solution has now been calculated and exported to NetCDF files in\n# `output_dir`.\n\n# Let's plot the horizontal statistics of `ρu` and `ρcT`, as well as the evolution of\n# `ρu` for the southwest nodal stack:\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_avg_fvm.png\");\n xlabel = \"Horizontal mean of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρu[1]\"),\n joinpath(output_dir, \"variance_vs_time_u_fvm.png\");\n xlabel = \"Horizontal variance of ρu\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_avg,\n (\"ρcT\"),\n joinpath(output_dir, \"solution_vs_time_T_avg_fvm.png\");\n xlabel = \"Horizontal mean of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_var,\n (\"ρcT\"),\n joinpath(output_dir, \"variance_vs_time_T_fvm.png\");\n xlabel = \"Horizontal variance of ρcT\",\n ylabel = z_label,\n);\nexport_plot(\n z,\n time_data,\n data_nodal,\n (\"ρu[1]\"),\n joinpath(output_dir, \"solution_vs_time_u_nodal_fvm.png\");\n xlabel = \"ρu at southwest node\",\n ylabel = z_label,\n);\n# ![](solution_vs_time_u_avg_fvm.png)\n# ![](variance_vs_time_u_fvm.png)\n# ![](solution_vs_time_T_avg_fvm.png)\n# ![](variance_vs_time_T_fvm.png)\n# ![](solution_vs_time_u_nodal_fvm.png)\n\n# Rayleigh friction returns the horizontal velocity to the objective\n# profile on the timescale of the simulation (1 second), since `γ`∼1. The horizontal viscosity\n# and 2D divergence damping act to reduce the horizontal variance over the same timescale.\n# The initial Gaussian noise is propagated to the temperature field through advection.\n# The horizontal diffusivity acts to reduce this `ρcT` variance in time, although in a longer\n# timescale.\n\n# To run this file, and\n# inspect the solution, include this tutorial in the Julia REPL\n# with:\n\n# ```julia\n# include(joinpath(\"tutorials\", \"Atmos\", \"burgers_single_stack.jl\"))\n# ```\n" }, { "path": "tutorials/Atmos/densitycurrent.jl", "content": "# # Density Current\n#\n# In this example, we demonstrate the usage of the `ClimateMachine`\n# to solve the density current test by Straka 1993.\n# We solve a flow in a box configuration, which is\n# representative of a large-eddy simulation. Several versions of the problem\n# setup may be found in literature, but the general idea is to examine the\n# vertical ascent of a thermal _bubble_ (we can interpret these as simple\n# representation of convective updrafts).\n#\n# ## Description of experiment\n# The setup described below is such that the simulation reaches completion\n# (timeend = 900 s) in approximately 4 minutes of wall-clock time on 1 GPU\n#\n# 1) Dry Density Current (circular potential temperature perturbation)\n# 2) Boundaries\n# - `Impenetrable(FreeSlip())` - no momentum flux, no mass flux through\n# walls.\n# - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through\n# walls.\n# 3) Domain - 25600m (horizontal) x 10000m (horizontal) x 6400m (vertical)\n# 4) Resolution - 100m effective resolution\n# 5) Total simulation time - 900s\n# 6) Mesh Aspect Ratio (Effective resolution) 1:1\n# 7) Overrides defaults for\n# - CPU Initialisation\n# - Time integrator\n# - Sources\n# - Smagorinsky Coefficient _C_smag\n# 8) Default settings can be found in `src/Driver/.jl`\n\n#md # !!! note\n#md # This experiment setup assumes that you have installed the\n#md # `ClimateMachine` according to the instructions on the landing page.\n#md # We assume the users' familiarity with the conservative form of the\n#md # equations of motion for a compressible fluid\n#md #\n#md # The following topics are covered in this example\n#md # - Package requirements\n#md # - Defining a `model` subtype for the set of conservation equations\n#md # - Defining the initial conditions\n#md # - Applying boundary conditions\n#md # - Applying source terms\n#md # - Choosing a turbulence model\n#md # - Adding tracers to the model\n#md # - Choosing a time-integrator\n#md #\n#md # The following topics are not covered in this example\n#md # - Defining new boundary conditions\n#md # - Defining new turbulence models\n#md # - Building new time-integrators\n#\n# ## Boilerplate (Using Modules)\n#\n# #### [Skip Section](@ref init-dc)\n#\n# Before setting up our experiment, we recognize that we need to import some\n# pre-defined functions from other packages. Julia allows us to use existing\n# modules (variable workspaces), or write our own to do so. Complete\n# documentation for the Julia module system can be found\n# [here](https://docs.julialang.org/en/v1/manual/modules/#).\n\n# We need to use the `ClimateMachine` module! This imports all functions\n# specific to atmospheric and ocean flow modeling. While we do not cover the\n# ins-and-outs of the contents of each of these we provide brief descriptions\n# of the utility of each of the loaded packages.\nusing ClimateMachine\nClimateMachine.init(parse_clargs = true)\n\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\n# - Required so that we inherit the appropriate model types for the large-eddy\n# simulation (LES) and global-circulation-model (GCM) configurations.\nusing ClimateMachine.ConfigTypes\n# - Required so that we may define diagnostics configurations, e.g. choice of\n# file-writer, choice of output variable sets, output-frequency and directory,\nusing ClimateMachine.Diagnostics\n# - Required so that we may define (or utilise existing functions) functions\n# that are `called-back` or executed at frequencies of either timesteps,\n# simulation-time, or wall-clock time.\nusing ClimateMachine.GenericCallbacks\n# - Required so we load the appropriate functions for the time-integration\n# component. Contains ODESolver methods.\nusing ClimateMachine.ODESolvers\n# - Required for utility of spatial filtering functions (e.g. positivity\n# preservation)\nusing ClimateMachine.Mesh.Filters\n# - Required so functions for computation of temperature profiles.\nusing Thermodynamics.TemperatureProfiles\n# - Required so functions for computation of moist thermodynamic quantities and turbulence closures\n# are available.\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\n# - Required so we may access our variable arrays by a sensible naming\n# convention rather than by numerical array indices.\nusing ClimateMachine.VariableTemplates\n# - Required so we may access planet parameters\n# ([CLIMAParameters](https://github.com/CliMA/CLIMAParameters.jl)\n# specific to this problem include the gas constant, specific heats,\n# mean-sea-level pressure, gravity and the Smagorinsky coefficient)\n\n# In ClimateMachine we use `StaticArrays` for our variable arrays.\nusing StaticArrays\n# We also use the `Test` package to help with unit tests and continuous\n# integration systems to design sensible tests for our experiment to ensure new\n# / modified blocks of code don't damage the fidelity of the physics. The test\n# defined within this experiment is not a unit test for a specific\n# subcomponent, but ensures time-integration of the defined problem conditions\n# within a reasonable tolerance. Immediately useful macros and functions from\n# this include `@test` and `@testset` which will allow us to define the testing\n# parameter sets.\nusing Test\n\nusing CLIMAParameters\nusing CLIMAParameters.Atmos.SubgridScale: C_smag\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# ## [Initial Conditions](@id init-dc)\n#md # !!! note\n#md # The following variables are assigned in the initial condition\n#md # - `state.ρ` = Scalar quantity for initial density profile\n#md # - `state.ρu`= 3-component vector for initial momentum profile\n#md # - `state.energy.ρe`= Scalar quantity for initial total-energy profile\n#md # humidity\n#md # - `state.tracers.ρχ` = Vector of four tracers (here, for demonstration\n#md # only; we can interpret these as dye injections for visualisation\n#md # purposes)\nfunction init_densitycurrent!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n ## Problem float-type\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n ## Unpack constant parameters\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n γ::FT = c_p / c_v\n\n ## Define bubble center and background potential temperature\n xc::FT = 0\n yc::FT = 0\n zc::FT = 3000\n rx::FT = 4000\n rz::FT = 2000\n r = sqrt(((x - xc)^2) / rx^2 + ((z - zc)^2) / rz^2)\n\n ## TODO: clean this up, or add convenience function:\n ## This is configured in the reference hydrostatic state\n ref_state = reference_state(bl)\n θ_ref::FT = ref_state.virtual_temperature_profile.T_surface\n Δθ::FT = 0\n θamplitude::FT = -15.0\n\n ## Compute temperature difference over bubble region\n if r <= 1\n Δθ = 0.5 * θamplitude * (1 + cospi(r))\n end\n\n ## Compute perturbed thermodynamic state:\n θ = θ_ref + Δθ ## potential temperature\n π_exner = FT(1) - _grav / (c_p * θ) * z ## exner pressure\n ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas) ## density\n T = θ * π_exner\n e_int = internal_energy(param_set, T)\n ts = PhaseDry(param_set, e_int, ρ)\n ρu = SVector(FT(0), FT(0), FT(0)) ## momentum\n ## State (prognostic) variable assignment\n e_kin = FT(0) ## kinetic energy\n e_pot = gravitational_potential(bl.orientation, aux)## potential energy\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts) ## total energy\n\n ## Assign State Variables\n state.ρ = ρ\n state.ρu = ρu\n state.energy.ρe = ρe_tot\nend\n\n# ## [Model Configuration](@id config-helper)\n# We define a configuration function to assist in prescribing the physical\n# model.\nfunction config_densitycurrent(\n ::Type{FT},\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n) where {FT}\n\n ## The model coefficient for the turbulence closure is defined via the\n ## [CLIMAParameters\n ## package](https://CliMA.github.io/CLIMAParameters.jl/dev/) A reference\n ## state for the linearisation step is also defined.\n T_surface = FT(300)\n T_min_ref = FT(0)\n T_profile = DryAdiabaticProfile{FT}(param_set, T_surface, T_min_ref)\n ref_state = HydrostaticState(T_profile)\n\n ## The fun part! Here we assemble the `AtmosModel`.\n ##md # !!! note\n ##md # Docs on model subcomponent options can be found here:\n ##md # - [`param_set`](https://CliMA.github.io/CLIMAParameters.jl/dev/)\n ##md # - [`turbulence`](@ref Turbulence-Closures-docs)\n ##md # - [`source`](@ref atmos-sources)\n ##md # - [`init_state`](@ref init-dc)\n\n _C_smag = FT(0.21)\n physics = AtmosPhysics{FT}(\n param_set; # Parameter set corresponding to earth parameters\n ref_state = ref_state, # Reference state\n turbulence = Vreman(_C_smag), # Turbulence closure model\n moisture = DryModel(), # Exclude moisture variables\n tracers = NoTracers(), # Tracer model with diffusivity coefficients\n )\n model = AtmosModel{FT}(\n AtmosLESConfigType, # Flow in a box, requires the AtmosLESConfigType\n physics; # Atmos physics\n init_state_prognostic = init_densitycurrent!, # Apply the initial condition\n source = (Gravity(),), # Gravity is the only source term here\n )\n\n ## Finally, we pass a `Problem Name` string, the mesh information, and the\n ## model type to the [`AtmosLESConfiguration`](@ref ClimateMachine.AtmosLESConfiguration) object.\n config = ClimateMachine.AtmosLESConfiguration(\n \"DryDensitycurrent\", # Problem title [String]\n N, # Polynomial order [Int]\n resolution, # (Δx, Δy, Δz) effective resolution [m]\n xmax, # Domain maximum size [m]\n ymax, # Domain maximum size [m]\n zmax, # Domain maximum size [m]\n param_set, # Parameter set.\n init_densitycurrent!, # Function specifying initial condition\n model = model, # Model type\n periodicity = (false, false, false),\n boundary = ((1, 1), (1, 1), (1, 1)), # Set all boundaries to solid walls\n )\n return config\nend\n\n#md # !!! note\n#md # `Keywords` are used to specify some arguments (see appropriate source\n#md # files).\n\nfunction main()\n ## These are essentially arguments passed to the\n ## [`config_densitycurrent`](@ref config-helper) function. For type\n ## consistency we explicitly define the problem floating-precision.\n FT = Float64\n ## We need to specify the polynomial order for the DG discretization,\n ## effective resolution, simulation end-time, the domain bounds, and the\n ## courant-number for the time-integrator. Note how the time-integration\n ## components `solver_config` are distinct from the spatial / model\n ## components in `driver_config`. `init_on_cpu` is a helper keyword argument\n ## that forces problem initialisation on CPU (thereby allowing the use of\n ## random seeds, spline interpolants and other special functions at the\n ## initialisation step.)\n N = 4\n Δx = FT(100)\n Δy = FT(250)\n Δv = FT(100)\n resolution = (Δx, Δy, Δv)\n xmax = FT(25600)\n ymax = FT(1000)\n zmax = FT(6400)\n t0 = FT(0)\n timeend = FT(100)\n CFL = FT(1.5)\n\n ## Assign configurations so they can be passed to the `invoke!` function\n driver_config = config_densitycurrent(FT, N, resolution, xmax, ymax, zmax)\n\n ## Choose an Explicit Single-rate Solver LSRK144 from the existing [ODESolvers](@ref\n ## ODESolvers-docs) options Apply the outer constructor to define the\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFL,\n )\n\n ## Invoke solver (calls `solve!` function for time-integrator), pass the driver, solver and diagnostic config\n ## information.\n result =\n ClimateMachine.invoke!(solver_config; check_euclidean_distance = true)\n\n ## Check that the solution norm is reasonable.\n @test isapprox(result, FT(1); atol = 1.5e-2)\nend\n\n# The experiment definition is now complete. Time to run it.\n# `julia --project=$CLIMA_HOME tutorials/Atmos/densitycurrent.jl --vtk 1smins`\n# to run with VTK output enabled at intervals of 1 simulation minute.\n#\n# ## References\n#\n# - [Straka1993](@cite)\n# - [Carpenter1990](@cite)\n\nmain()\n" }, { "path": "tutorials/Atmos/dry_rayleigh_benard.jl", "content": "# # Dry Rayleigh Benard\n\n# ## Problem description\n#\n# 1) Dry Rayleigh Benard Convection (re-entrant channel configuration)\n# 2) Boundaries - `Sides` : Periodic (Default `bctuple` used to identify bot,top walls)\n# `Top` : Prescribed temperature, no-slip\n# `Bottom`: Prescribed temperature, no-slip\n# 3) Domain - 250m[horizontal] x 250m[horizontal] x 500m[vertical]\n# 4) Timeend - 100s\n# 5) Mesh Aspect Ratio (Effective resolution) 1:1\n# 6) Random values in initial condition (Requires `init_on_cpu=true` argument)\n# 7) Overrides defaults for\n# `C_smag`\n# `Courant_number`\n# `init_on_cpu`\n# `ref_state`\n# `solver_type`\n# `bc`\n# `sources`\n# 8) Default settings can be found in src/Driver/Configurations.jl\n\n# ## Loading code\nusing Distributions\nusing Random\nusing StaticArrays\nusing Test\nusing DocStringExtensions\nusing Printf\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.Mesh.Filters\nusing Thermodynamics: PhaseEquil_pTq, internal_energy\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\nusing CLIMAParameters\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, grav, MSLP\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# Convenience struct for sharing data between kernels\nstruct DryRayleighBenardConvectionDataConfig{FT}\n xmin::FT\n ymin::FT\n zmin::FT\n xmax::FT\n ymax::FT\n zmax::FT\n T_bot::FT\n T_lapse::FT\n T_top::FT\nend\n\n# Define initial condition kernel\nfunction init_problem!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n dc = bl.data_config\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n _R_d::FT = R_d(param_set)\n _cp_d::FT = cp_d(param_set)\n _grav::FT = grav(param_set)\n _cv_d::FT = cv_d(param_set)\n _MSLP::FT = MSLP(param_set)\n\n γ::FT = _cp_d / _cv_d\n δT =\n sinpi(6 * z / (dc.zmax - dc.zmin)) *\n cospi(6 * z / (dc.zmax - dc.zmin)) + rand()\n δw =\n sinpi(6 * z / (dc.zmax - dc.zmin)) *\n cospi(6 * z / (dc.zmax - dc.zmin)) + rand()\n ΔT = _grav / _cv_d * z + δT\n T = dc.T_bot - ΔT\n P = _MSLP * (T / dc.T_bot)^(_grav / _R_d / dc.T_lapse)\n ρ = P / (_R_d * T)\n\n q_tot = FT(0)\n e_pot = gravitational_potential(bl.orientation, aux)\n ts = PhaseEquil_pTq(param_set, P, T, q_tot)\n\n ρu, ρv, ρw = FT(0), FT(0), ρ * δw\n\n e_int = internal_energy(ts)\n e_kin = FT(1 / 2) * δw^2\n\n ρe_tot = ρ * (e_int + e_pot + e_kin)\n state.ρ = ρ\n state.ρu = SVector(ρu, ρv, ρw)\n state.energy.ρe = ρe_tot\n state.moisture.ρq_tot = FT(0)\n ρχ = zero(FT)\n if z <= 100\n ρχ += FT(0.1) * (cospi(z / 2 / 100))^2\n end\n state.tracers.ρχ = SVector{1, FT}(ρχ)\nend\n\n# Define problem configuration kernel\nfunction config_problem(::Type{FT}, N, resolution, xmax, ymax, zmax) where {FT}\n\n ## Boundary conditions\n T_bot = FT(299)\n\n _cp_d::FT = cp_d(param_set)\n _grav::FT = grav(param_set)\n\n T_lapse = FT(_grav / _cp_d)\n T_top = T_bot - T_lapse * zmax\n\n ntracers = 1\n δ_χ = SVector{ntracers, FT}(1)\n\n ## Turbulence\n C_smag = FT(0.23)\n data_config = DryRayleighBenardConvectionDataConfig{FT}(\n 0,\n 0,\n 0,\n xmax,\n ymax,\n zmax,\n T_bot,\n T_lapse,\n FT(T_bot - T_lapse * zmax),\n )\n\n ## Define the physics\n physics = AtmosPhysics{FT}(\n param_set;\n turbulence = Vreman(C_smag),\n tracers = NTracers{ntracers, FT}(δ_χ),\n )\n\n ## Set up the problem\n problem = AtmosProblem(;\n physics = physics,\n boundaryconditions = (\n AtmosBC(\n physics;\n momentum = Impenetrable(NoSlip()),\n energy = PrescribedTemperature((state, aux, t) -> T_bot),\n ),\n AtmosBC(\n physics;\n momentum = Impenetrable(NoSlip()),\n energy = PrescribedTemperature((state, aux, t) -> T_top),\n ),\n ),\n init_state_prognostic = init_problem!,\n )\n\n ## Set up the model\n model = AtmosModel{FT}(\n AtmosLESConfigType,\n physics;\n problem = problem,\n source = (Gravity(),),\n data_config = data_config,\n )\n\n config = ClimateMachine.AtmosLESConfiguration(\n \"DryRayleighBenardConvection\",\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n param_set,\n init_problem!,\n model = model,\n )\n return config\nend\n\n# Define diagnostics configuration kernel\nfunction config_diagnostics(driver_config)\n interval = \"10000steps\"\n dgngrp = setup_atmos_default_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n )\n return ClimateMachine.DiagnosticsConfiguration([dgngrp])\nend\n\n# Define main entry point kernel\nfunction main()\n FT = Float64\n ## DG polynomial order\n N = 4\n ## Domain resolution and size\n Δh = FT(10)\n ## Time integrator setup\n t0 = FT(0)\n ## Courant number\n CFLmax = FT(20)\n timeend = FT(1000)\n xmax, ymax, zmax = FT(250), FT(250), FT(500)\n\n @testset \"DryRayleighBenardTest\" begin\n for Δh in Δh\n Δv = Δh\n resolution = (Δh, Δh, Δv)\n driver_config = config_problem(FT, N, resolution, xmax, ymax, zmax)\n\n ## Set up the time-integrator, using a multirate infinitesimal step\n ## method. The option `splitting_type = ClimateMachine.SlowFastSplitting()`\n ## separates fast-slow modes by splitting away the acoustic waves and\n ## treating them via a sub-stepped explicit method.\n ode_solver_type = ClimateMachine.MISSolverType(;\n splitting_type = ClimateMachine.SlowFastSplitting(),\n mis_method = MIS2,\n fast_method = LSRK144NiegemannDiehlBusch,\n nsubsteps = (40,),\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFLmax,\n )\n dgn_config = config_diagnostics(driver_config)\n ## User defined callbacks (TMAR positivity preserving filter)\n cbtmarfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n (\"moisture.ρq_tot\",),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (cbtmarfilter,),\n check_euclidean_distance = true,\n )\n ## result == engf/eng0\n @test isapprox(result, FT(1); atol = 1.5e-2)\n end\n end\nend\n\n# Run\nmain()\n" }, { "path": "tutorials/Atmos/heldsuarez.jl", "content": "# # Dry atmosphere GCM with Held-Suarez forcing\n#\n# The Held-Suarez setup (Held and Suarez, 1994) is a textbook example for a\n# simplified atmospheric global circulation model configuration which has been\n# used as a benchmark experiment for development of the dynamical cores (i.e.,\n# GCMs without continents, moisture or parametrization schemes of the physics)\n# for atmospheric models. It is forced by a thermal relaxation to a reference\n# state and damped by linear (Rayleigh) friction. This example demonstrates how\n#\n# * to set up a ClimateMachine-Atmos GCM configuration;\n# * to select and save GCM diagnostics output.\n#\n# To begin, we load ClimateMachine and a few miscellaneous useful Julia packages.\nusing Distributions\nusing Random\nusing StaticArrays\nusing UnPack\n\n# ClimateMachine specific modules needed to make this example work (e.g., we will need\n# spectral filters, etc.).\nusing ClimateMachine\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Mesh.Filters\nusing Thermodynamics.TemperatureProfiles\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.BalanceLaws\n\nimport ClimateMachine.BalanceLaws: source, prognostic_vars\n\n# [ClimateMachine parameters](https://github.com/CliMA/CLIMAParameters.jl) are\n# needed to have access to Earth's physical parameters.\nusing CLIMAParameters\nusing CLIMAParameters.Planet: MSLP, R_d, day, grav, cp_d, cv_d, planet_radius\n\n# We need to load the physical parameters for Earth to have an Earth-like simulation :).\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n\n\"\"\"\n HeldSuarezForcingTutorial <: TendencyDef{Source}\n\nDefines a forcing that parametrises radiative and frictional effects using\nNewtonian relaxation and Rayleigh friction, following Held and Suarez (1994)\n\"\"\"\nstruct HeldSuarezForcingTutorial <: TendencyDef{Source} end\n\nprognostic_vars(::HeldSuarezForcingTutorial) = (Momentum(), Energy())\n\nfunction held_suarez_forcing_coefficients(bl, args)\n @unpack state, aux = args\n @unpack ts = args.precomputed\n FT = eltype(state)\n\n ## Parameters\n T_ref = FT(255)\n param_set = parameter_set(bl)\n\n _R_d = FT(R_d(param_set))\n _day = FT(day(param_set))\n _grav = FT(grav(param_set))\n _cp_d = FT(cp_d(param_set))\n _p0 = FT(MSLP(param_set))\n\n ## Held-Suarez parameters\n k_a = FT(1 / (40 * _day))\n k_f = FT(1 / _day)\n k_s = FT(1 / (4 * _day))\n ΔT_y = FT(60)\n Δθ_z = FT(10)\n T_equator = FT(315)\n T_min = FT(200)\n σ_b = FT(7 / 10)\n\n ## Held-Suarez forcing\n φ = latitude(bl, aux)\n p = air_pressure(ts)\n\n ##TODO: replace _p0 with dynamic surface pressure in Δσ calculations to account\n ##for topography, but leave unchanged for calculations of σ involved in T_equil\n σ = p / _p0\n exner_p = σ^(_R_d / _cp_d)\n Δσ = (σ - σ_b) / (1 - σ_b)\n height_factor = max(0, Δσ)\n T_equil = (T_equator - ΔT_y * sin(φ)^2 - Δθ_z * log(σ) * cos(φ)^2) * exner_p\n T_equil = max(T_min, T_equil)\n k_T = k_a + (k_s - k_a) * height_factor * cos(φ)^4\n k_v = k_f * height_factor\n return (k_v = k_v, k_T = k_T, T_equil = T_equil)\nend\n\nfunction source(::Energy, s::HeldSuarezForcingTutorial, atmos, args)\n @unpack state = args\n FT = eltype(state)\n @unpack ts = args.precomputed\n nt = held_suarez_forcing_coefficients(atmos, args)\n param_set = parameter_set(atmos)\n _cv_d = FT(cv_d(param_set))\n @unpack k_T, T_equil = nt\n T = air_temperature(ts)\n return -k_T * state.ρ * _cv_d * (T - T_equil)\nend\n\nfunction source(::Momentum, s::HeldSuarezForcingTutorial, atmos, args)\n @unpack state, aux = args\n nt = held_suarez_forcing_coefficients(atmos, args)\n return -nt.k_v * projection_tangential(atmos, aux, state.ρu)\nend\n\n\n# ## Set initial condition\n# When using ClimateMachine, we need to define a function that sets the initial\n# state of our model run. In our case, we use the reference state of the\n# simulation (defined below) and add a little bit of noise. Note that the\n# initial states includes a zero initial velocity field.\nfunction init_heldsuarez!(problem, balance_law, state, aux, localgeo, time)\n FT = eltype(state)\n\n ## Set initial state to reference state with random perturbation\n rnd = FT(1 + rand(Uniform(-1e-3, 1e-3)))\n state.ρ = aux.ref_state.ρ\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = rnd * aux.ref_state.ρe\nend;\n\n\n# ## Initialize ClimateMachine\n# Before we do anything further, we need to initialize ClimateMachine. Among\n# other things, this will initialize the MPI us.\nClimateMachine.init();\n\n\n# ## Setting the floating-type precision\n# ClimateMachine allows us to run a model with different floating-type\n# precisions, with lower precision we get our results faster, and with higher\n# precision, we may get more accurate results, depending on the questions we\n# are after.\nconst FT = Float64;\n\n\n# ## Setup model configuration\n# Now that we have defined our forcing and initialization functions, and have\n# initialized ClimateMachine, we can set up the model.\n#\n# ## Set up a reference state\n# We start by setting up a reference state. This is simply a vector field that\n# we subtract from the solutions to the governing equations to both improve\n# numerical stability of the implicit time stepper and enable faster model\n# spin-up. The reference state assumes hydrostatic balance and ideal gas law,\n# with a pressure $p_r(z)$ and density $\\rho_r(z)$ that only depend on altitude\n# $z$ and are in hydrostatic balance with each other.\n#\n# In this example, the reference temperature field smoothly transitions from a\n# linearly decaying profile near the surface to a constant temperature profile\n# at the top of the domain.\ntemp_profile_ref = DecayingTemperatureProfile{FT}(param_set)\nref_state = HydrostaticState(temp_profile_ref);\n\n# ## Set up a Rayleigh sponge layer\n# To avoid wave reflection at the top of the domain, the model applies a sponge\n# layer that linearly damps the momentum equations.\ndomain_height = FT(30e3) # height of the computational domain (m)\nz_sponge = FT(12e3) # height at which sponge begins (m)\nα_relax = FT(1 / 60 / 15) # sponge relaxation rate (1/s)\nexponent = FT(2) # sponge exponent for squared-sinusoid profile\nu_relax = SVector(FT(0), FT(0), FT(0)) # relaxation velocity (m/s)\nsponge = RayleighSponge{FT}(domain_height, z_sponge, α_relax, u_relax, exponent);\n\n# ## Set up turbulence models\n# In order to produce a stable simulation, we need to dissipate energy and\n# enstrophy at the smallest scales of the developed flow field. To achieve this\n# we set up diffusive forcing functions.\nc_smag = FT(0.21); # Smagorinsky constant\nτ_hyper = FT(4 * 3600); # hyperdiffusion time scale\nturbulence_model = SmagorinskyLilly(c_smag);\nhyperdiffusion_model = DryBiharmonic(FT(4 * 3600));\n\n\n# ## Instantiate the model\n# The Held Suarez setup was designed to produce an equilibrated state that is\n# comparable to the zonal mean of the Earth’s atmosphere.\nphysics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = turbulence_model,\n hyperdiffusion = hyperdiffusion_model,\n moisture = DryModel(),\n);\n\nmodel = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = init_heldsuarez!,\n source = (Gravity(), Coriolis(), HeldSuarezForcingTutorial(), sponge),\n);\n\n# This concludes the setup of the physical model!\n\n# ## Set up the driver\n# We just need to set up a few parameters that define the resolution of the\n# discontinuous Galerkin method and for how long we want to run our model\n# setup.\npoly_order = 5; ## discontinuous Galerkin polynomial order\nn_horz = 2; ## horizontal element number\nn_vert = 2; ## vertical element number\nresolution = (n_horz, n_vert)\nn_days = 0.1; ## experiment day number\ntimestart = FT(0); ## start time (s)\ntimeend = FT(n_days * day(param_set)); ## end time (s);\n\n\n# The next lines set up the spatial grid.\ndriver_config = ClimateMachine.AtmosGCMConfiguration(\n \"HeldSuarez\",\n poly_order,\n resolution,\n domain_height,\n param_set,\n init_heldsuarez!;\n model = model,\n);\n\n# The next lines set up the time stepper. Since the resolution\n# in the vertical is much finer than in the horizontal,\n# the 'stiff' parts of the PDE will be in the vertical.\n# Setting `splitting_type = HEVISplitting()` will treat\n# vertical acoustic waves implicitly, while all other dynamics\n# are treated explicitly.\node_solver_type = ClimateMachine.IMEXSolverType(\n splitting_type = HEVISplitting(),\n implicit_model = AtmosAcousticGravityLinearModel,\n implicit_solver = ManyColumnLU,\n solver_method = ARK2GiraldoKellyConstantinescu,\n);\n\nsolver_config = ClimateMachine.SolverConfiguration(\n timestart,\n timeend,\n driver_config,\n Courant_number = FT(0.1),\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n CFL_direction = HorizontalDirection(),\n diffdir = HorizontalDirection(),\n);\n\n# ## Set up spectral exponential filter\n# After every completed time step we apply a spectral filter to remove\n# remaining small-scale noise introduced by the numerical procedures. This\n# assures that our run remains stable.\nfilterorder = 10;\nfilter = ExponentialFilter(solver_config.dg.grid, 0, filterorder);\ncbfilter = GenericCallbacks.EveryXSimulationSteps(1) do\n Filters.apply!(\n solver_config.Q,\n AtmosFilterPerturbations(model),\n solver_config.dg.grid,\n filter,\n state_auxiliary = solver_config.dg.state_auxiliary,\n )\n nothing\nend;\n\n# ## Setup diagnostic output\n#\n# Choose frequency and resolution of output, and a diagnostics group (dgngrp)\n# which defines output variables. This needs to be defined in\n# [`Diagnostics`](@ref ClimateMachine.Diagnostics).\ninterval = \"1000steps\";\n_planet_radius = FT(planet_radius(param_set));\ninfo = driver_config.config_info;\nboundaries = [\n FT(-90.0) FT(-180.0) _planet_radius\n FT(90.0) FT(180.0) FT(_planet_radius + info.domain_height)\n];\nresolution = (FT(10), FT(10), FT(1000)); # in (deg, deg, m)\ninterpol = ClimateMachine.InterpolationConfiguration(\n driver_config,\n boundaries,\n resolution,\n);\n\ndgngrps = [\n setup_dump_state_diagnostics(\n AtmosGCMConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n setup_dump_aux_diagnostics(\n AtmosGCMConfigType(),\n interval,\n driver_config.name,\n interpol = interpol,\n ),\n];\ndgn_config = ClimateMachine.DiagnosticsConfiguration(dgngrps);\n\n# ## Run the model\n# Finally, we can run the model using the physical setup and solvers from\n# above. We use the spectral filter in our callbacks after every time step, and\n# collect the diagnostics output.\nresult = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (cbfilter,),\n check_euclidean_distance = true,\n);\n\n\n# ## References\n#\n# - [Held1994](@cite)\n" }, { "path": "tutorials/Atmos/risingbubble.jl", "content": "# # Rising Thermal Bubble\n#\n# In this example, we demonstrate the usage of the `ClimateMachine`\n# [AtmosModel](@ref AtmosModel-docs) machinery to solve the fluid\n# dynamics of a thermal perturbation in a neutrally stratified background state\n# defined by its uniform potential temperature. We solve a flow in a box configuration -\n# this is representative of a large-eddy simulation. Several versions of the problem\n# setup may be found in literature, but the general idea is to examine the\n# vertical ascent of a thermal bubble (we can interpret these as simple\n# representation of convective updrafts).\n#\n# ## Description of experiment\n# 1) Dry Rising Bubble (circular potential temperature perturbation)\n# 2) Boundaries\n# Top and Bottom boundaries:\n# - `Impenetrable(FreeSlip())` - Top and bottom: no momentum flux, no mass flux through\n# walls.\n# - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through\n# walls.\n# Lateral boundaries\n# - Laterally periodic\n# 3) Domain - 2500m (horizontal) x 2500m (horizontal) x 2500m (vertical)\n# 4) Resolution - 50m effective resolution\n# 5) Total simulation time - 1000s\n# 6) Mesh Aspect Ratio (Effective resolution) 1:1\n# 7) Overrides defaults for\n# - CPU Initialisation\n# - Time integrator\n# - Sources\n# - Smagorinsky Coefficient\n\n#md # !!! note\n#md # This experiment setup assumes that you have installed the\n#md # `ClimateMachine` according to the instructions on the landing page.\n#md # We assume the users' familiarity with the conservative form of the\n#md # equations of motion for a compressible fluid (see the\n#md # [AtmosModel](@ref AtmosModel-docs) page).\n#md #\n#md # The following topics are covered in this example\n#md # - Package requirements\n#md # - Defining a `model` subtype for the set of conservation equations\n#md # - Defining the initial conditions\n#md # - Applying source terms\n#md # - Choosing a turbulence model\n#md # - Adding tracers to the model\n#md # - Choosing a time-integrator\n#md # - Choosing diagnostics (output) configurations\n#md #\n#md # The following topics are not covered in this example\n#md # - Defining new boundary conditions\n#md # - Defining new turbulence models\n#md # - Building new time-integrators\n#md # - Adding diagnostic variables (beyond a standard pre-defined list of\n#md # variables)\n\n# ## [Loading code](@id Loading-code-rtb)\n\n# Before setting up our experiment, we recognize that we need to import some\n# pre-defined functions from other packages. Julia allows us to use existing\n# modules (variable workspaces), or write our own to do so. Complete\n# documentation for the Julia module system can be found\n# [here](https://docs.julialang.org/en/v1/manual/modules/#).\n\n# We need to use the `ClimateMachine` module! This imports all functions\n# specific to atmospheric and ocean flow modeling.\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\n# In ClimateMachine we use `StaticArrays` for our variable arrays.\n# We also use the `Test` package to help with unit tests and continuous\n# integration systems to design sensible tests for our experiment to ensure new\n# / modified blocks of code don't damage the fidelity of the physics. The test\n# defined within this experiment is not a unit test for a specific\n# subcomponent, but ensures time-integration of the defined problem conditions\n# within a reasonable tolerance. Immediately useful macros and functions from\n# this include `@test` and `@testset` which will allow us to define the testing\n# parameter sets.\nusing StaticArrays\nusing Test\nusing CLIMAParameters\nusing CLIMAParameters.Atmos.SubgridScale: C_smag\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n\n# ## [Initial Conditions](@id init-rtb)\n# This example demonstrates the use of functions defined\n# in the [`Thermodynamics`](@ref Thermodynamics) package to\n# generate the appropriate initial state for our problem.\n\n#md # !!! note\n#md # The following variables are assigned in the initial condition\n#md # - `state.ρ` = Scalar quantity for initial density profile\n#md # - `state.ρu`= 3-component vector for initial momentum profile\n#md # - `state.energy.ρe`= Scalar quantity for initial total-energy profile\n#md # humidity\n#md # - `state.tracers.ρχ` = Vector of four tracers (here, for demonstration\n#md # only; we can interpret these as dye injections for visualization\n#md # purposes)\nfunction init_risingbubble!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n ## Problem float-type\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n ## Unpack constant parameters\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n γ::FT = c_p / c_v\n\n ## Define bubble center and background potential temperature\n xc::FT = 5000\n yc::FT = 1000\n zc::FT = 2000\n r = sqrt((x - xc)^2 + (z - zc)^2)\n rc::FT = 2000\n θamplitude::FT = 2\n\n ## This is configured in the reference hydrostatic state\n ref_state = reference_state(bl)\n θ_ref::FT = ref_state.virtual_temperature_profile.T_surface\n\n ## Add the thermal perturbation:\n Δθ::FT = 0\n if r <= rc\n Δθ = θamplitude * (1.0 - r / rc)\n end\n\n ## Compute perturbed thermodynamic state:\n θ = θ_ref + Δθ ## potential temperature\n π_exner = FT(1) - _grav / (c_p * θ) * z ## exner pressure\n ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas) ## density\n T = θ * π_exner\n e_int = internal_energy(param_set, T)\n ts = PhaseDry(param_set, e_int, ρ)\n ρu = SVector(FT(0), FT(0), FT(0)) ## momentum\n ## State (prognostic) variable assignment\n e_kin = FT(0) ## kinetic energy\n e_pot = gravitational_potential(bl, aux) ## potential energy\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts) ## total energy\n\n ρχ = FT(0) ## tracer\n\n ## We inject tracers at the initial condition at some specified z coordinates\n if 500 < z <= 550\n ρχ += FT(0.05)\n end\n\n ## We want 4 tracers\n ntracers = 4\n\n ## Define 4 tracers, (arbitrary scaling for this demo problem)\n ρχ = SVector{ntracers, FT}(ρχ, ρχ / 2, ρχ / 3, ρχ / 4)\n\n ## Assign State Variables\n state.ρ = ρ\n state.ρu = ρu\n state.energy.ρe = ρe_tot\n state.tracers.ρχ = ρχ\nend\n\n# ## [Model Configuration](@id config-helper)\n# We define a configuration function to assist in prescribing the physical\n# model. The purpose of this is to populate the\n# `ClimateMachine.AtmosLESConfiguration` with arguments\n# appropriate to the problem being considered.\nfunction config_risingbubble(\n ::Type{FT},\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n) where {FT}\n ## Since we want four tracers, we specify this and include the appropriate\n ## diffusivity scaling coefficients (normally these would be physically\n ## informed but for this demonstration we use integers corresponding to the\n ## tracer index identifier)\n ntracers = 4\n δ_χ = SVector{ntracers, FT}(1, 2, 3, 4)\n ## To assemble `AtmosModel` with no tracers, set `tracers = NoTracers()`.\n\n ## The model coefficient for the turbulence closure is defined via the\n ## [CLIMAParameters\n ## package](https://CliMA.github.io/CLIMAParameters.jl/latest/) A reference\n ## state for the linearisation step is also defined.\n T_surface = FT(300)\n T_min_ref = FT(0)\n T_profile = DryAdiabaticProfile{FT}(param_set, T_surface, T_min_ref)\n ref_state = HydrostaticState(T_profile)\n\n ## Here we assemble the `AtmosModel`.\n _C_smag = FT(C_smag(param_set))\n physics = AtmosPhysics{FT}(\n param_set; ## Parameter set corresponding to earth parameters\n ref_state = ref_state, ## Reference state\n turbulence = SmagorinskyLilly(_C_smag), ## Turbulence closure model\n moisture = DryModel(), ## Exclude moisture variables\n tracers = NTracers{ntracers, FT}(δ_χ), ## Tracer model with diffusivity coefficients\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType, ## Flow in a box, requires the AtmosLESConfigType\n physics; ## Atmos physics\n init_state_prognostic = init_risingbubble!, ## Apply the initial condition\n source = (Gravity(),), ## Gravity is the only source term here\n )\n\n ## Finally, we pass a `Problem Name` string, the mesh information, and the\n ## model type to the [`AtmosLESConfiguration`] object.\n config = ClimateMachine.AtmosLESConfiguration(\n \"DryRisingBubble\", ## Problem title [String]\n N, ## Polynomial order [Int]\n resolution, ## (Δx, Δy, Δz) effective resolution [m]\n xmax, ## Domain maximum size [m]\n ymax, ## Domain maximum size [m]\n zmax, ## Domain maximum size [m]\n param_set, ## Parameter set.\n init_risingbubble!, ## Function specifying initial condition\n model = model, ## Model type\n )\n return config\nend\n\n#md # !!! note\n#md # `Keywords` are used to specify some arguments (see appropriate source\n#md # files).\n\n# ## [Diagnostics](@id config_diagnostics)\n# Here we define the diagnostic configuration specific to this problem.\nfunction config_diagnostics(driver_config)\n interval = \"10000steps\"\n dgngrp = setup_atmos_default_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n )\n return ClimateMachine.DiagnosticsConfiguration([dgngrp])\nend\n\nfunction main()\n ## These are essentially arguments passed to the\n ## [`config_risingbubble`](@ref config-helper) function. For type\n ## consistency we explicitly define the problem floating-precision.\n FT = Float64\n ## We need to specify the polynomial order for the DG discretization,\n ## effective resolution, simulation end-time, the domain bounds, and the\n ## courant-number for the time-integrator. Note how the time-integration\n ## components `solver_config` are distinct from the spatial / model\n ## components in `driver_config`. `init_on_cpu` is a helper keyword argument\n ## that forces problem initialization on CPU (thereby allowing the use of\n ## random seeds, spline interpolants and other special functions at the\n ## initialization step.)\n N = 4\n Δh = FT(125)\n Δv = FT(125)\n resolution = (Δh, Δh, Δv)\n xmax = FT(10000)\n ymax = FT(500)\n zmax = FT(10000)\n t0 = FT(0)\n timeend = FT(100)\n ## For full simulation set `timeend = 1000`\n\n ## Use up to 1.7 if ode_solver is the single rate LSRK144.\n CFL = FT(1.7)\n\n ## Assign configurations so they can be passed to the `invoke!` function\n driver_config = config_risingbubble(FT, N, resolution, xmax, ymax, zmax)\n\n ## Choose an Explicit Single-rate Solver from the existing [`ODESolvers`](@ref ClimateMachine.ODESolvers) options.\n ## Apply the outer constructor to define the `ode_solver`.\n ## The 1D-IMEX method is less appropriate for the problem given the current\n ## mesh aspect ratio (1:1).\n ode_solver_type = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n )\n ## If the user prefers a multi-rate explicit time integrator,\n ## the ode_solver above can be replaced with\n ##\n ## `ode_solver = ClimateMachine.MultirateSolverType(\n ## fast_model = AtmosAcousticGravityLinearModel,\n ## slow_method = LSRK144NiegemannDiehlBusch,\n ## fast_method = LSRK144NiegemannDiehlBusch,\n ## timestep_ratio = 10,\n ## )`\n ## See [ODESolvers](@ref ODESolvers-docs) for all of the available solvers.\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_solver_type = ode_solver_type,\n init_on_cpu = true,\n Courant_number = CFL,\n )\n dgn_config = config_diagnostics(driver_config)\n\n ## Invoke solver (calls `solve!` function for time-integrator), pass the driver,\n ## solver and diagnostic config information.\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (),\n check_euclidean_distance = true,\n )\n\n ## Check that the solution norm is reasonable.\n @test isapprox(result, FT(1); atol = 1.5e-3)\nend\n\n# The experiment definition is now complete. Time to run it.\n\n# ## Running the file\n# `julia --project tutorials/Atmos/risingbubble.jl` will run the\n# experiment from the main ClimateMachine.jl directory, with diagnostics output\n# at the intervals specified in [`config_diagnostics`](@ref\n# config_diagnostics). You can also prescribe command line arguments for\n# simulation update and output specifications. For\n# rapid turnaround, we recommend that you run this experiment on a GPU.\n\n# VTK output can be controlled via command line by\n# setting `parse_clargs=true` in the `ClimateMachine.init`\n# arguments, and then using `--vtk=`.\n\n# ## [Output Visualisation](@id output-viz)\n# See the `ClimateMachine` API interface documentation\n# for generating output.\n#\n#\n# - [VisIt](https://wci.llnl.gov/simulation/computer-codes/visit/)\n# - [Paraview](https://www.paraview.org/)\n# are two commonly used programs for `.vtu` files.\n#\n# For NetCDF or JLD2 diagnostics you may use any of the following tools:\n# Julia's\n# [`NCDatasets`](https://github.com/Alexander-Barth/NCDatasets.jl) and\n# [`JLD2`](https://github.com/JuliaIO/JLD2.jl) packages with a suitable\n#\n# or the known and quick NCDF visualization tool:\n# [`ncview`](http://meteora.ucsd.edu/~pierce/ncview_home_page.html)\n# plotting program.\n\nmain()\n" }, { "path": "tutorials/BalanceLaws/tendency_specification_layer.jl", "content": "# # A functional tendency specification layer\n\n# In the balance law (mutating) functions, where we specify fluxes and sources,\n\n# - [`flux_first_order!`](@ref ClimateMachine.BalanceLaws.flux_first_order!)\n# - [`flux_second_order!`](@ref ClimateMachine.BalanceLaws.flux_second_order!)\n# - and [`source!`](@ref ClimateMachine.BalanceLaws.source!),\n\n# an additional (functional) tendency specification\n# layer can be placed on-top that has several nice\n# properties. The functional layer:\n\n# - Separates tendency definitions from which tendencies are included in a particular model.\n# - Reduces duplicate implementations of tendency definitions (e.g., in optional submodel variants)\n# - Allows a more flexible combination of tendencies\n# - Allows a simple way to loop over all tendencies for all prognostic variables and recover\n# _each_ flux / source term. This will allow us a simple way to evaluate, for example, the energy budget.\n\n# ## Used modules / imports\n\n# Make running locally easier from ClimateMachine.jl/:\nif !(\".\" in LOAD_PATH)\n push!(LOAD_PATH, \".\")\n nothing\nend\n\n# First, using necessary modules:\nusing ClimateMachine.BalanceLaws\nusing ClimateMachine.VariableTemplates\nusing StaticArrays, Test\n\n# Import methods to overload\nimport ClimateMachine.BalanceLaws: prognostic_vars, eq_tends, flux\n\n# ## Define a balance law\n\n# Here, we define a simple balance law:\n\nstruct MyBalanceLaw <: BalanceLaw end\n\n# ## Define prognostic variable types\n\n# Here, we'll define some prognostic variable types,\n# by sub-typing [`AbstractPrognosticVariable`](@ref ClimateMachine.BalanceLaws.AbstractPrognosticVariable),\n# for mass and energy:\nstruct Mass <: AbstractPrognosticVariable end\nstruct Energy <: AbstractPrognosticVariable end\n\n# Define [`prognostic_vars`](@ref ClimateMachine.BalanceLaws.prognostic_vars),\n# which returns _all_ prognostic variables\nprognostic_vars(::MyBalanceLaw) = (Mass(), Energy());\n\n# ## Define tendency definition types\n\n# Tendency definitions types are made by subtyping\n# [`TendencyDef`](@ref ClimateMachine.BalanceLaws.TendencyDef).\n# `TendencyDef` has one type parameters: the\n# `AbstractTendencyType`, which can be either\n# `Flux{FirstOrder}`, `Flux{SecondOrder}`, or `Source`.\nstruct Advection <: TendencyDef{Flux{FirstOrder}} end\nstruct Source1 <: TendencyDef{Source} end\nstruct Source2 <: TendencyDef{Source} end\nstruct Diffusion <: TendencyDef{Flux{SecondOrder}} end\n\n# Define [`eq_tends`](@ref ClimateMachine.BalanceLaws.eq_tends),\n# which returns a tuple of tendency definitions (those sub-typed\n# by [`TendencyDef`](@ref ClimateMachine.BalanceLaws.TendencyDef)),\n# given\n# - the prognostic variable\n# - the model (balance law)\n# - the tendency type ([`Flux`](@ref ClimateMachine.BalanceLaws.Flux) or\n# [`Source`](@ref ClimateMachine.BalanceLaws.Source))\neq_tends(::Mass, ::MyBalanceLaw, ::Flux{FirstOrder}) = (Advection(),);\neq_tends(::Energy, ::MyBalanceLaw, ::Flux{FirstOrder}) = (Advection(),);\neq_tends(::Mass, ::MyBalanceLaw, ::Flux{SecondOrder}) = ();\neq_tends(::Energy, ::MyBalanceLaw, ::Flux{SecondOrder}) = (Diffusion(),);\neq_tends(::Mass, ::MyBalanceLaw, ::Source) = (Source1(), Source2());\neq_tends(::Energy, ::MyBalanceLaw, ::Source) = (Source1(), Source2());\n\n# ## Testing `prognostic_vars` `eq_tends`\n\n# To test that `prognostic_vars` and `eq_tends` were\n# implemented correctly, we'll create a balance law\n# instance and call [`show_tendencies`](@ref ClimateMachine.BalanceLaws.show_tendencies),\n# to make sure that the tendency table is accurate.\n\nbl = MyBalanceLaw()\nshow_tendencies(bl; table_complete = true)\n\n# The table looks correct. Now we're ready to\n# add the specification layer.\n\n# ## Adding the tendency specification layer\n\n# For the purpose of this tutorial, we'll only focus\n# on adding the layer to the first order flux, since\n# doing so for the second order flux and source\n# functions follow the same exact pattern. In other words,\n# we'll add a layer that tests the `Flux{FirstOrder}` column\n# in the table above. First, we'll define individual\n# [`flux`](@ref ClimateMachine.BalanceLaws.flux) kernels:\nflux(::Mass, ::Advection, bl::MyBalanceLaw, args) =\n args.state.ρ * SVector(1, 1, 1);\nflux(::Energy, ::Advection, bl::MyBalanceLaw, args) =\n args.state.ρe * SVector(1, 1, 1);\n\n# !!! note\n# - `flux` should return a 3-componet vector for scalar equations\n# - `flux` should return a 3xN-componet tensor for N-component vector equations\n# - `source` should return a scalar for scalar equations\n# - `source` should return a N-componet vector for N-component vector equations\n\n# Define `flux_first_order!` and utilize `eq_tends`\nfunction flux_first_order!(\n bl::MyBalanceLaw,\n flx::Grad,\n state::Vars,\n aux,\n t,\n direction,\n)\n\n tend_type = Flux{FirstOrder}()\n args = (; state, aux, t, direction)\n\n ## `Σfluxes(Mass(), eq_tends(Mass(), bl, tend_type), bl, args)` calls\n ## `flux(::Mass, ::Advection, ...)` defined above:\n eqt_ρ = eq_tends(Mass(), bl, tend_type)\n flx.ρ = Σfluxes(Mass(), eqt_ρ, bl, args)\n\n ## `Σfluxes(Energy(), eq_tends(Energy(), bl, tend_type), bl, args)` calls\n ## `flux(::Energy, ::Advection, ...)` defined above:\n eqt_ρe = eq_tends(Energy(), bl, tend_type)\n flx.ρe = Σfluxes(Energy(), eqt_ρe, bl, args)\n return nothing\nend;\n\n# ## Testing the tendency specification layer\n\n# Now, let's test `flux_first_order!` we need to initialize\n# some dummy data to call it first:\n\nFT = Float64; # float type\naux = (); # auxiliary fields\nt = 0.0; # time\ndirection = nothing; # Direction\n\nstate = Vars{@vars(ρ::FT, ρe::FT)}([1, 2]);\nflx = Grad{@vars(ρ::FT, ρe::FT)}(zeros(MArray{Tuple{3, 2}, FT}));\n\n# call `flux_first_order!`\nflux_first_order!(bl, flx, state, aux, t, direction);\n\n# Test that `flx` has been properly mutated:\n@testset \"Test results\" begin\n @test flx.ρ == [1, 1, 1]\n @test flx.ρe == [2, 2, 2]\nend\n\nnothing\n" }, { "path": "tutorials/Diagnostics/Debug/StateCheck.jl", "content": "# # State debug statistics\n#\n# This page shows how to use the `StateCheck` functions to get basic\n# statistics for nodal values of fields held in ClimateMachine `MPIStateArray`\n# data structures. The `StateCheck` functions can be used to\n#\n# 1. Generate statistics on `MPIStateArray` holding the state of a ClimateMachine experiment.\n#\n# and to\n#\n# 2. Compare against saved reference statistics from ClimateMachine `MPIStateArray`\n# variables. This can enable simple automated regression test checks for\n# detecting unexpected changes introduced into numerical experiments\n# by code updates.\n#\n# These two cases are shown below:\n\n\n# ## 1. Generating statistics for a set of MPIStateArrays\n#\n# Here we create a callback that can generate statistics for an arbitrary\n# set of the MPIStateArray type variables of the sort that hold persistent state for\n# ClimateMachine models. We then invoke the call back to show the statistics.\n#\n# In regular use the `MPIStateArray` variables will come from model configurations.\n# Here we create a dummy set of `MPIStateArray` variables for use in stand alone\n# examples.\n\n# ### Create a dummy set of MPIStateArrays\n#\n# First we set up two `MPIStateArray` variables. This need a few packages to be in placeT,\n# and utilizes some utility functions to create the array and add named\n# persistent state variables.\n# This is usually handled automatically as part of model definition in regular\n# ClimateMachine activity.\n# Calling `ClimateMachine.init()` includes initializing GPU CUDA and MPI parallel\n# processing options that match the hardware/software system in use.\n\n# Set up a basic environment\nusing MPI\nusing StaticArrays\nusing Random\nusing ClimateMachine\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.StateCheck\n\nClimateMachine.init()\nFT = Float64\n\n# Define some dummy vector and tensor abstract variables with associated types\n# and dimensions\nF1 = @vars begin\n ν∇u::SMatrix{3, 2, FT, 6}\n κ∇θ::SVector{3, FT}\nend\nF2 = @vars begin\n u::SVector{2, FT}\n θ::SVector{1, FT}\nend\nnothing # hide\n\n# Create `MPIStateArray` variables with arrays to hold elements of the\n# vectors and tensors\nQ1 = MPIStateArray{Float32, F1}(\n MPI.COMM_WORLD,\n ClimateMachine.array_type(),\n 4,\n 9,\n 8,\n)\nQ2 = MPIStateArray{Float64, F2}(\n MPI.COMM_WORLD,\n ClimateMachine.array_type(),\n 4,\n 3,\n 8,\n)\nnothing # hide\n\n# ### Create a call-back\n#\n# Now we can create a `StateCheck` call-back, _cb_, tied to the `MPIStateArray`\n# variables _Q1_ and _Q2_. Each `MPIStateArray` in the array\n# of `MPIStateArray` variables tracked is paired with a label\n# to identify it. The call-back is also given a frequency (in time step numbers) and\n# precision for printing summary tables.\ncb = ClimateMachine.StateCheck.sccreate(\n [(Q1, \"My gradients\"), (Q2, \"My fields\")],\n 1;\n prec = 15,\n)\nGenericCallbacks.init!(cb, nothing, nothing, nothing, nothing)\nnothing # hide\n\n# ### Invoke the call-back\n#\n# The call-back is of type `ClimateMachine.GenericCallbacks.EveryXSimulationSteps`\n# and in regular use is designed to be passed to a ClimateMachine timestepping\n# solver e.g.\ntypeof(cb)\n\n# Here, for demonstration purposes, we can invoke\n# the call-back after simply initializing the `MPIStateArray` fields to a random\n# set of values e.g.\nQ1.data .= rand(MersenneTwister(0), Float32, size(Q1.data))\nQ2.data .= rand(MersenneTwister(0), Float64, size(Q2.data))\nGenericCallbacks.call!(cb, nothing, nothing, nothing, nothing)\n\n# ## 2. Comparing to reference values\n\n# ### Generate arrays of reference values\n#\n# StateCheck functions can generate text that can be used to set the value of stored\n# arrays that can be used in a reference test for subsequent regression testing. This\n# involves 3 steps.\n#\n# **Step 1.** First a reference array setting program code is generated from the latest\n# state of a given callback e.g.\n\nClimateMachine.StateCheck.scprintref(cb)\n\n# **Step 2.** Next the array setting program code is executed (see below). At this stage the _parr[]_ array\n# context may be hand edited. The parr[] array sets a target number of decimal places for\n# matching against reference values in _varr[]_. For different experiments and different fields\n# the degree of precision that constitutes failing a regression test may vary. Choosing the\n# _parr[]_ values requires some sense as to the stability of the particular numerical\n# and physical scenario an experiment represents. In the example below some precision\n# settings have been hand edited from the default of 16 to illustrate the process.\n\n#! format: off\nvarr = [\n [ \"My gradients\", \"ν∇u[1]\", 1.34348869323730468750e-04, 9.84732866287231445313e-01, 5.23545503616333007813e-01, 3.08209930764271777814e-01 ],\n [ \"My gradients\", \"ν∇u[2]\", 1.16317868232727050781e-01, 9.92088317871093750000e-01, 4.83800649642944335938e-01, 2.83350456014221541157e-01 ],\n [ \"My gradients\", \"ν∇u[3]\", 1.05845928192138671875e-03, 9.51775908470153808594e-01, 4.65474426746368408203e-01, 2.73615551085745090099e-01 ],\n [ \"My gradients\", \"ν∇u[4]\", 5.97668886184692382813e-02, 9.68048095703125000000e-01, 5.42618036270141601563e-01, 2.81570862027933854765e-01 ],\n [ \"My gradients\", \"ν∇u[5]\", 8.31030607223510742188e-02, 9.35931921005249023438e-01, 5.05405902862548828125e-01, 2.46073509972619536290e-01 ],\n [ \"My gradients\", \"ν∇u[6]\", 3.09681892395019531250e-02, 9.98341441154479980469e-01, 4.54375565052032470703e-01, 3.09461067853178561915e-01 ],\n [ \"My gradients\", \"κ∇θ[1]\", 8.47448110580444335938e-02, 9.94180679321289062500e-01, 5.27157366275787353516e-01, 2.92455951648181833313e-01 ],\n [ \"My gradients\", \"κ∇θ[2]\", 1.20514631271362304688e-02, 9.93527650833129882813e-01, 4.71063584089279174805e-01, 2.96449027197666359346e-01 ],\n [ \"My gradients\", \"κ∇θ[3]\", 8.14980268478393554688e-02, 9.55443382263183593750e-01, 5.05038917064666748047e-01, 2.77201022741208891187e-01 ],\n [ \"My fields\", \"u[1]\", 4.31410233294131639781e-02, 9.97140933049696531754e-01, 4.62139750850942054861e-01, 3.23076684924287371725e-01 ],\n [ \"My fields\", \"u[2]\", 1.01416659908237782872e-02, 9.14712023896926407218e-01, 4.76160523012988778913e-01, 2.71443440757963339038e-01 ],\n [ \"My fields\", \"θ[1]\", 6.58965491052394547467e-02, 9.73216404386510802738e-01, 4.60007166313864512830e-01, 2.87310472114545079059e-01 ],\n]\nparr = [\n [ \"My gradients\", \"ν∇u[1]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[2]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[3]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[4]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[5]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"ν∇u[6]\", 16, 7, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[1]\", 16, 16, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[2]\", 16, 16, 16, 0 ],\n [ \"My gradients\", \"κ∇θ[3]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"u[1]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"u[2]\", 16, 16, 16, 0 ],\n [ \"My fields\", \"θ[1]\", 16, 16, 16, 0 ],\n]\n#! format: on\n\n# **Step 3.** Finally a call-back stored value can be compared for consistency to with _parr[]_ decimal places\n\nClimateMachine.StateCheck.scdocheck(cb, (varr, parr))\nnothing # hide\n\n# In this trivial case the match is guaranteed. The function will return _true_ to the calling\n# routine and this can be passed to an `@test` block.\n#\n# However we can modify the reference test values to\n# see the effect of a mismatch e.g.\nvarr[1][3] = varr[1][3] * 10.0\nClimateMachine.StateCheck.scdocheck(cb, (varr, parr))\nnothing # hide\n\n# Here the mis-matching field is highlighted with _N(0)_ indicating that the precision\n# was not met and actual match length was (in this case) 0. If any field fails the test returns false\n# for use in any regression testing control logic.\n" }, { "path": "tutorials/Land/Heat/heat_equation.jl", "content": "# # Heat equation tutorial\n\n# In this tutorial, we'll be solving the [heat\n# equation](https://en.wikipedia.org/wiki/Heat_equation):\n\n# ``\n# \\frac{∂ ρcT}{∂ t} + ∇ ⋅ (-α ∇ρcT) = 0\n# ``\n\n# where\n# - `t` is time\n# - `α` is the thermal diffusivity\n# - `T` is the temperature\n# - `ρ` is the density\n# - `c` is the heat capacity\n# - `ρcT` is the thermal energy\n\n# To put this in the form of ClimateMachine's [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw), we'll re-write the equation as:\n\n# ``\n# \\frac{∂ ρcT}{∂ t} + ∇ ⋅ (F(α, ρcT, t)) = 0\n# ``\n\n# where\n# - ``F(α, ρcT, t) = -α ∇ρcT`` is the second-order flux\n\n# with boundary conditions\n# - Fixed temperature ``T_{surface}`` at ``z_{min}`` (non-zero Dirichlet)\n# - No thermal flux at ``z_{min}`` (zero Neumann)\n\n# Solving these equations is broken down into the following steps:\n# 1) Preliminary configuration\n# 2) PDEs\n# 3) Space discretization\n# 4) Time discretization / solver\n# 5) Solver hooks / callbacks\n# 6) Solve\n# 7) Post-processing\n\n# # Preliminary configuration\n\n# ## [Loading code](@id Loading-code-heat)\n\n# First, we'll load our pre-requisites:\n# - load external packages:\nusing MPI\nusing OrderedCollections\nusing Plots\nusing StaticArrays\nusing OrdinaryDiffEq\nusing DiffEqBase\n\n# - load CLIMAParameters and set up to use it:\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\n# - load necessary ClimateMachine modules:\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux\n\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\n\n# - import necessary ClimateMachine modules: (`import`ing enables us to\n# provide implementations of these structs/methods)\nimport ClimateMachine.BalanceLaws:\n vars_state,\n source!,\n flux_second_order!,\n flux_first_order!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n nodal_update_auxiliary_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n BoundaryCondition,\n boundary_conditions,\n boundary_state!\n\nimport ClimateMachine.DGMethods: calculate_dt\n\n# ## Initialization\n\n# Define the float type (`Float64` or `Float32`)\nconst FT = Float64;\n# Initialize ClimateMachine for CPU.\nClimateMachine.init(; disable_gpu = true);\n\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\n\n# Load some helper functions for plotting\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Define the set of Partial Differential Equations (PDEs)\n\n# ## Define the model\n\n# Model parameters can be stored in the particular [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw), in this case, a `HeatModel`:\n\nBase.@kwdef struct HeatModel{FT, APS} <: BalanceLaw\n \"Parameters\"\n param_set::APS\n \"Heat capacity\"\n ρc::FT = 1\n \"Thermal diffusivity\"\n α::FT = 0.01\n \"Initial conditions for temperature\"\n initialT::FT = 295.15\n \"Bottom boundary value for temperature (Dirichlet boundary conditions)\"\n T_bottom::FT = 300.0\n \"Top flux (α∇ρcT) at top boundary (Neumann boundary conditions)\"\n flux_top::FT = 0.0\nend\n\n# Create an instance of the `HeatModel`:\nm = HeatModel{FT, typeof(param_set)}(; param_set = param_set);\n\n# This model dictates the flow control, using [Dynamic Multiple\n# Dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch), for which\n# kernels are executed.\n\n# ## Define the variables\n\n# All of the methods defined in this section were `import`ed in\n# [Loading code](@ref Loading-code-heat) to let us provide\n# implementations for our `HeatModel` as they will be used by\n# the solver.\n\n# Specify auxiliary variables for `HeatModel`\nvars_state(::HeatModel, ::Auxiliary, FT) = @vars(z::FT, T::FT);\n\n# Specify prognostic variables, the variables solved for in the PDEs, for\n# `HeatModel`\nvars_state(::HeatModel, ::Prognostic, FT) = @vars(ρcT::FT);\n\n# Specify state variables whose gradients are needed for `HeatModel`\nvars_state(::HeatModel, ::Gradient, FT) = @vars(ρcT::FT);\n\n# Specify gradient variables for `HeatModel`\nvars_state(::HeatModel, ::GradientFlux, FT) = @vars(α∇ρcT::SVector{3, FT});\n\n# ## Define the compute kernels\n\n# Specify the initial values in `aux::Vars`, which are available in\n# `init_state_prognostic!`. Note that\n# - this method is only called at `t=0`\n# - `aux.z` and `aux.T` are available here because we've specified `z` and `T`\n# in `vars_state` given `Auxiliary`\n# in `vars_state`\nfunction nodal_init_state_auxiliary!(\n m::HeatModel,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.z = geom.coord[3]\n aux.T = m.initialT\nend;\n\n# Specify the initial values in `state::Vars`. Note that\n# - this method is only called at `t=0`\n# - `state.ρcT` is available here because we've specified `ρcT` in\n# `vars_state` given `Prognostic`\nfunction init_state_prognostic!(\n m::HeatModel,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n state.ρcT = m.ρc * aux.T\nend;\n\n# The remaining methods, defined in this section, are called at every\n# time-step in the solver by the [`BalanceLaw`](@ref\n# ClimateMachine.BalanceLaws.BalanceLaw) framework.\n\n# Compute/update all auxiliary variables at each node. Note that\n# - `aux.T` is available here because we've specified `T` in\n# `vars_state` given `Auxiliary`\nfunction nodal_update_auxiliary_state!(\n m::HeatModel,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n aux.T = state.ρcT / m.ρc\nend;\n\n# Since we have second-order fluxes, we must tell `ClimateMachine` to compute\n# the gradient of `ρcT`. Here, we specify how `ρcT` is computed. Note that\n# - `transform.ρcT` is available here because we've specified `ρcT` in\n# `vars_state` given `Gradient`\nfunction compute_gradient_argument!(\n m::HeatModel,\n transform::Vars,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n transform.ρcT = state.ρcT\nend;\n\n# Specify where in `diffusive::Vars` to store the computed gradient from\n# `compute_gradient_argument!`. Note that:\n# - `diffusive.α∇ρcT` is available here because we've specified `α∇ρcT` in\n# `vars_state` given `Gradient`\n# - `∇transform.ρcT` is available here because we've specified `ρcT` in\n# `vars_state` given `Gradient`\nfunction compute_gradient_flux!(\n m::HeatModel,\n diffusive::Vars,\n ∇transform::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n)\n diffusive.α∇ρcT = -m.α * ∇transform.ρcT\nend;\n\n# We have no sources, nor non-diffusive fluxes.\nfunction source!(m::HeatModel, _...) end;\nfunction flux_first_order!(m::HeatModel, _...) end;\n\n# Compute diffusive flux (``F(α, ρcT, t) = -α ∇ρcT`` in the original PDE).\n# Note that:\n# - `diffusive.α∇ρcT` is available here because we've specified `α∇ρcT` in\n# `vars_state` given `GradientFlux`\nfunction flux_second_order!(\n m::HeatModel,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n)\n flux.ρcT += diffusive.α∇ρcT\nend;\n\n# ### Boundary conditions\n\n# Second-order terms in our equations, ``∇⋅(F)`` where ``F = -α∇ρcT``, are\n# internally reformulated to first-order unknowns.\n# Boundary conditions must be specified for all unknowns, both first-order and\n# second-order unknowns which have been reformulated.\nstruct DirichletBC <: BoundaryCondition end;\nstruct NeumannBC <: BoundaryCondition end;\n\nboundary_conditions(::HeatModel) = (DirichletBC(), NeumannBC())\n\n# The boundary conditions for `ρcT` (first order unknown)\nfunction boundary_state!(\n nf,\n bc::DirichletBC,\n m::HeatModel,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n state⁻::Vars,\n aux⁻::Vars,\n t,\n _...,\n)\n ## Apply Dirichlet BCs\n state⁺.ρcT = m.ρc * m.T_bottom\nend;\nfunction boundary_state!(\n nf,\n bc::NeumannBC,\n m::HeatModel,\n state⁺::Vars,\n aux⁺::Vars,\n n⁻,\n state⁻::Vars,\n aux⁻::Vars,\n t,\n _...,\n)\n nothing\nend;\n\n# The boundary conditions for `ρcT` are specified here for second-order\n# unknowns\nfunction boundary_state!(\n nf,\n bc::DirichletBC,\n m::HeatModel,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n state⁻::Vars,\n diff⁻::Vars,\n hyperdiff⁻::Vars,\n aux⁻::Vars,\n t,\n _...,\n)\n nothing\nend;\nfunction boundary_state!(\n nf,\n bc::NeumannBC,\n m::HeatModel,\n state⁺::Vars,\n diff⁺::Vars,\n hyperdiff⁺::Vars,\n aux⁺::Vars,\n n⁻,\n state⁻::Vars,\n diff⁻::Vars,\n hyperdiff⁻::Vars,\n aux⁻::Vars,\n t,\n _...,\n)\n ## Apply Neumann BCs\n diff⁺.α∇ρcT = n⁻ * m.flux_top\nend;\n\n# # Spatial discretization\n\n# Prescribe polynomial order of basis functions in finite elements\nN_poly = 5;\n\n# Specify the number of vertical elements\nnelem_vert = 10;\n\n# Specify the domain height\nzmax = FT(1);\n\n# Establish a `ClimateMachine` single stack configuration\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"HeatEquation\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n# # Time discretization / solver\n\n# Specify simulation time (SI units)\nt0 = FT(0)\ntimeend = FT(40)\n\n# In this section, we initialize the state vector and allocate memory for\n# the solution in space (`dg` has the model `m`, which describes the PDEs\n# as well as the function used for initialization). `SolverConfiguration`\n# initializes the ODE solver, by default an explicit Low-Storage\n# [Runge-Kutta](https://en.wikipedia.org/wiki/Runge%E2%80%93Kutta_methods)\n# method. In this tutorial, we prescribe an option for an implicit\n# `Kvaerno3` method.\n\n# First, let's define how the time-step is computed, based on the\n# [Fourier number](https://en.wikipedia.org/wiki/Fourier_number)\n# (i.e., diffusive Courant number) is defined. Because\n# the `HeatModel` is a custom model, we must define how both are computed.\n# First, we must define our own implementation of `DGMethods.calculate_dt`,\n# (which we imported):\nfunction calculate_dt(dg, model::HeatModel, Q, Courant_number, t, direction)\n Δt = one(eltype(Q))\n CFL = DGMethods.courant(diffusive_courant, dg, model, Q, Δt, t, direction)\n return Courant_number / CFL\nend\n\n# Next, we'll define our implementation of `diffusive_courant`:\nfunction diffusive_courant(\n m::HeatModel,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n t,\n direction,\n)\n return Δt * m.α / (Δx * Δx)\nend\n\n# Finally, we initialize the state vector and solver\n# configuration based on the given Fourier number.\n# Note that, we can use a much larger Fourier number\n# for implicit solvers as compared to explicit solvers.\nuse_implicit_solver = false\nif use_implicit_solver\n given_Fourier = FT(30)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config;\n ode_solver_type = ImplicitSolverType(OrdinaryDiffEq.Kvaerno3(\n autodiff = false,\n linsolve = LinSolveGMRES(),\n )),\n Courant_number = given_Fourier,\n CFL_direction = VerticalDirection(),\n )\nelse\n given_Fourier = FT(0.7)\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config;\n Courant_number = given_Fourier,\n CFL_direction = VerticalDirection(),\n )\nend;\n\n\ngrid = solver_config.dg.grid;\nQ = solver_config.Q;\naux = solver_config.dg.state_auxiliary;\n\n# ## Inspect the initial conditions\n\n# Let's export a plot of the initial state\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nz_scale = 100; # convert from meters to cm\nz_key = \"z\";\nz_label = \"z [cm]\";\nz = get_z(grid; z_scale = z_scale, rm_dupes = true);\n\n# Create an array to store the solution:\ndons_arr = Dict[dict_of_nodal_states(solver_config; interp = true)] # store initial condition at ``t=0``\ntime_data = FT[0] # store time data\n\nexport_plot(\n z,\n time_data,\n dons_arr,\n (\"ρcT\",),\n joinpath(output_dir, \"initial_condition.png\");\n xlabel = \"ρcT\",\n ylabel = z_label,\n xlims = (m.initialT - 1, m.T_bottom + 1),\n);\n# ![](initial_condition.png)\n\n# It matches what we have in `init_state_prognostic!(m::HeatModel, ...)`, so\n# let's continue.\n\n# # Solver hooks / callbacks\n\n# Define the number of outputs from `t0` to `timeend`\nconst n_outputs = 5;\n\n# This equates to exports every ceil(Int, timeend/n_outputs) time-step:\nconst every_x_simulation_time = ceil(Int, timeend / n_outputs);\n\n# The `ClimateMachine`'s time-steppers provide hooks, or callbacks, which\n# allow users to inject code to be executed at specified intervals. In this\n# callback, a dictionary of prognostic and auxiliary states are appended to\n# `dons_arr` for time the callback is executed. In addition, time is collected\n# and appended to `time_data`.\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n push!(dons_arr, dict_of_nodal_states(solver_config; interp = true))\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Solve\n\n# This is the main `ClimateMachine` solver invocation. While users do not have\n# access to the time-stepping loop, code may be injected via `user_callbacks`,\n# which is a `Tuple` of callbacks in [`GenericCallbacks`](@ref ClimateMachine.GenericCallbacks).\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,));\n\n# Append result at the end of the last time step:\npush!(dons_arr, dict_of_nodal_states(solver_config; interp = true));\npush!(time_data, gettime(solver_config.solver));\n\n# # Post-processing\n\n# Our solution is stored in the array of dictionaries `dons_arr` whose keys are\n# the output interval. The next level keys are the variable names, and the\n# values are the values along the grid:\n\n# To get `T` at ``t=0``, we can use `T_at_t_0 = dons_arr[1][\"T\"][:]`\n@show keys(dons_arr[1])\n\n# Let's plot the solution:\n\nexport_plot(\n z,\n time_data,\n dons_arr,\n (\"ρcT\",),\n joinpath(output_dir, \"solution_vs_time.png\");\n xlabel = \"ρcT\",\n ylabel = z_label,\n);\n# ![](solution_vs_time.png)\n\nexport_contour(\n z,\n time_data,\n dons_arr,\n \"ρcT\",\n joinpath(output_dir, \"solution_contour.png\");\n ylabel = \"z [cm]\",\n)\n# ![](solution_contour.png)\n\n# The results look as we would expect: a fixed temperature at the bottom is\n# resulting in heat flux that propagates up the domain. To run this file, and\n# inspect the solution in `dons_arr`, include this tutorial in the Julia REPL\n# with:\n\n# ```julia\n# include(joinpath(\"tutorials\", \"Land\", \"Heat\", \"heat_equation.jl\"))\n# ```\n" }, { "path": "tutorials/Land/Soil/Artifacts.toml", "content": "[bonan_soil_heat]\ngit-tree-sha1 = \"9f0933ccd6d902d0b75cc5bc9eb9dea3aa3708d7\"\n" }, { "path": "tutorials/Land/Soil/Coupled/equilibrium_test.jl", "content": "# # Coupled heat and water equations tending towards equilibrium\n\n# Other tutorials, such as the [soil heat tutorial](../Heat/bonan_heat_tutorial.md)\n# and [Richards equation tutorial](../Water/equilibrium_test.md)\n# demonstrate how to solve the heat\n# equation or Richard's equation without considering\n# dynamic interactions between the two. As an example, the user could\n# prescribe a fixed function of space and time for the liquid water content,\n# and use that to drive the heat equation, but without allowing the water\n# content to dynamically evolve according to Richard's equation and without\n# allowing the changing temperature of the soil to affect the water\n# evolution.\n\n# Here we show how to solve the interacting heat and water equations,\n# in sand, but without phase changes. This allows us to capture\n# behavior that is not present in the decoupled equations.\n\n# The equations\n# are:\n\n# ``\n# \\frac{∂ ρe_{int}}{∂ t} = ∇ ⋅ κ(θ_l, θ_i; ν, ...) ∇T + ∇ ⋅ ρe_{int_{liq}} K (T,θ_l, θ_i; ν, ...) \\nabla h( ϑ_l, z; ν, ...)\n# ``\n\n# ``\n# \\frac{ ∂ ϑ_l}{∂ t} = ∇ ⋅ K (T,θ_l, θ_i; ν, ...) ∇h( ϑ_l, z; ν, ...).\n# ``\n\n# Here\n\n# ``t`` is the time (s),\n\n# ``z`` is the location in the vertical (m),\n\n# ``ρe_{int}`` is the volumetric internal energy of the soil (J/m^3),\n\n# ``T`` is the temperature of the soil (K),\n\n# ``κ`` is the thermal conductivity (W/m/K),\n\n# ``ρe_{int_{liq}}`` is the volumetric internal energy of liquid water (J/m^3),\n\n# ``K`` is the hydraulic conductivity (m/s),\n\n# ``h`` is the hydraulic head (m),\n\n# ``ϑ_l`` is the augmented volumetric liquid water fraction,\n\n# ``θ_i`` is the volumetric ice fraction, and\n\n# ``ν, ...`` denotes parameters relating to soil type, such as porosity.\n\n\n# We will solve this equation in an effectively 1-d domain with ``z ∈ [-1,0]``,\n# and with the following boundary and initial conditions:\n\n# ``- κ ∇T(t, z = 0) = 0 ẑ``\n\n# `` -κ ∇T(t, z = -1) = 0 ẑ ``\n\n# `` T(t = 0, z) = T_{min} + (T_{max}-T_{min}) e^{Cz}``\n\n# ``- K ∇h(t, z = 0) = 0 ẑ ``\n\n# `` -K ∇h(t, z = -1) = 0 ẑ``\n\n# `` ϑ(t = 0, z) = ϑ_{min} + (ϑ_{max}-ϑ_{min}) e^{Cz}, ``\n\n# where ``C, T_{min}, T_{max}, ϑ_{min},`` and ``ϑ_{max}`` are\n# constants.\n\n\n# If we evolve this system for times long compared to the dynamical timescales\n# of the system, we expect it to reach an equilibrium where\n# the LHS of these equations tends to zero.\n# Assuming zero fluxes at the boundaries, the resulting equilibrium state\n# should satisfy ``∂h/∂z = 0`` and ``∂T/∂z = 0``. Physically, this means that\n# the water settles into a vertical profile in which\n# the resulting pressure balances gravity and that the temperature\n# is constant across the domain.\n\n# We verify that the system is approaching this equilibrium, and we also sketch out\n# an analytic calculation for the final temperature in equilibrium.\n\n# # Import necessary modules\n# External (non - CliMA) modules\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Plots\n\n# CliMA Parameters\nusing CLIMAParameters\nusing CLIMAParameters.Planet: ρ_cloud_liq, ρ_cloud_ice, cp_l, cp_i, T_0, LH_f0\n\n# ClimateMachine modules\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw,\n Prognostic,\n Auxiliary,\n Gradient,\n GradientFlux,\n vars_state,\n parameter_set\n\n# # Preliminary set-up\n\n# Get the parameter set, which holds constants used across CliMA models:\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n# Initialize and pick a floating point precision:\nClimateMachine.init()\nconst FT = Float64;\n\n# Load plot helpers:\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# Set soil parameters to be consistent with sand.\n# Please see e.g. the [soil heat tutorial](../Heat/bonan_heat_tutorial.md)\n# for other soil type parameters, or [Cosby1984](@cite).\n\n# The porosity:\nporosity = FT(0.395);\n# Soil solids\n# are the components of soil besides water, ice, gases, and air.\n# We specify the soil component fractions, relative to all soil solids.\n# These should sum to unity; they do not account for pore space.\nν_ss_quartz = FT(0.92)\nν_ss_minerals = FT(0.08)\nν_ss_om = FT(0.0)\nν_ss_gravel = FT(0.0);\n# Other parameters include the hydraulic conductivity at saturation, the specific\n# storage, and the van Genuchten parameters for sand.\n# We recommend Chapter 8 of [Bonan19a](@cite) for finding parameters\n# for other soil types.\nKsat = FT(4.42 / 3600 / 100) # m/s\nS_s = FT(1e-3) #inverse meters\nvg_n = FT(1.89)\nvg_α = FT(7.5); # inverse meters\n\n\n# Other constants needed:\nκ_quartz = FT(7.7) # W/m/K\nκ_minerals = FT(2.5) # W/m/K\nκ_om = FT(0.25) # W/m/K\nκ_liq = FT(0.57) # W/m/K\nκ_ice = FT(2.29); # W/m/K\n# The particle density of organic material-free soil is\n# equal to the particle density of quartz and other minerals ([BallandArp2005](@cite)):\nρp = FT(2700); # kg/m^3\n# We calculate the thermal conductivities for the solid material\n# and for saturated soil. These functions are taken from [BallandArp2005](@cite).\nκ_solid = k_solid(ν_ss_om, ν_ss_quartz, κ_quartz, κ_minerals, κ_om)\nκ_sat_frozen = ksat_frozen(κ_solid, porosity, κ_ice)\nκ_sat_unfrozen = ksat_unfrozen(κ_solid, porosity, κ_liq);\n# Next, we calculate the volumetric heat capacity of dry soil. Dry soil\n# refers to soil that has no water content.\nρc_ds = FT((1 - porosity) * 1.926e06); # J/m^3/K\n# We collect the majority of the parameters needed\n# for modeling heat and water flow in soil in `soil_param_functions`,\n# an object of type [`SoilParamFunctions`](@ref ClimateMachine.Land.SoilParamFunctions).\n# Parameters used only for hydrology are stored in `water`, which\n# is of type\n# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions).\n\nsoil_param_functions = SoilParamFunctions(\n FT;\n porosity = porosity,\n ν_ss_gravel = ν_ss_gravel,\n ν_ss_om = ν_ss_om,\n ν_ss_quartz = ν_ss_quartz,\n ρc_ds = ρc_ds,\n ρp = ρp,\n κ_solid = κ_solid,\n κ_sat_unfrozen = κ_sat_unfrozen,\n κ_sat_frozen = κ_sat_frozen,\n water = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s),\n);\n\n# # Initial and Boundary conditions\n\n# As we are not including the equations for phase changes in this tutorial,\n# we chose temperatures that are above the freezing point of water.\n\n# The initial temperature profile:\nfunction T_init(aux)\n FT = eltype(aux)\n zmax = FT(0)\n zmin = FT(-1)\n T_max = FT(289.0)\n T_min = FT(288.0)\n c = FT(20.0)\n z = aux.z\n output = T_min + (T_max - T_min) * exp(-(z - zmax) / (zmin - zmax) * c)\n return output\nend;\n\n\n# The initial water profile:\nfunction ϑ_l0(aux)\n FT = eltype(aux)\n zmax = FT(0)\n zmin = FT(-1)\n theta_max = FT(porosity * 0.5)\n theta_min = FT(porosity * 0.4)\n c = FT(20.0)\n z = aux.z\n output =\n theta_min +\n (theta_max - theta_min) * exp(-(z - zmax) / (zmin - zmax) * c)\n return output\nend;\n\n\n\n# The boundary value problem in this case\n# requires a boundary condition at the top and the bottom of the domain\n# for each equation being solved. These conditions can be Dirichlet, or Neumann.\n\n# Dirichlet boundary conditions are on `ϑ_l` and\n# `T`, while Neumann boundary conditions are on `-κ∇T` and `-K∇h`. For Neumann\n# conditions, the user supplies a scalar, which is multiplied by `ẑ` within the code.\n\n# Water boundary conditions:\nsurface_water_flux = (aux, t) -> eltype(aux)(0.0)\nbottom_water_flux = (aux, t) -> eltype(aux)(0.0);\n\n# The boundary conditions for the heat equation:\nsurface_heat_flux = (aux, t) -> eltype(aux)(0.0)\nbottom_heat_flux = (aux, t) -> eltype(aux)(0.0);\n\nbc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_heat = Neumann(bottom_heat_flux),\n soil_water = Neumann(bottom_water_flux),\n ),\n surface_bc = LandComponentBC(\n soil_heat = Neumann(surface_heat_flux),\n soil_water = Neumann(surface_water_flux),\n ),\n);\n\n# Next, we define the required `init_soil!` function, which takes the user\n# specified functions of space for `T_init` and `ϑ_l0` and initializes the state\n# variables of volumetric internal energy and augmented liquid fraction. This requires\n# a conversion from `T` to `ρe_int`.\nfunction init_soil!(land, state, aux, localgeo, time)\n myFT = eltype(state)\n ϑ_l = myFT(land.soil.water.initialϑ_l(aux))\n θ_i = myFT(land.soil.water.initialθ_i(aux))\n state.soil.water.ϑ_l = ϑ_l\n state.soil.water.θ_i = θ_i\n param_set = parameter_set(land)\n\n θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_ds = land.soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n\n state.soil.heat.ρe_int = volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n )\nend;\n\n\n# # Create the soil model structure\n# First, for water (this is where the hydrology parameters\n# are supplied):\nsoil_water_model = SoilWaterModel(\n FT;\n viscosity_factor = TemperatureDependentViscosity{FT}(),\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = vanGenuchten(FT; α = vg_α, n = vg_n),\n initialϑ_l = ϑ_l0,\n);\n\n\n# Note that the viscosity of water depends on temperature.\n# We account for the effect that has on the hydraulic conductivity\n# by specifying `viscosity_factor = TemperatureDependentViscosity{FT}()`.\n# The default, if no `viscosity_factor` keyword argument is supplied,\n# is to not include the effect of `T` on viscosity. More guidance about\n# specifying the\n# hydraulic conductivity, and the `hydraulics` model,\n# can be found in the [`hydraulic functions`](../Water/hydraulic_functions.md)\n# tutorial.\n\n# Repeat for heat:\nsoil_heat_model = SoilHeatModel(FT; initialT = T_init)\n\n\n# Combine into a single soil model:\nm_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);\n\n\n# We aren't using any sources or sinks in the equations here, but this is where\n# freeze/thaw terms, runoff, root extraction, etc. would go.\nsources = ();\n\n\n# Create the LandModel - without other components (canopy, carbon, etc):\nm = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n);\n\n\n# # Specify the numerical details\n# Choose a resolution, domain boundaries, integration time,\n# timestep, and ODE solver.\n\nN_poly = 1\nnelem_vert = 50\nzmin = FT(-1)\nzmax = FT(0)\n\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n)\n\nt0 = FT(0)\ntimeend = FT(60 * 60 * 72)\ndt = FT(30.0)\n\n\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\n\n# Determine how often you want output:\nconst n_outputs = 4\nconst every_x_simulation_time = ceil(Int, timeend / n_outputs);\n\n# Store initial condition at ``t=0``,\n# including prognostic, auxiliary, and\n# gradient flux variables:\nstate_types = (Prognostic(), Auxiliary(), GradientFlux())\ndons_arr = Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\ntime_data = FT[0] # store time data\n\n# We specify a function which evaluates `every_x_simulation_time` and returns\n# the state vector, appending the variables we are interested in into\n# `dons_arr`.\n\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Run the integration\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,));\n\n# Get z-coordinate\nz = get_z(solver_config.dg.grid; rm_dupes = true);\n\n# Let's export a plot of the initial state\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nexport_plot(\n z,\n time_data ./ (60 * 60 * 24),\n dons_arr,\n (\"soil.water.ϑ_l\",),\n joinpath(output_dir, \"eq_moisture_plot.png\");\n xlabel = \"ϑ_l\",\n ylabel = \"z (m)\",\n time_units = \"(days)\",\n)\n# ![](eq_moisture_plot.png)\n\nexport_plot(\n z,\n time_data[2:end] ./ (60 * 60 * 24),\n dons_arr[2:end],\n (\"soil.water.K∇h[3]\",),\n joinpath(output_dir, \"eq_hydraulic_head_plot.png\");\n xlabel = \"K∇h (m/s)\",\n ylabel = \"z (m)\",\n time_units = \"(days)\",\n)\n# ![](eq_hydraulic_head_plot.png)\n\nexport_plot(\n z,\n time_data ./ (60 * 60 * 24),\n dons_arr,\n (\"soil.heat.T\",),\n joinpath(output_dir, \"eq_temperature_plot.png\");\n xlabel = \"T (K)\",\n ylabel = \"z (m)\",\n time_units = \"(days)\",\n)\n# ![](eq_temperature_plot.png)\n\nexport_plot(\n z,\n time_data[2:end] ./ (60 * 60 * 24),\n dons_arr[2:end],\n (\"soil.heat.κ∇T[3]\",),\n joinpath(output_dir, \"eq_heat_plot.png\");\n xlabel = \"κ∇T\",\n ylabel = \"z (m)\",\n time_units = \"(days)\",\n)\n# ![](eq_heat_plot.png)\n\n# # Analytic Expectations\n\n# We can determine a priori what we expect the final temperature to be in\n# equilibrium.\n\n# Regardless of the final water profile in equilibrium, we know that\n# the final temperature `T_f` will be a constant across the domain. All\n# water that began with a temperature above this point will cool to `T_f`,\n# and water that began with a temperature below this point will warm to\n# `T_f`. The initial function `T(z)` is equal to `T_f` at a value of\n# `z = z̃`. This is the location in space which divides these two groups\n# (water that warms over time and water that cools over time) spatially.\n# We can solve for `z̃(T_f)` using `T_f = T(z̃)`.\n\n# Next, we can determine the change in energy required to cool\n# the water above `z̃` to `T_f`: it is the integral from `z̃` to the surface\n# at `z = 0` of ` c θ(z) T(z) `, where `c` is the volumetric heat capacity -\n# a constant here - and `θ(z)` is the initial water profile. Compute the energy\n# required to warm the water below `z̃` to `T_f` in a similar way, set equal, and solve\n# for `T_f`. This results in `T_f = 288.056`, which is very close to the mean `T` we observe\n# after 3 days, of `288.054`.\n\n# One could also solve the equation for `ϑ_l` specified by\n# ``∂ h/∂ z = 0`` to determine the functional form of the\n# equilibrium profile of the liquid water.\n\n# # References\n# - [Bonan19a](@cite)\n# - [BallandArp2005](@cite)\n# - [Cosby1984](@cite)\n" }, { "path": "tutorials/Land/Soil/Heat/bonan_heat_tutorial.jl", "content": "# # Solving the heat equation in soil\n\n# This tutorial shows how to use CliMA code to solve the heat\n# equation in soil.\n# For background on the heat equation in general,\n# and how to solve it using CliMA code, please see the\n# [`heat_equation.jl`](../../Heat/heat_equation.md)\n# tutorial.\n\n# The version of the heat equation we are solving here assumes no\n# sources or sinks and no flow of liquid water. It takes the form\n\n# ``\n# \\frac{∂ ρe_{int}}{∂ t} = ∇ ⋅ κ(θ_l, θ_i; ν, ...) ∇T\n# ``\n\n# Here\n\n# ``t`` is the time (s),\n\n# ``z`` is the location in the vertical (m),\n\n# ``ρe_{int}`` is the volumetric internal energy of the soil (J/m^3),\n\n# ``T`` is the temperature of the soil (K),\n\n# ``κ`` is the thermal conductivity (W/m/K),\n\n# ``ϑ_l`` is the augmented volumetric liquid water fraction,\n\n# ``θ_i`` is the volumetric ice fraction, and\n\n# ``ν, ...`` denotes parameters relating to soil type, such as porosity.\n\n\n# We will solve this equation in an effectively 1-d domain with ``z ∈ [-1,0]``,\n# and with the following boundary and initial conditions:\n\n# ``T(t=0, z) = 277.15^\\circ K``\n\n# ``T(t, z = 0) = 288.15^\\circ K ``\n\n# `` -κ ∇T(t, z = -1) = 0 ẑ``\n\n\n# The temperature ``T`` and\n# volumetric internal energy ``ρe_{int}`` are related as\n\n# ``\n# ρe_{int} = ρc_s (θ_l, θ_i; ν, ...) (T - T_0) - θ_i ρ_i LH_{f0}\n# ``\n\n# where\n\n# ``ρc_s`` is the volumetric heat capacity of the soil (J/m^3/K),\n\n# ``T_0`` is the freezing temperature of water,\n\n# ``ρ_i`` is the density of ice (kg/m^3), and\n\n# ``LH_{f0}`` is the latent heat of fusion at ``T_0``.\n\n# In this tutorial, we will use a [`PrescribedWaterModel`](@ref\n# ClimateMachine.Land.PrescribedWaterModel). This option allows\n# the user to specify a function for the spatial and temporal\n# behavior of `θ_i` and `θ_l`; it does not solve Richard's equation\n# for the evolution of moisture. Please see the tutorials\n# in the `Soil/Coupled/` folder or the `Soil/Water/`\n# folder for information on solving\n# Richard's equation, either coupled or uncoupled from the heat equation, respectively.\n\n\n\n# # Import necessary modules\n\n# External (non - CliMA) modules\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Dierckx\nusing Plots\nusing DelimitedFiles\n\n# CliMA Parameters\nusing CLIMAParameters\nusing CLIMAParameters.Planet: ρ_cloud_liq, ρ_cloud_ice, cp_l, cp_i, T_0, LH_f0\n\n\n# ClimateMachine modules\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw,\n Prognostic,\n Auxiliary,\n Gradient,\n GradientFlux,\n vars_state,\n parameter_set\nimport ClimateMachine.DGMethods: calculate_dt\nusing ArtifactWrappers\n\n# # Preliminary set-up\n# Get the parameter set, which holds constants used across CliMA models.\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n# Initialize and pick a floating point precision.\nClimateMachine.init()\nconst FT = Float32;\n# Load functions that will help with plotting\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Determine soil parameters\n\n# Below are the soil component fractions for various soil\n# texture classes, from [Cosby1984](@cite) and [Bonan19a](@cite).\n# Note that these fractions are volumetric fractions, relative\n# to other soil solids, i.e. not including pore space. These are denoted `ν_ss_i`; the CliMA\n# Land Documentation uses the symbol `ν_i` to denote the volumetric fraction\n# of a soil component `i` relative to the soil, including pore space.\n\nν_ss_silt_array =\n FT.(\n [5.0, 12.0, 32.0, 70.0, 39.0, 15.0, 56.0, 34.0, 6.0, 47.0, 20.0] ./\n 100.0,\n )\nν_ss_quartz_array =\n FT.(\n [92.0, 82.0, 58.0, 17.0, 43.0, 58.0, 10.0, 32.0, 52.0, 6.0, 22.0] ./\n 100.0,\n )\nν_ss_clay_array =\n FT.(\n [3.0, 6.0, 10.0, 13.0, 18.0, 27.0, 34.0, 34.0, 42.0, 47.0, 58.0] ./\n 100.0,\n )\nporosity_array =\n FT.([\n 0.395,\n 0.410,\n 0.435,\n 0.485,\n 0.451,\n 0.420,\n 0.477,\n 0.476,\n 0.426,\n 0.492,\n 0.482,\n ]);\n\n\n# The soil types that correspond to array elements above are, in order,\n# sand, loamy sand, sandy loam, silty loam, loam, sandy clay loam,\n# silty clay loam, clay loam, sandy clay, silty clay, and clay.\n\n# Here we choose the soil type to be sandy.\n# The soil column is uniform in space and time.\nsoil_type_index = 1\nν_ss_minerals =\n ν_ss_clay_array[soil_type_index] + ν_ss_silt_array[soil_type_index]\nν_ss_quartz = ν_ss_quartz_array[soil_type_index]\nporosity = porosity_array[soil_type_index];\n\n\n# This tutorial additionally compares the output of a ClimateMachine simulation with that\n# of Supplemental Program 2, Chapter 5, of [Bonan19a](@cite).\n# We found this useful as it\n# allows us compare results from our code against a published version.\n\n\n# The simulation code of [Bonan19a](@cite) employs a formalism for the thermal\n# conductivity `κ` based on [Johanson1975](@cite). It assumes\n# no organic matter, and only requires the volumetric\n# fraction of soil solids for quartz and other minerals.\n# ClimateMachine employs the formalism of [BallandArp2005](@cite),\n# which requires the\n# fraction of soil solids for quartz, gravel,\n# organic matter, and other minerals. [Dai2019a](@cite) found\n# the model of [BallandArp2005](@cite) to better match\n# measured soil properties across a range of soil types.\n\n# To compare the output of the two simulations, we set the organic\n# matter content and gravel content to zero in the CliMA model.\n# The remaining soil components (quartz and other minerals) match between\n# the two. We also run the simulation for relatively wet soil\n# (water content at 80% of porosity). Under these conditions,\n# the two formulations for `κ`, though taking different functional forms,\n# are relatively consistent.\n# The differences between models are important\n# for soil with organic material and for soil that is relatively dry.\nν_ss_om = FT(0.0)\nν_ss_gravel = FT(0.0);\n\n\n# We next calculate a few intermediate quantities needed for the\n# determination of the thermal conductivity ([BallandArp2005](@cite)). These include\n# the conductivity of the solid material, the conductivity\n# of saturated soil, and the conductivity of frozen saturated soil.\n\nκ_quartz = FT(7.7) # W/m/K\nκ_minerals = FT(2.5) # W/m/K\nκ_om = FT(0.25) # W/m/K\nκ_liq = FT(0.57) # W/m/K\nκ_ice = FT(2.29); # W/m/K\n\n# The particle density of soil solids in moisture-free soil\n# is taken as a constant, across soil types, as in [Bonan19a](@cite).\n# This is a good estimate for organic material free soil. The user is referred to\n# [BallandArp2005](@cite) for a more general expression.\nρp = FT(2700) # kg/m^3\nκ_solid = k_solid(ν_ss_om, ν_ss_quartz, κ_quartz, κ_minerals, κ_om)\nκ_sat_frozen = ksat_frozen(κ_solid, porosity, κ_ice)\nκ_sat_unfrozen = ksat_unfrozen(κ_solid, porosity, κ_liq);\n\n# The thermal conductivity of dry soil is also required, but this is\n# calculated internally using the expression of [3].\n\n# The volumetric specific heat of dry soil is chosen so as to match Bonan's simulation.\n# The user could instead compute this using a volumetric fraction weighted average\n# across soil components.\nρc_ds = FT((1 - porosity) * 1.926e06) # J/m^3/K\n\n# Finally, we store the soil-specific parameters and functions\n# in a place where they will be accessible to the model\n# during integration.\nsoil_param_functions = SoilParamFunctions(\n FT;\n porosity = porosity,\n ν_ss_gravel = ν_ss_gravel,\n ν_ss_om = ν_ss_om,\n ν_ss_quartz = ν_ss_quartz,\n ρc_ds = ρc_ds,\n ρp = ρp,\n κ_solid = κ_solid,\n κ_sat_unfrozen = κ_sat_unfrozen,\n κ_sat_frozen = κ_sat_frozen,\n);\n\n\n# # Initial and Boundary conditions\n# We will be using a [`PrescribedWaterModel`](@ref\n# ClimateMachine.Land.PrescribedWaterModel), where the user supplies the augmented\n# liquid fraction and ice fraction as functions of space and time. Since we are not\n# implementing phase changes, it makes sense to either have entirely liquid or\n# frozen water. This tutorial shows liquid water.\n\n# Because the two models for thermal conductivity agree well for wetter soil, we'll\n# choose that here. However, the user could also explore how they differ by choosing\n# drier soil.\n\n# Please note that if the user uses a mix of liquid and frozen water, that they must\n# ensure that the total water content does not exceed porosity.\nprescribed_augmented_liquid_fraction = FT(porosity * 0.8)\nprescribed_volumetric_ice_fraction = FT(0.0);\n\n\n# Choose boundary and initial conditions for heat that will not lead to freezing of water:\nheat_surface_state = (aux, t) -> eltype(aux)(288.15)\nheat_bottom_flux = (aux, t) -> eltype(aux)(0.0)\nT_init = (aux) -> eltype(aux)(275.15);\n\n# The boundary value problem in this case, with two spatial derivatives,\n# requires a boundary condition at the top of the domain and the bottom.\n# Here we choose to specify a bottom flux condition, and a top state condition.\n# Our problem is effectively 1D, so we do not need to specify lateral boundary\n# conditions.\nbc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_heat = Neumann(heat_bottom_flux)),\n surface_bc = LandComponentBC(soil_heat = Dirichlet(heat_surface_state)),\n);\n\n# We also need to define a function `init_soil!`, which\n# initializes all of the prognostic variables (here, we\n# only have `ρe_int`, the volumetric internal energy).\n# The initialization is based on user-specified\n# initial conditions. Note that the user provides initial\n# conditions for heat based on the temperature - `init_soil!` also\n# converts between `T` and `ρe_int`.\n\nfunction init_soil!(land, state, aux, localgeo, time)\n param_set = parameter_set(land)\n ϑ_l, θ_i = get_water_content(land.soil.water, aux, state, time)\n θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_ds = land.soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n\n state.soil.heat.ρe_int = volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n )\nend;\n\n\n# # Create the model structure\nsoil_water_model = PrescribedWaterModel(\n (aux, t) -> prescribed_augmented_liquid_fraction,\n (aux, t) -> prescribed_volumetric_ice_fraction,\n);\n\nsoil_heat_model = SoilHeatModel(FT; initialT = T_init);\n\n# The full soil model requires a heat model and a water model, as well as the\n# soil parameter functions:\nm_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);\n\n# The equations being solved in this tutorial have no sources or sinks:\nsources = ();\n\n# Finally, we create the `LandModel`. In more complex land models, this would\n# include the canopy, carbon state of the soil, etc.\nm = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n);\n\n\n# # Specify the numerical details\n# These include the resolution, domain boundaries, integration time,\n# Courant number, and ODE solver.\n\nN_poly = 1\nnelem_vert = 100\n\nzmax = FT(0)\nzmin = FT(-1)\n\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n\n# In this tutorial, we determine a timestep based on a Courant number (\n# also called a Fourier number in the context of the heat equation).\n# In short, we can use the parameters of the model (`κ` and `ρc_s`),\n# along with with the size of\n# elements of the grid used for discretizing the PDE, to estimate\n# a natural timescale for heat transfer across a grid cell.\n# Because we are using an explicit ODE solver, the timestep should\n# be a fraction of this in order to resolve the dynamics.\n\n# This allows us to automate, to a certain extent, choosing a value for\n# the timestep, even as we switch between soil types.\n\nfunction calculate_dt(dg, model::LandModel, Q, Courant_number, t, direction)\n Δt = one(eltype(Q))\n CFL = DGMethods.courant(diffusive_courant, dg, model, Q, Δt, t, direction)\n return Courant_number / CFL\nend\n\nfunction diffusive_courant(\n m::LandModel,\n state::Vars,\n aux::Vars,\n diffusive::Vars,\n Δx,\n Δt,\n t,\n direction,\n)\n param_set = parameter_set(m)\n soil = m.soil\n ϑ_l, θ_i = get_water_content(soil.water, aux, state, t)\n θ_l = volumetric_liquid_fraction(ϑ_l, soil.param_functions.porosity)\n κ_dry = k_dry(param_set, soil.param_functions)\n S_r = relative_saturation(θ_l, θ_i, soil.param_functions.porosity)\n kersten = kersten_number(θ_i, S_r, soil.param_functions)\n κ_sat = saturated_thermal_conductivity(\n θ_l,\n θ_i,\n soil.param_functions.κ_sat_unfrozen,\n soil.param_functions.κ_sat_frozen,\n )\n κ = thermal_conductivity(κ_dry, kersten, κ_sat)\n ρc_ds = soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n return Δt * κ / (Δx * Δx * ρc_ds)\nend\n\n\nt0 = FT(0)\ntimeend = FT(60 * 60 * 3)\nCourant_number = FT(0.5) # much bigger than this leads to domain errors\n\nsolver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config;\n Courant_number = Courant_number,\n CFL_direction = VerticalDirection(),\n);\n\n\n# # Run the integration\nClimateMachine.invoke!(solver_config);\nstate_types = (Prognostic(), Auxiliary())\ndons = dict_of_nodal_states(solver_config, state_types; interp = true)\n\n# # Plot results and comparison data from [Bonan19a](@cite)\n\nz = get_z(solver_config.dg.grid; rm_dupes = true);\nT = dons[\"soil.heat.T\"];\nplot(\n T,\n z,\n label = \"ClimateMachine\",\n ylabel = \"z (m)\",\n xlabel = \"T (K)\",\n title = \"Heat transfer in sand\",\n)\nplot!(T_init.(z), z, label = \"Initial condition\")\nfilename = \"bonan_heat_data.csv\"\nbonan_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"bonan_soil_heat\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/99vm8q8tlyoulext6c35lnd3355tx6bu.csv\",\n filename = filename,\n ),],\n)\nbonan_dataset_path = get_data_folder(bonan_dataset)\ndata = joinpath(bonan_dataset_path, filename)\nds_bonan = readdlm(data, ',')\nbonan_T = reverse(ds_bonan[:, 2])\nbonan_z = reverse(ds_bonan[:, 1])\nbonan_T_continuous = Spline1D(bonan_z, bonan_T)\nbonan_at_clima_z = bonan_T_continuous.(z)\nplot!(bonan_at_clima_z, z, label = \"Bonan simulation\")\nplot!(legend = :bottomleft)\nsavefig(\"thermal_conductivity_comparison.png\")\n# ![](thermal_conductivity_comparison.png)\n\n# The plot shows that the temperature at the top of the\n# soil is gradually increasing. This is because the surface\n# temperature is held fixed at a value larger than\n# the initial temperature. If we ran this for longer,\n# we would see that the bottom of the domain would also\n# increase in temperature because there is no heat\n# leaving the bottom (due to zero heat flux specified in\n# the boundary condition).\n\n# # References\n# - [Bonan19a](@cite)\n# - [Johanson1975](@cite)\n# - [BallandArp2005](@cite)\n# - [Dai2019a](@cite)\n# - [Cosby1984](@cite)\n" }, { "path": "tutorials/Land/Soil/PhaseChange/freezing_front.jl", "content": "# # Modeling a freezing front in unsaturated soil\n\n# Before reading this tutorial, \n# we recommend that you look over the coupled energy\n# and water [tutorial](../Coupled/equilibrium_test.md).\n# That tutorial showed how to solve the heat equation for soil volumetric\n# internal energy `ρe_int`, simultaneously\n# with Richards equation for volumetric liquid water fraction `ϑ_l`, assuming zero\n# volumetric ice fraction `θ_i` for all time, everywhere in the domain[^a].\n# In this example, we add in a source term to the right hand side for both `θ_i`\n# and `ϑ_l` which models freezing and thawing and conserves water mass during the process.\n# The equations are\n\n\n# ``\n# \\frac{∂ ρe_{int}}{∂ t} = ∇ ⋅ κ(θ_l, θ_i; ν, ...) ∇T + ∇ ⋅ ρe_{int_{liq}} K (T,θ_l, θ_i; ν, ...) \\nabla h( ϑ_l, z; ν, ...)\n# ``\n\n# ``\n# \\frac{ ∂ ϑ_l}{∂ t} = ∇ ⋅ K (T,θ_l, θ_i; ν, ...) ∇h( ϑ_l, z; ν, ...) -\\frac{F_T}{ρ_l}\n# ``\n\n# ``\n# \\frac{ ∂ θ_i}{∂ t} = \\frac{F_T}{ρ_i}\n# ``\n\n# Here\n\n# ``t`` is the time (s),\n\n# ``z`` is the location in the vertical (m),\n\n# ``ρe_{int}`` is the volumetric internal energy of the soil (J/m^3),\n\n# ``T`` is the temperature of the soil (K),\n\n# ``κ`` is the thermal conductivity (W/m/K),\n\n# ``ρe_{int_{liq}}`` is the volumetric internal energy of liquid water (J/m^3),\n\n# ``K`` is the hydraulic conductivity (m/s),\n\n# ``h`` is the hydraulic head (m),\n\n# ``ϑ_l`` is the augmented volumetric liquid water fraction,\n\n# ``θ_i`` is the volumetric ice fraction,\n\n# ``ν, ...`` denotes parameters relating to soil type, such as porosity, and\n\n# ``F_T`` is the freeze-thaw term.\n\n# To begin, we will show how to implement adding in this source term. After the results are obtained,\n# we will [explain](#Discussion-and-Model-Explanation) how our model parameterizes this effect and\n# compare the results with some analytic expections.\n\n# We solve these equations in an effectively 1-d domain with ``z ∈ [-0.2,0]``,\n# and with the following boundary and initial conditions:\n\n# ``- κ ∇T(t, z = 0) = 28 W/m^2/K (T - 267.15K) ẑ``\n\n# ``- κ ∇T(t, z= -0.2) = 0 ẑ ``\n\n# `` T(t = 0, z) = 279.85 K``\n\n# ``- K ∇h(t, z = 0) = 0 ẑ ``\n\n# `` -K ∇h(t, z = -0.2) = 0 ẑ``\n\n# `` ϑ_l(t = 0, z) = 0.33 ``.\n\n# The problem setup and soil properties are chosen to match the lab experiment of [Mizoguchi1990](@cite), as detailed in [Hansson2004](@cite) and [DallAmico2011](@cite)].\n\n# # Import necessary modules\n# External (non - CliMA) modules\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\nusing DelimitedFiles\nusing Plots\n\n# CliMA Parameters\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\nusing CLIMAParameters.Planet: ρ_cloud_liq\nusing CLIMAParameters.Planet: ρ_cloud_ice\n\n# ClimateMachine modules\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\nusing ArtifactWrappers\n\n# # Preliminary set-up\n\n# Get the parameter set, which holds constants used across CliMA models:\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n# Initialize and pick a floating point precision:\nClimateMachine.init()\nconst FT = Float64;\n\n# Load plot helpers:\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n\n# # Simulation specific parameters\nN_poly = 1\nnelem_vert = 20\nzmax = FT(0)\nzmin = FT(-0.2)\nt0 = FT(0)\ndt = FT(6)\ntimeend = FT(3600 * 50)\nn_outputs = 50\nevery_x_simulation_time = ceil(Int, timeend / n_outputs)\nΔ = abs(zmin - zmax) / FT(nelem_vert);\n\n# # Soil properties.\n# All are given in mks units.\nν = FT(0.535)\nθ_r = FT(0.05)\nS_s = FT(1e-3)\nKsat = FT(3.2e-6)\nvg_α = 1.11\nvg_n = 1.48;\n\nν_ss_quartz = FT(0.7)\nν_ss_minerals = FT(0.0)\nν_ss_om = FT(0.3)\nν_ss_gravel = FT(0.0);\nκ_quartz = FT(7.7)\nκ_minerals = FT(2.4)\nκ_om = FT(0.25)\nκ_liq = FT(0.57)\nκ_ice = FT(2.29);\nκ_solid = k_solid(ν_ss_om, ν_ss_quartz, κ_quartz, κ_minerals, κ_om)\nκ_sat_frozen = ksat_frozen(κ_solid, ν, κ_ice)\nκ_sat_unfrozen = ksat_unfrozen(κ_solid, ν, κ_liq);\nρp = FT(3200)\nρc_ds = FT((1 - ν) * 2.3e6);\n\n\nsoil_param_functions = SoilParamFunctions(\n FT;\n porosity = ν,\n ν_ss_gravel = ν_ss_gravel,\n ν_ss_om = ν_ss_om,\n ν_ss_quartz = ν_ss_quartz,\n ρc_ds = ρc_ds,\n ρp = ρp,\n κ_solid = κ_solid,\n κ_sat_unfrozen = κ_sat_unfrozen,\n κ_sat_frozen = κ_sat_frozen,\n water = WaterParamFunctions(FT; Ksat = Ksat, S_s = S_s, θ_r = θ_r),\n);\n\n\n# # Build the model\n# Initial and Boundary conditions. The default initial condition for\n# `θ_i` is zero everywhere, so we don't modify that. Furthermore, since\n# the equation for `θ_i` does not involve spatial derivatives, we don't need\n# to supply boundary conditions for it. Note that Neumann fluxes, when chosen,\n# are specified by giving the magnitude of the normal flux *into* the domain.\n# In this case, the normal vector at the surface n̂ = ẑ. Internally, we multiply\n# the flux magnitude by -n̂.\nzero_flux = (aux, t) -> eltype(aux)(0.0)\nsurface_heat_flux =\n (aux, t) -> eltype(aux)(-28) * (aux.soil.heat.T - eltype(aux)(273.15 - 6))\nT_init = aux -> eltype(aux)(279.85)\nϑ_l0 = (aux) -> eltype(aux)(0.33);\n\nbc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_heat = Neumann(zero_flux),\n soil_water = Neumann(zero_flux),\n ),\n surface_bc = LandComponentBC(\n soil_heat = Neumann(surface_heat_flux),\n soil_water = Neumann(zero_flux),\n ),\n);\n\n# Create the [`SoilWaterModel`](@ref ClimateMachine.Land.SoilWaterModel),\n# [`SoilHeatModel`](@ref ClimateMachine.Land.SoilHeatModel),\n# and the [`SoilModel`](@ref ClimateMachine.Land.SoilModel) instances.\n# Note that we are allowing for the hydraulic conductivity to be affected by\n# both temperature and ice fraction by choosing the following\n# [`viscosity_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.viscosity_factor)\n# and [`impedance_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.impedance_factor).\n# To turn these off - the default - just remove these lines. These factors are explained more\n# [here](../Water/hydraulic_functions.md).\nsoil_water_model = SoilWaterModel(\n FT;\n viscosity_factor = TemperatureDependentViscosity{FT}(),\n moisture_factor = MoistureDependent{FT}(),\n impedance_factor = IceImpedance{FT}(Ω = 7.0),\n hydraulics = vanGenuchten(FT; α = vg_α, n = vg_n),\n initialϑ_l = ϑ_l0,\n)\n\nsoil_heat_model = SoilHeatModel(FT; initialT = T_init);\n\nm_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);\n\n# Create the source term instance. Our phase change model requires\n# knowledge of the vertical spacing, so we pass\n# that information in via an attribute of the\n# [`PhaseChange`](@ref ClimateMachine.Land.PhaseChange) structure.\nfreeze_thaw_source = PhaseChange{FT}(Δz = Δ);\n\n# Sources are added as elements of a list of sources. Here we just add freezing\n# and thawing.\nsources = (freeze_thaw_source,);\n\n# Next, we define the required `init_soil!` function, which takes the user\n# specified functions of space for `T_init` and `ϑ_l0` and initializes the state\n# variables of volumetric internal energy and augmented liquid fraction. This requires\n# a conversion from `T` to `ρe_int`.\nfunction init_soil!(land, state, aux, localgeo, time)\n myFT = eltype(state)\n ϑ_l = myFT(land.soil.water.initialϑ_l(aux))\n θ_i = myFT(land.soil.water.initialθ_i(aux))\n state.soil.water.ϑ_l = ϑ_l\n state.soil.water.θ_i = θ_i\n param_set = land.param_set\n\n θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_ds = land.soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n\n state.soil.heat.ρe_int = volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n )\nend;\n\n# Lastly, package it all up in the `LandModel`:\nm = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n);\n\n# # Build the simulation domain, solver, and callbacks\n\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\ndg = solver_config.dg\nQ = solver_config.Q\nstate_types = (Prognostic(), Auxiliary(), GradientFlux())\ndons_arr = Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\ntime_data = FT[0]\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Run the simulation, and plot the output\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,));\n\nz = get_z(solver_config.dg.grid; rm_dupes = true);\n\noutput_dir = @__DIR__;\n\nmkpath(output_dir);\n\nexport_plot(\n z,\n time_data[[1, 16, 31, 46]] ./ (60 * 60),\n dons_arr[[1, 16, 31, 46]],\n (\"soil.water.ϑ_l\",),\n joinpath(output_dir, \"moisture_plot.png\");\n xlabel = \"ϑ_l\",\n ylabel = \"z (m)\",\n time_units = \"hrs \",\n)\n# ![](moisture_plot.png)\n\nexport_plot(\n z,\n time_data[[1, 16, 31, 46]] ./ (60 * 60),\n dons_arr[[1, 16, 31, 46]],\n (\"soil.water.θ_i\",),\n joinpath(output_dir, \"ice_plot.png\");\n xlabel = \"θ_i\",\n ylabel = \"z (m)\",\n time_units = \"hrs \",\n legend = :bottomright,\n)\n# ![](ice_plot.png)\nexport_plot(\n z,\n time_data[[1, 16, 31, 46]] ./ (60 * 60),\n dons_arr[[1, 16, 31, 46]],\n (\"soil.heat.T\",),\n joinpath(output_dir, \"T_plot.png\");\n xlabel = \"T (K)\",\n ylabel = \"z (m)\",\n time_units = \"hrs \",\n)\n# ![](T_plot.png)\n\n# # Comparison to data\n# This data was obtained by us from the figures of [Hansson2004](@cite), but was originally obtained\n# by [Mizoguchi1990](@cite). No error bars were reported, and we haven't quantified the error in our\n# estimation of the data from images.\ndataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"mizoguchi\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/3xbo4rlam8u390vmucc498cao6wmqlnd.csv\",\n filename = \"mizoguchi_all_data.csv\",\n ),],\n);\ndataset_path = get_data_folder(dataset);\ndata = joinpath(dataset_path, \"mizoguchi_all_data.csv\")\nds = readdlm(data, ',')\nhours = ds[:, 1][2:end]\nvwc = ds[:, 2][2:end] ./ 100.0\ndepth = ds[:, 3][2:end]\nmask_12h = hours .== 12\nmask_24h = hours .== 24\nmask_50h = hours .== 50;\n\nplot_12h =\n scatter(vwc[mask_12h], -depth[mask_12h], label = \"\", color = \"purple\")\nplot!(\n dons_arr[13][\"soil.water.θ_i\"] .+ dons_arr[13][\"soil.water.ϑ_l\"],\n z,\n label = \"\",\n color = \"green\",\n)\nplot!(title = \"12h\")\nplot!(xlim = [0.2, 0.55])\nplot!(xticks = [0.2, 0.3, 0.4, 0.5])\nplot!(ylabel = \"Depth (m)\");\n\nplot_24h =\n scatter(vwc[mask_24h], -depth[mask_24h], label = \"Data\", color = \"purple\")\nplot!(\n dons_arr[25][\"soil.water.θ_i\"] .+ dons_arr[25][\"soil.water.ϑ_l\"],\n z,\n label = \"Simulation\",\n color = \"green\",\n)\nplot!(title = \"24h\")\nplot!(legend = :bottomright)\nplot!(xlim = [0.2, 0.55])\nplot!(xticks = [0.2, 0.3, 0.4, 0.5]);\n\nplot_50h =\n scatter(vwc[mask_50h], -depth[mask_50h], label = \"\", color = \"purple\")\nplot!(\n dons_arr[51][\"soil.water.θ_i\"] .+ dons_arr[51][\"soil.water.ϑ_l\"],\n z,\n label = \"\",\n color = \"green\",\n)\nplot!(title = \"50h\")\nplot!(xlim = [0.2, 0.55])\nplot!(xticks = [0.2, 0.3, 0.4, 0.5]);\n\nplot(plot_12h, plot_24h, plot_50h, layout = (1, 3))\nplot!(xlabel = \"θ_l+θ_i\")\nsavefig(\"mizoguchi_data_comparison.png\")\n# ![](mizoguchi_data_comparison.png)\n\n# # Discussion and Model Explanation\n\n# To begin, let's observe that the freeze thaw source term alone conserves water mass, as\n# it satisfies\n\n# ``\n# ρ_l \\partial_tϑ_l + ρ_i \\partial_tθ_i = -F_T + F_T = 0\n# ``\n\n# Next, we describe how we define `F_T`.\n# The Clausius-Clapeyron (CC) equation defines a pressure-temperature curve along which two\n# phases can co-exist. It assumes that the phases are at equal temperature and pressures.\n# For water in soil, however, the liquid water experiences pressure `ρ_l g ψ`, where\n# `ψ` is the matric potential. A more general form of the CC equation allows for different\n# pressures in the two phases. Usually the ice pressure is taken to be zero, which is reasonable\n# for unsaturated freezing soils. In saturated soils, freezing can lead to heaving of the soil which\n# we do not model. After that assumption is made, we obtain that, below freezing (``T < T_f``)\n\n# ``\n# \\frac{dp_l}{ρ_l} = L_f \\frac{dT}{T},\n# ``\n\n# or\n\n# ``\n# p_l = p_{l,0} + L_f ρ_l \\frac{T-T_f}{T_f} \\mathcal{H}(T_f-T)\n# ``\n\n# where we have assumed that assumed `T` is near the freezing point, and then\n# performed a Taylor explansion of the logarithm,\n# and we are ignoring the freezing point depression, which is small (less than one degree) for\n# non-clay soils. What we have sketched is further explained in [DallAmico2011](@cite) and [KurylykWatanabe2013](@cite).\n\n# What this implies is that above the freezing point, the pressure is equal to ``p_{l,0}``,\n# which is independent of temperature. Once the temperature drops below the freezing point,\n# the pressure drops. Since prior to freezing, the pressure ``p_{l,0}`` is equal to\n# `ρ_l g ψ(θ_l)`, water undergoing freezing alone (without flowing) should satisfy ([DallAmico2011](@cite)):\n\n# ``\n# p_{l,0} = ρ_l g ψ(θ_l+ρ_iθ_i/ρ_l)\n# ``\n\n# where `ψ` is the matric potential function of van Genuchten. At each step, we know both\n# the water and ice contents, as well as the temperature, and can then solve for\n\n# ``\n# θ_{l}^* = (ν-θ_r) ψ^{-1}(p_l/(ρ_l g)) + θ_r.\n# ``\n\n# For freezing, the freeze thaw function `F_T` is equal to\n\n# ``\n# F_T = \\frac{1}{τ} ρ_l (θ_l-θ_{l}^*) \\mathcal{H}(T_f-T) \\mathcal{H}(θ_l-θ_{l}^*)\n# ``\n\n# which brings the `θ_l` to a value which satisfies `p_l = ρ_l g ψ(θ_l)`.\n# This is why, in our simulation, we see the liquid\n# water fraction approaches a constant around 0.075 in the frozen region, rather than the residual fraction\n# of 0.019, or 0. This behavior is observed, for example, in the experiments of [Watanabe2011](@cite).\n\n# Although this approach may indicate that we should replace the pressure head appearing in the\n# diffusive water flux term in Richards equation ([DallAmico2011](@cite)), we do not do so at present. As such, we may not be modeling\n# the flow of water around the freezing front properly. However, we still observe cryosuction, which\n# is the flow of water towards the freezing front, from the unfrozen side. As the water freezes, the liquid\n# water content drops,\n# setting up a larger gradient in matric potential across the freezing front, which generates upward flow\n# against gravity. This is evident because the total water content at the top is larger at the end of the\n# simulation\n# than it was at `t=0` (when it was 0.33).\n\n# This model differs from others (e.g. [Painter2011](@cite), [Hansson2004](@cite), [DallAmico2011](@cite)) in that it requires us to set a timescale for the phase change, `τ`.\n# In a first-order\n# phase transition, the temperature is fixed while the necessary latent heat is either lost or gained by the\n# system. Ignoring\n# changes in internal energy due to flowing water, we would expect\n\n\n# ``\n# \\partial_t ρe_{int} \\approx (ρ_l c_l \\partial_t θ_l + ρ_i c_i \\partial_t θ_i) (T-T_0) -ρ_i L_f \\partial_t θ_i\n# ``\n\n# ``\n# = [(c_i-c_l) (T-T_0) -L_f]F_T \\approx -L_f F_T\n# ``\n\n# or\n\n# ``\n# F_T ∼ \\frac{κ|∇²T|}{L_f} ∼\\frac{κ}{c̃ Δz²}\\frac{c̃ |∂zT| Δz}{L_f}\n# ``\n\n# suggesting\n\n# ``\n# τ ∼ τ_{LTE}\\frac{ρ_lL_f (ν-θ_r)}{c̃ |∂zT| Δz}\n# ``\n\n# with\n\n# ``\n# τ_{LTE}= c̃ Δz²/κ\n# ``\n\n# This is the value we use. This seems to work adequately for modeling freezing front propagation and\n# cryosuction, via comparisons with [Mizoguchi1990](@cite), but we plan to revisit it in the future. For example,\n# we do not see a strong temperature plateau at the freezing point ([Watanabe2011](@cite)),\n# which we would expect while the phase change is occuring. Experimentally, this timescale also affects the abruptness of the freezing front, which our simulation softens.\n\n# # References\n# - [Mizoguchi1990](@cite)\n# - [Hansson2004](@cite)\n# - [DallAmico2011](@cite)\n# - [KurylykWatanabe2013](@cite)\n# - [Watanabe2011](@cite)\n# - [Painter2011](@cite)\n\n# [^a]:\n# Note that `θ_i` is always treated as a prognostic variable\n# in the `SoilWaterModel`, but with\n# zero terms on the RHS unless freezing and thawing is turn on, as demonstrated in this\n# tutorial. That means that the user could, in principle, set the initial condition to be nonzero\n# (`θ_i(x, y, z ,t=0) = 0` is the default), which in turn would allow a nonzero `θ_i`\n# profile to affect things like thermal conductivity, etc,\n# in a consistent way. However, it would not be enforced that ``θ_l+θ_i \\leq ν``, because there would\n# be no physics linking the liquid and water content to each other, and they are independent\n# variables in our model. We don't envision this being a common use case.\n" }, { "path": "tutorials/Land/Soil/PhaseChange/phase_change_analytic_test.jl", "content": "# # Comparison to Neumann analytic solution\n\n# Before reading this tutorial, \n# we recommend that you look over the coupled energy\n# and water [tutorial](../Coupled/equilibrium_test.md)\n# and the freezing front [tutorial](freezing_front.md).\n# The former shows how to solve the heat equation for soil volumetric\n# internal energy `ρe_int` simultaneously\n# with Richards equation for volumetric liquid water fraction `ϑ_l`, assuming zero\n# volumetric ice fraction `θ_i` for all time, everywhere in the domain[^a].\n# The latter shows how to include freezing and thawing, and explains the freeze thaw model employed\n# by CliMA Land. This tutorial compares a simulated temperature profile\n# from a freezing front with an analytic solution.\n\n# The analytic solution applies to a freezing front propagating under certain assumptions. It assumes\n# that there is no water movement, which we mimic by setting `K_sat=0`. It also assumes a semi-infinite\n# domain, which we approximate by making the domain larger than the extent to which the front propagates.\n# The solution also assumes that all water freezes. Our model does not satisfy that assumption, as discussed\n# [here](freezing_front.md#Discussion-and-Model-Explanation), but as we will see, the model still matches the\n# analytic expectation well in the frozen region.\n\n# As such, our set of equations is\n\n\n# ``\n# \\frac{∂ ρe_{int}}{∂ t} = ∇ ⋅ κ(θ_l, θ_i; ν, ...) ∇T \n# ``\n\n# ``\n# \\frac{ ∂ ϑ_l}{∂ t} = -\\frac{F_T}{ρ_l}\n# ``\n\n# ``\n# \\frac{ ∂ θ_i}{∂ t} = \\frac{F_T}{ρ_i}\n# ``\n\n# Here\n\n# ``t`` is the time (s),\n\n# ``z`` is the location in the vertical (m),\n\n# ``ρe_{int}`` is the volumetric internal energy of the soil (J/m^3),\n\n# ``T`` is the temperature of the soil (K),\n\n# ``κ`` is the thermal conductivity (W/m/K),\n\n# ``ϑ_l`` is the augmented volumetric liquid water fraction,\n\n# ``θ_i`` is the volumetric ice fraction,\n\n# ``ν, ...`` denotes parameters relating to soil type, such as porosity, and\n\n# ``F_T`` is the freeze-thaw term.\n\n# Our domain is effectively 1-d, with ``z ∈ [-3,0]``,\n# and with the following boundary and initial conditions:\n\n# `` T(t, z=0) = 263.15 K``\n\n# ``- κ ∇T(t, z= -3) = 0 ẑ ``\n\n# `` T(t = 0, z) = 275.15 K``\n\n# ``- K ∇h(t, z = 0) = 0 ẑ ``\n\n# `` -K ∇h(t, z = -3) = 0 ẑ``\n\n# `` ϑ(t = 0, z) = 0.33 ``.\n\n# # Import necessary modules\n# External (non - CliMA) modules\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\nusing Test\nusing DelimitedFiles\nusing Plots\n\n# CliMA Parameters\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\nusing CLIMAParameters.Planet: ρ_cloud_liq\nusing CLIMAParameters.Planet: ρ_cloud_ice\nusing CLIMAParameters.Planet: LH_f0\n\n# ClimateMachine modules\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Land.SoilHeatParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.SystemSolvers\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\nusing SpecialFunctions\nusing ArtifactWrappers\n\n# # Preliminary set-up\n\n# Get the parameter set, which holds constants used across CliMA models:\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n# Initialize and pick a floating point precision:\nClimateMachine.init()\nconst FT = Float64;\n\n# # Simulation specific parameters\nN_poly = 1\nnelem_vert = 40\nzmax = FT(0)\nzmin = FT(-3)\nt0 = FT(0)\ndt = FT(50)\ntimeend = FT(3600 * 24 * 20)\nn_outputs = 540\nevery_x_simulation_time = ceil(Int, timeend / n_outputs)\nΔ = abs(zmin - zmax) / FT(nelem_vert);\n\n# # Soil properties\n# All units are mks.\nporosity = FT(0.535)\nvg_α = 1.11\nvg_n = 1.48\nKsat = 0.0\nν_ss_quartz = FT(0.2)\nν_ss_minerals = FT(0.6)\nν_ss_om = FT(0.2)\nν_ss_gravel = FT(0.0);\nκ_quartz = FT(7.7)\nκ_minerals = FT(2.5)\nκ_om = FT(0.25)\nκ_liq = FT(0.57)\nκ_ice = FT(2.29)\nρp = FT(2700)\nκ_solid = k_solid(ν_ss_om, ν_ss_quartz, κ_quartz, κ_minerals, κ_om)\nκ_sat_frozen = ksat_frozen(κ_solid, porosity, κ_ice)\nκ_sat_unfrozen = ksat_unfrozen(κ_solid, porosity, κ_liq)\nρc_ds = FT((1 - porosity) * 2.3e6)\nsoil_param_functions = SoilParamFunctions(\n FT;\n porosity = porosity,\n ν_ss_gravel = ν_ss_gravel,\n ν_ss_om = ν_ss_om,\n ν_ss_quartz = ν_ss_quartz,\n ρc_ds = ρc_ds,\n ρp = ρp,\n κ_solid = κ_solid,\n κ_sat_unfrozen = κ_sat_unfrozen,\n κ_sat_frozen = κ_sat_frozen,\n water = WaterParamFunctions(FT; Ksat = Ksat, S_s = 1e-3),\n);\n\n# # Build the model\n# Initial and Boundary conditions. The default initial condition for\n# `θ_i` is zero everywhere, so we don't modify that. Furthermore, since\n# the equation for `θ_i` does not involve spatial derivatives, we don't need\n# to supply boundary conditions for it.\nzero_flux = (aux, t) -> eltype(aux)(0.0)\nϑ_l0 = (aux) -> eltype(aux)(0.33)\nsurface_state = (aux, t) -> eltype(aux)(273.15 - 10.0)\nT_init = aux -> eltype(aux)(275.15)\nbc = LandDomainBC(\n bottom_bc = LandComponentBC(\n soil_heat = Neumann(zero_flux),\n soil_water = Neumann(zero_flux),\n ),\n surface_bc = LandComponentBC(\n soil_heat = Dirichlet(surface_state),\n soil_water = Neumann(zero_flux),\n ),\n);\n\n# Create the [`SoilWaterModel`](@ref ClimateMachine.Land.SoilWaterModel),\n# [`SoilHeatModel`](@ref ClimateMachine.Land.SoilHeatModel),\n# and the [`SoilModel`](@ref ClimateMachine.Land.SoilModel) instances.\n# Note that we are still specifying a hydraulics model, because the matric potential\n# and hydraulic conductivity functions are still evaluated (though they don't affect\n# the outcome). Setting `Ksat =0` is just a\n# hack for turning off water flow. \nsoil_water_model = SoilWaterModel(\n FT;\n hydraulics = vanGenuchten(FT; α = vg_α, n = vg_n),\n initialϑ_l = ϑ_l0,\n);\n\n\nsoil_heat_model = SoilHeatModel(FT; initialT = T_init);\n\nm_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);\n# Create the source term instance. Our phase change model requires\n# knowledge of the vertical spacing, so we pass\n# that information in via an attribute of the\n# [`PhaseChange`](@ref ClimateMachine.Land.PhaseChange) structure.\nfreeze_thaw_source = PhaseChange{FT}(Δz = Δ);\n\n# Sources are added as elements of a list of sources. Here we just add freezing\n# and thawing.\nsources = (freeze_thaw_source,);\n\n# Next, we define the required `init_soil!` function, which takes the user\n# specified functions of space for `T_init` and `ϑ_l0` and initializes the state\n# variables of volumetric internal energy and augmented liquid fraction. This requires\n# a conversion from `T` to `ρe_int`.\nfunction init_soil!(land, state, aux, localgeo, time)\n myFT = eltype(state)\n ϑ_l = myFT(land.soil.water.initialϑ_l(aux))\n θ_i = myFT(land.soil.water.initialθ_i(aux))\n state.soil.water.ϑ_l = ϑ_l\n state.soil.water.θ_i = θ_i\n param_set = land.param_set\n\n θ_l = volumetric_liquid_fraction(ϑ_l, land.soil.param_functions.porosity)\n ρc_ds = land.soil.param_functions.ρc_ds\n ρc_s = volumetric_heat_capacity(θ_l, θ_i, ρc_ds, param_set)\n\n state.soil.heat.ρe_int = volumetric_internal_energy(\n θ_i,\n ρc_s,\n land.soil.heat.initialT(aux),\n param_set,\n )\nend;\n\n# Lastly, package it all up in the `LandModel`:\nm = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil!,\n);\n\n# # Set up and run the simulation\n\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\nstate_types = (Prognostic(), Auxiliary(), GradientFlux())\nall_data = Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\ntime_data = FT[0]\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(all_data, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,));\nz = get_z(solver_config.dg.grid; rm_dupes = true);\n\n# # Analytic Solution of Neumann\n# All details here are taken from [DallAmico2011](@cite) (see also [CarslawJaeger](@cite)), and the reader is referred to that\n# for further information on the solution. It takes the form of a function\n# for T(z) on each side of the freezing front interface, which depends on\n# the thermal properties in that region, and which is also parameterized by\n# a parameter (ζ), which we show how to solve for below. In computing the\n# thermal properties, we evaluate the conductivity and heat capacity\n# assuming that all of the water is either in liquid or frozen form, with total\n# mass proportional to ``θ_{l,0}ρ_l`` (as we have no water flow).\n\n# Compute the thermal conductivity and heat capacity in the frozen region - subscript 1.\nθ_l0 = FT(0.33)\nkdry = k_dry(param_set, soil_param_functions)\nksat = saturated_thermal_conductivity(\n FT(0.0),\n θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),\n κ_sat_unfrozen,\n κ_sat_frozen,\n)\nkersten = kersten_number(\n θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),\n θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set) / porosity,\n soil_param_functions,\n)\nλ1 = thermal_conductivity(kdry, kersten, ksat)\nc1 = volumetric_heat_capacity(\n FT(0.0),\n θ_l0 * ρ_cloud_liq(param_set) / ρ_cloud_ice(param_set),\n ρc_ds,\n param_set,\n)\nd1 = λ1 / c1;\n\n# Compute the thermal conductivity and heat capacity in the region\n# with liquid water - subscript 2.\nksat =\n saturated_thermal_conductivity(θ_l0, FT(0.0), κ_sat_unfrozen, κ_sat_frozen)\nkersten = kersten_number(FT(0.0), θ_l0 / porosity, soil_param_functions)\nλ2 = thermal_conductivity(kdry, kersten, ksat)\nc2 = volumetric_heat_capacity(θ_l0, FT(0.0), ρc_ds, param_set)\nd2 = λ2 / c2;\n\n# Initial T and surface T, in Celsius\nTi = FT(2)\nTs = FT(-10.0);\n\n# The solution requires the root of the implicit equation below\n# (you'll need to install `Roots` via the package manager).\n# ``` julia\n# using Roots\n# function implicit(ζ)\n# term1 = exp(-ζ^2) / ζ / erf(ζ)\n# term2 =\n# -λ2 * sqrt(d1) * (Ti - 0) /\n# (λ1 * sqrt(d2) * (0 - Ts) * ζ * erfc(ζ * sqrt(d1 / d2))) *\n# exp(-d1 / d2 * ζ^2)\n# term3 =\n# -LH_f0(param_set) * ρ_cloud_liq(param_set) * θ_l0 * sqrt(π) / c1 /\n# (0 - Ts)\n# return (term1 + term2 + term3)\n# end\n# find_zero(implicit, (0.25,0.27), Bisection())\n# ```\n\n# The root is\nζ = 0.26447353269809687;\n\n\n\n# This function plots the analytic solution\n# and the simulated result, at the output time index\n# `k`.\nfunction f(k; ζ = ζ)\n T = all_data[k][\"soil.heat.T\"][:] .- 273.15\n plot(\n T,\n z[:],\n xlim = [-10.5, 3],\n ylim = [zmin, zmax],\n xlabel = \"T\",\n label = \"simulation\",\n )\n t = time_data[k]\n zf = 2.0 * ζ * sqrt(d1 * t)\n myz = zmin:0.001:0\n spatially_varying =\n (erfc.(abs.(myz) ./ (zf / ζ / (d1 / d2)^0.5))) ./\n erfc(ζ * (d1 / d2)^0.5)\n mask = abs.(myz) .>= zf\n plot!(\n Ti .- (Ti - 0.0) .* spatially_varying[mask],\n myz[mask],\n label = \"analytic\",\n color = \"green\",\n )\n spatially_varying = ((erf.(abs.(myz) ./ (zf / ζ)))) ./ erf(ζ)\n mask = abs.(myz) .< zf\n plot!(\n Ts .+ (0.0 - Ts) .* spatially_varying[mask],\n myz[mask],\n label = \"\",\n color = \"green\",\n )\nend\n\n# Now we will plot this and compare to other methods for modeling phase change without water movement.\n# These solutions were produced by modifying the Supplemental Program 5:3 from [Bonan19a](@cite).\n\n# Excess heat solution:\neh_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"eh\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/6xs1r98wk7u1b0xjhpkdvt80q4sh16un.csv\",\n filename = \"bonan_data.csv\",\n ),],\n);\neh_dataset_path = get_data_folder(eh_dataset);\neh_data = joinpath(eh_dataset_path, \"bonan_data.csv\")\nds = readdlm(eh_data, ',');\n# Apparent heat capacity solution:\nahc_dataset = ArtifactWrapper(\n @__DIR__,\n isempty(get(ENV, \"CI\", \"\")),\n \"ahc\",\n ArtifactFile[ArtifactFile(\n url = \"https://caltech.box.com/shared/static/d6xciskzl2djwi03xxfo8wu4obrptjmy.csv\",\n filename = \"bonan_data_ahc.csv\",\n ),],\n);\nahc_dataset_path = get_data_folder(ahc_dataset);\nahc_data = joinpath(ahc_dataset_path, \"bonan_data_ahc.csv\")\nds2 = readdlm(ahc_data, ',');\n\nk = n_outputs\nf(k)\nplot!(ds[:, 2], ds[:, 1], label = \"Excess Heat\")\nplot!(ds2[:, 2], ds2[:, 1], label = \"Apparent heat capacity\")\nplot!(legend = :bottomleft)\nsavefig(\"analytic_comparison.png\")\n# ![](analytic_comparison.png)\n\n# # References\n# - [DallAmico2011](@cite)\n# - [CarslawJaeger](@cite)\n# - [Bonan19a](@cite)\n\n# [^a]:\n# Note that `θ_i` is always treated as a prognostic variable\n# in the `SoilWaterModel`, but with\n# zero terms on the RHS unless freezing and thawing is turn on, as demonstrated in this\n# tutorial. That means that the user could, in principle, set the initial condition to be nonzero\n# (`θ_i(x, y, z ,t=0) = 0` is the default), which in turn would allow a nonzero `θ_i`\n# profile to affect things like thermal conductivity, etc,\n# in a consistent way. However, it would not be enforced that ``θ_l+θ_i \\leq ν``, because there would\n# be no physics linking the liquid and water content to each other, and they are independent\n# variables in our model. We don't envision this being a common use case.\n" }, { "path": "tutorials/Land/Soil/Water/equilibrium_test.jl", "content": "# # Hydrostatic Equilibrium test for Richards Equation\n\n# This tutorial shows how to use `ClimateMachine` code to solve\n# Richards equation in a column of soil. We choose boundary\n# conditions of zero flux at the top and bottom of the column,\n# and then run the simulation long enough to see that the system\n# is approaching hydrostatic equilibrium, where the gradient of the\n# pressure head is equal and opposite the gradient of the\n# gravitational head. Note that the [`SoilWaterModel`](@ref\n# ClimateMachine.Land.SoilWaterModel) includes\n# a prognostic equation for the volumetric ice fraction,\n# as ice is a form of water that must be accounted for to1 ensure\n# water mass conservation. If freezing and thawing are not turned on\n# (the default), the amount of ice in the model is zero for all space and time\n# (again by default). \n\n# The equations are:\n\n# ``\n# \\frac{ ∂ ϑ_l}{∂ t} = ∇ ⋅ K (T, ϑ_l, θ_i; ν, ...) ∇h( ϑ_l, z; ν, ...).\n# ``\n\n# ``\n# \\frac{ ∂ θ_i}{∂ t} = 0\n# ``\n\n# Here\n\n# ``t`` is the time (s),\n\n# ``z`` is the location in the vertical (m),\n\n# ``T`` is the temperature of the soil (K),\n\n# ``K`` is the hydraulic conductivity (m/s),\n\n# ``h`` is the hydraulic head (m),\n\n# ``ϑ_l`` is the augmented volumetric liquid water fraction,\n\n# ``θ_i`` is the volumetric ice fraction, and\n\n# ``ν, ...`` denotes parameters relating to soil type, such as porosity.\n\n\n# We will solve this equation in an effectively 1-d domain with ``z ∈ [-10,0]``,\n# and with the following boundary and initial conditions:\n\n# ``- K ∇h(t, z = 0) = 0 ẑ ``\n\n# `` -K ∇h(t, z = -10) = 0 ẑ``\n\n# `` ϑ(t = 0, z) = ν-0.001 ``\n\n# `` θ_i(t = 0, z) = 0.0. ``\n\n# where ``\\nu`` is the porosity.\n\n# A word about the hydraulic conductivity: please see the\n# [`hydraulic functions`](./hydraulic_functions.md) tutorial\n# for options regarding this function. The user can choose to make it depend\n# on the temperature and the amount of ice in the soil; the default, which we use\n# here, is that `K` only depends on the liquid moisture content.\n\n# Lastly, our formulation of this equation allows for a continuous solution in both\n# saturated and unsaturated areas, following [Woodward00a](@cite).\n\n# # Preliminary setup\n\n# - Load external packages\n\nusing MPI\nusing OrderedCollections\nusing StaticArrays\nusing Statistics\n\n# - Load CLIMAParameters and ClimateMachine modules\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.DGMethods: BalanceLaw, LocalGeometry\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux, vars_state\n\n\n# - Define the float type desired (`Float64` or `Float32`)\nconst FT = Float64;\n\n# - Initialize ClimateMachine for CPU\nClimateMachine.init(; disable_gpu = true);\n\n# Load plot helpers:\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\n# # Set up the soil model\n\n# We want to solve Richards equation alone, without simultaneously\n# solving the heat equation. Because of that, we choose a\n# [`PrescribedTemperatureModel`](@ref\n# ClimateMachine.Land.PrescribedTemperatureModel).\n# The user can supply a function for temperature,\n# depending on time and space; if this option is desired, one could also\n# choose to model the temperature dependence of viscosity, or to drive\n# a freeze/thaw cycle, for example. If the user simply wants to model\n# Richards equation for liquid water, the defaults will allow for that.\n# Here we ignore the effects of temperature and freezing and thawing,\n# using the defaults.\n\nsoil_heat_model = PrescribedTemperatureModel();\n\n# Define the porosity, Ksat, and specific storage values for the soil. Note\n# that all values must be given in mks units. The soil parameters chosen\n# roughly correspond to Yolo light clay, and are stored in \n# [`SoilParamFunctions`](@ref ClimateMachine.Land.SoilParamFunctions).\n# Hydrology specific parameters are further organized and stored in\n# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions),\n# with the exception of the hydraulic model and hydraulic conductivity model -\n# see the [`hydraulics`](./hydraulic_functions.md) tutorial.\nwpf = WaterParamFunctions(FT; Ksat = 0.0443 / (3600 * 100), S_s = 1e-3);\nsoil_param_functions = SoilParamFunctions(FT; porosity = 0.495, water = wpf);\n\n# Define the boundary conditions. The user can specify two conditions,\n# either at the top or at the bottom, and they can either be Dirichlet\n# (on `ϑ_l`) or Neumann (on `-K∇h`). Note that fluxes are supplied as\n# scalars, inside the code they are multiplied by ẑ.\n\nsurface_flux = (aux, t) -> eltype(aux)(0.0)\nbottom_flux = (aux, t) -> eltype(aux)(0.0);\n\n# Our problem is effectively 1D, so we do not need to specify lateral boundary\n# conditions.\nbc = LandDomainBC(\n bottom_bc = LandComponentBC(soil_water = Neumann(bottom_flux)),\n surface_bc = LandComponentBC(soil_water = Neumann(surface_flux)),\n);\n\n# Define the initial state function. The default for `θ_i` is zero.\nϑ_l0 = (aux) -> eltype(aux)(0.494);\n\n# Create the SoilWaterModel. The defaults are a temperature independent\n# viscosity, and no impedance factor due to ice. We choose to make the\n# hydraulic conductivity a function of the moisture content `ϑ_l`,\n# and employ the vanGenuchten hydraulic model with `n` = 2.0. The van\n# Genuchten parameter `m` is calculated from `n`, and we use the default\n# value for `α`.\nsoil_water_model = SoilWaterModel(\n FT;\n moisture_factor = MoistureDependent{FT}(),\n hydraulics = vanGenuchten(FT; n = 2.0),\n initialϑ_l = ϑ_l0,\n);\n\n# Create the soil model - the coupled soil water and soil heat models.\nm_soil = SoilModel(soil_param_functions, soil_water_model, soil_heat_model);\n\n# We are ignoring sources and sinks here, like runoff or freezing and thawing.\nsources = ();\n\n# Define the function that initializes the prognostic variables. This\n# in turn calls the functions supplied to `soil_water_model`.\nfunction init_soil_water!(land, state, aux, localgeo, time)\n state.soil.water.ϑ_l = eltype(state)(land.soil.water.initialϑ_l(aux))\n state.soil.water.θ_i = eltype(state)(land.soil.water.initialθ_i(aux))\nend\n\n\n# Create the land model - in this tutorial, it only includes the soil.\nm = LandModel(\n param_set,\n m_soil;\n boundary_conditions = bc,\n source = sources,\n init_state_prognostic = init_soil_water!,\n);\n\n# # Specify the numerical configuration and output data.\n\n# Specify the polynomial order and vertical resolution.\nN_poly = 2;\nnelem_vert = 20;\n\n# Specify the domain boundaries.\nzmax = FT(0);\nzmin = FT(-10);\n\n# Create the driver configuration.\ndriver_config = ClimateMachine.SingleStackConfiguration(\n \"LandModel\",\n N_poly,\n nelem_vert,\n zmax,\n param_set,\n m;\n zmin = zmin,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n);\n\n# Choose the initial and final times, as well as a timestep.\nt0 = FT(0)\ntimeend = FT(60 * 60 * 24 * 36)\ndt = FT(100);\n\n# Create the solver configuration.\nsolver_config =\n ClimateMachine.SolverConfiguration(t0, timeend, driver_config, ode_dt = dt);\n\n# Determine how often you want output.\nconst n_outputs = 6;\n\nconst every_x_simulation_time = ceil(Int, timeend / n_outputs);\n\n# Create a place to store this output.\nstate_types = (Prognostic(), Auxiliary(), GradientFlux())\ndons_arr = Dict[dict_of_nodal_states(solver_config, state_types; interp = true)]\ntime_data = FT[0] # store time data\n\ncallback = GenericCallbacks.EveryXSimulationTime(every_x_simulation_time) do\n dons = dict_of_nodal_states(solver_config, state_types; interp = true)\n push!(dons_arr, dons)\n push!(time_data, gettime(solver_config.solver))\n nothing\nend;\n\n# # Run the integration\nClimateMachine.invoke!(solver_config; user_callbacks = (callback,));\n\n# Get z-coordinate\nz = get_z(solver_config.dg.grid; rm_dupes = true);\n\n# # Create some plots\n# We'll plot the moisture content vs depth in the soil, as well as\n# the expected profile of `ϑ_l` in hydrostatic equilibrium.\n# For `ϑ_l` values above porosity, the soil is\n# saturated, and the pressure head changes from being equal to the matric\n# potential to the pressure generated by compression of water and the soil\n# matrix. The profile can be solved\n# for analytically by (1) solving for the form that `ϑ_l(z)` must take\n# in both the saturated and unsaturated zones to satisfy the steady-state\n# requirement with zero flux boundary conditions, (2) requiring that\n# at the interface between saturated and unsaturated zones, the water content\n# equals porosity, and (3) solving for the location of the interface by\n# requiring that the integrated water content at the end matches that\n# at the beginning (yielding an interface location of `z≈-0.56m`).\noutput_dir = @__DIR__;\n\nt = time_data ./ (60 * 60 * 24);\n\nplot(\n dons_arr[1][\"soil.water.ϑ_l\"],\n dons_arr[1][\"z\"],\n label = string(\"t = \", string(t[1]), \"days\"),\n xlim = [0.47, 0.501],\n ylabel = \"z\",\n xlabel = \"ϑ_l\",\n legend = :bottomleft,\n title = \"Equilibrium test\",\n);\nplot!(\n dons_arr[2][\"soil.water.ϑ_l\"],\n dons_arr[2][\"z\"],\n label = string(\"t = \", string(t[2]), \"days\"),\n);\nplot!(\n dons_arr[7][\"soil.water.ϑ_l\"],\n dons_arr[7][\"z\"],\n label = string(\"t = \", string(t[7]), \"days\"),\n);\nfunction expected(z, z_interface)\n ν = 0.495\n S_s = 1e-3\n α = 2.6\n n = 2.0\n m = 0.5\n if z < z_interface\n return -S_s * (z - z_interface) + ν\n else\n return ν * (1 + (α * (z - z_interface))^n)^(-m)\n end\nend\nplot!(expected.(dons_arr[1][\"z\"], -0.56), dons_arr[1][\"z\"], label = \"expected\");\n\nplot!(\n 1e-3 .+ dons_arr[1][\"soil.water.ϑ_l\"],\n dons_arr[1][\"z\"],\n label = \"porosity\",\n);\n# save the output.\nsavefig(joinpath(output_dir, \"equilibrium_test_ϑ_l_vG.png\"))\n# ![](equilibrium_test_ϑ_l_vG.png)\n# # References\n# - [Woodward00a](@cite)\n" }, { "path": "tutorials/Land/Soil/Water/hydraulic_functions.jl", "content": "# # Hydraulic functions\n\n# This tutorial shows how to specify the hydraulic functions\n# used in Richard's equation. In particular,\n# we show how to choose the formalism for matric potential and hydraulic\n# conductivity, and how to make the hydraulic conductivity account for\n# the presence of ice as well as the temperature dependence of the\n# viscosity of liquid water.\n\n# # Preliminary setup\n\n# External modules\nusing Plots\n\n# ClimateMachine modules\nusing ClimateMachine\nusing ClimateMachine.Land\nusing ClimateMachine.Land.SoilWaterParameterizations\n\nFT = Float32;\n# # Specifying a hydraulics model\n# ClimateMachine's Land model allows the user to pick between two hydraulics models,\n# that of van Genuchten [vanGenuchten1980](@cite) or that of Brooks and Corey, see [BrooksCorey1964](@cite) or [Corey1977](@cite). The\n# same model is consistently used for the matric potential\n# and hydraulic conductivity.\n\n# The van Genuchten model requires two free parameters, `α` and `n`.\n# A third parameter, `m`, is computed from `n`. Of these, only `α` carries\n# units, of inverse meters. The Brooks and Corey model also uses\n# two free parameters, `ψ_b`, the magnitude of the matric potential at saturation,\n# and a constant `M`. `ψ_b` carries units of meters. These parameter sets are stored in\n# either the [`vanGenuchten`](@ref ClimateMachine.Land.SoilWaterParameterizations.vanGenuchten) or the\n# [`BrooksCorey`](@ref ClimateMachine.Land.SoilWaterParameterizations.BrooksCorey)\n# hydraulics model structures (more details below). These parameters are enough to compute the matric potential.\n\n# The hydraulic conductivity\n# requires an additional parameter, `Ksat` (m/s), which is the hydraulic conductivity\n# in saturated soil. This parameter is\n# not stored in the hydraulics model, but rather as part of the\n# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions), which stores\n# other parameters needed for the soil water modeling.\n\n# Below we show how to create two concrete examples of these hydraulics models,\n# for sandy loam ([Bonan19a](@cite)). Note that the parameters chosen are a function of soil type,\n# and that the parameters are converted to type `FT` internally.\nvg_α = 7.5 # m^-1\nvg_n = 1.89\nhydraulics = vanGenuchten(FT; α = vg_α, n = vg_n);\n\nψ_sat = 0.218 # m\nMval = 0.2041\nhydraulics_bc = BrooksCorey(FT; ψb = ψ_sat, m = Mval);\n# # Matric Potential\n# The matric potential `ψ` reflects the negative pressure of water\n# in unsaturated soil. The negative pressure (suction) of water arises\n# because of adhesive forces between water and soil.\n\n# The van Genuchten expression for matric potential is\n# ``\n# ψ = -\\frac{1}{α} S_l^{-1/(nm)}\\times (1-S_l^{1/m})^{1/n},\n# ``\n\n# and the Brooks and Corey expression is\n# ``\n# ψ = -ψ_b S_l^{-M}.\n# ``\n\n# Here `S_l` is the effective saturation of liquid water, `θ_l/ν`, where `ν` is\n# porosity of the soil. We generally neglect the residual pore space in the CliMA model,\n# but the user can set the parameter in the\n# [`WaterParamFunctions`](@ref ClimateMachine.Land.WaterParamFunctions) structure if it is\n# desired.\n\n# In the CliMA code, we use [multiple dispatch](https://en.wikipedia.org/wiki/Multiple_dispatch).\n# With multiple dispatch, a function can have many\n# ways of executing (called methods), depending on the *type* of the\n# variables passed in. A simple example of multiple dispatch is the division operation.\n# Integer division takes two numbers as input, and returns an integer - ignoring the decimal.\n# Float division takes two numbers as input, and returns a floating point number, including the decimal.\n# In Julia, we might write these as:\n\n# ```julia\n# function division(a::Int, b::Int)\n# return floor(Int, a/b)\n# end\n# ```\n# ```julia\n# function division(a::Float64, b::Float64)\n# return a/b\n# end\n# ```\n\n\n# We can see that `division` is now a function with two methods.\n\n# ```julia\n# julia> division\n# division (generic function with 2 methods)\n# ```\n\n# Now, using the same function signature, we can carry out integer\n# division or floating point division, depending on the types of the\n# arguments:\n\n# ```julia\n# julia> division(1,2)\n# 0\n#\n# julia> division(1.0,2.0)\n# 0.5\n# ```\n\n\n# Here are more pertinent examples:\n# Based on our choice of `FT = Float32`,\n\n# ```julia\n# julia> typeof(hydraulics)\n# vanGenuchten{Float32,Float32,Float32,Float32}\n# ```\n\n\n# but meanwhile,\n\n# ```julia\n# julia> typeof(hydraulics_bc)\n# BrooksCorey{Float32,Float32,Float32}\n# ```\n\n\n# The function `matric_potential` will execute different methods\n# depending on if we pass a hydraulics model of type `vanGenuchten` or\n# `BrooksCorey`. In both cases, it will return the correct value\n# for `ψ`.\n\n# Let's plot the matric potential as a function of the effective saturation `S_l = θ_l/ν`,\n# which can range from zero to one.\nS_l = FT.(0.01:0.01:0.99)\nψ = matric_potential.(Ref(hydraulics), S_l)\nψ_bc = matric_potential.(Ref(hydraulics_bc), S_l)\nplot(\n S_l,\n log10.(-ψ),\n xlabel = \"effective saturation\",\n ylabel = \"Log10(|ψ|)\",\n label = \"van Genuchten\",\n)\nplot!(S_l, log10.(-ψ_bc), label = \"Brooks and Corey\")\nsavefig(\"bc_vg_matric_potential.png\")\n# ![](bc_vg_matric_potential.png)\n\n# The steep slope in\n# `ψ` near saturated and completely dry soil are part of the reason\n# why Richard's equation is such a challenging numerical problem.\n\n\n# # Hydraulic conductivity\n# The hydraulic conductivity is a more complex function than the matric potential,\n# as it depends on the temperature of the water, the volumetric ice fraction, and\n# the volumetric liquid water fraction. It also depends on the hydraulics model\n# chosen.\n\n# We represent the hydraulic conductivity `K` as the product of four factors:\n# `Ksat`, an impedance factor (which accounts for the effect of ice on conductivity)\n# a viscosity factor (which accounts for the effect of temperature on the\n# viscosity of liquid water, and how that in turn affects conductivity)\n# and a moisture factor (which accounts for the effect of liquid water, and is determined by the hydraulics model).\n# We are going to calculate `K = Ksat × viscosity factor × impedance factor × moisture factor`.\n# In the code, each of these factors is\n# computed by a function with multiple methods, except for `Ksat`.\n# Like we defined new type\n# classes for `vanGenuchten` and `BrooksCorey`, we also created new type classes\n# for the impedance choice, the viscosity choice, and the moisture choice.\n\n# The function [`viscosity_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.viscosity_factor)\n# takes as arguments the temperature of the soil and the\n# viscosity model desired, and returns the factor `k_v` by which the hydraulic conductivity is scaled.\n# One option is to account for this effect:\n\n# ``\n# k_v = e^{γ (T-T_{\\rm ref})}\n# ``\n\n# where γ = 0.0264/K and ``T_{\\rm ref}`` = 288K.\n\n# For example, at the freezing point of water, using the default values\n# for γ and T_ref, viscosity reduces the conductivity by a third:\nviscous_effect_model = TemperatureDependentViscosity{FT}();\nviscosity_factor(viscous_effect_model, FT(273.15))\n\n# The other option is to ignore this effect:\n\n# ``\n# k_v = 1\n# ``\n\n# This is the default approach.\nno_viscous_effect_model = ConstantViscosity{FT}();\nviscosity_factor(no_viscous_effect_model, FT(273.15))\n\n# Very similarly, the function\n# [`impedance_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.impedance_factor)\n# takes as arguments the liquid water and ice\n# volumetric fractions in the soil, as well as the impedance model being used, and returns\n# the factor `k_i` by which the hydraulic conductivity is scaled.\n# One option is to account for this effect:\n\n# ``\n# k_i = 10^{-Ω f_i},\n# ``\n\n# where `Ω = 7` is an empirical factor and\n# `f_i` is the ratio of the volumetric\n# ice fraction to total volumetric water fraction ([Lundin1990](@cite)).\n\n# For example, with ``\\theta_i = \\theta_l``, or f_i = 0.5, ice reduces the conductivity by over 1000x.\nimpedance_effect_model = IceImpedance{FT}();\nimpedance_factor(impedance_effect_model, FT(0.5))\n\n# The other option is to ignore this effect:\n\n# ``\n# k_i = 1\n# ``\n\n# This is the default approach.\nno_impedance_effect_model = NoImpedance{FT}();\nimpedance_factor(no_impedance_effect_model, FT(0.5))\n\n# As for the moisture dependence of hydraulic conductivity, it can also be either\n# independent of moisture, or dependent on moisture. If it is dependent on moisture,\n# the specific function evaluated is dictated by the hydraulics model.\n# The [`moisture_factor`](@ref ClimateMachine.Land.SoilWaterParameterizations.moisture_factor)\n# for the van Genuchten model is (denoting it as ``k_m``)\n\n# ``\n# k_m = \\sqrt{S_l}[1-(1-S_l^{1/m})^m]^2,\n# ``\n\n# for ``S_l < 1``,\n\n# and for the Brooks and Corey model it is\n\n# ``\n# k_m = S_l^{2M+3},\n# ``\n\n# also for ``S_l<1``. When ``S_l\\geq 1``, ``k_m = 1`` for each model.\n\n\n\n# Let's put all these factors together now. Below\n# we choose additional parameters, consistent with the hydraulics parameters\n# for sandy loam ([Bonan19a](@cite)), and show how hydraulic conductivity varies with\n# liquid water content, in the case without ice impedance or temperature effects.\nKsat = FT(4.42 / (3600 * 100))\nT = FT(0.0)\nf_i = FT(0.0)\nK =\n hydraulic_conductivity.(\n Ref(Ksat),\n Ref(impedance_factor(NoImpedance{FT}(), f_i)),\n Ref(viscosity_factor(ConstantViscosity{FT}(), T)),\n moisture_factor.(Ref(MoistureDependent{FT}()), Ref(hydraulics), S_l),\n );\n\n# Let's also compute `K` when we include the effects of temperature\n# and ice on the hydraulic conductivity.\n# In the cases where a\n# [`TemperatureDependentViscosity`](@ref ClimateMachine.Land.SoilWaterParameterizations.TemperatureDependentViscosity)\n# or\n# [`IceImpedance`](@ref ClimateMachine.Land.SoilWaterParameterizations.IceImpedance)\n# type is passed, the correct factors are calculated,\n# based on the temperature `T` and volumetric ice fraction `θ_i`.\n\nT = FT(273.15)\nS_i = FT(0.1); # = θ_i/ν\n# The total volumetric water fraction cannot\n# exceed unity, so the effective liquid water saturation\n# should have a max of 1-S_i.\nS_l_accounting_for_ice = FT.(0.01:0.01:(0.99 - S_i))\nf_i = S_i ./ (S_l_accounting_for_ice .+ S_i)\nK_w_factors =\n hydraulic_conductivity.(\n Ref(Ksat),\n impedance_factor.(Ref(NoImpedance{FT}()), f_i),\n Ref(viscosity_factor(ConstantViscosity{FT}(), T)),\n moisture_factor.(\n Ref(MoistureDependent{FT}()),\n Ref(hydraulics),\n S_l_accounting_for_ice,\n ),\n );\nplot(\n S_l,\n log10.(K),\n xlabel = \"total effective saturation, (θ_i+θ_l)/ν\",\n ylabel = \"Log10(K)\",\n label = \"Base case\",\n legend = :bottomright,\n)\nplot!(\n S_l_accounting_for_ice .+ S_i,\n log10.(K_w_factors),\n label = \"θ_i = 0.1, T = 273.15\",\n)\nsavefig(\"T_ice_K.png\")\n# ![](T_ice_K.png)\n\n# If the user is not considering phase transitions\n# and does not add in Freeze/Thaw source terms, the default is for zero\n# ice in the model, for all time and space. In this case the ice impedance\n# factor evaluates to 1 regardless of which type is passed.\n\n# # Other features\n# The user also has the choice of making the conductivity constant by choosing\n# [`MoistureIndependent`](@ref ClimateMachine.Land.SoilWaterParameterizations.MoistureIndependent)\n# along with\n# [`ConstantViscosity`](@ref ClimateMachine.Land.SoilWaterParameterizations.ConstantViscosity)\n# and\n# [`NoImpedance`](@ref ClimateMachine.Land.SoilWaterParameterizations.NoImpedance).\n# This is useful for debugging!\nno_moisture_dependence = MoistureIndependent{FT}()\nK_constant =\n hydraulic_conductivity.(\n Ref(Ksat),\n Ref(FT(1.0)),\n Ref(FT(1.0)),\n moisture_factor.(Ref(no_moisture_dependence), Ref(hydraulics), S_l),\n );\n# ```julia\n# julia> unique(K_constant)\n# 1-element Array{Float32,1}:\n# 1.2277777f-5\n# ```\n\n# Note that choosing this option does not mean the matric potential\n# is constant, as a hydraulics model is still required and employed.\n\n\n# And, lastly, you might also find it helpful in debugging\n# to be able to turn off the flow of water by setting `Ksat = 0`.\n\n# # References\n# - [vanGenuchten1980](@cite)\n# - [BrooksCorey1964](@cite)\n# - [Corey1977](@cite)\n# - [Lundin1990](@cite)\n# - [Bonan19a](@cite)\n" }, { "path": "tutorials/Land/Soil/interpolation_helper.jl", "content": "\"\"\"\n function create_interpolation_grid(xbnd, xres, thegrid)\nGiven boundaries of a domain, and a resolution in each direction, create \nan interpolation grid. This is where the interpolation functions for \na set of variables will be evaluated.\n\"\"\"\nfunction create_interpolation_grid(xbnd, xres, thegrid)\n x1g = collect(range(xbnd[1, 1], xbnd[2, 1], step = xres[1]))\n x2g = collect(range(xbnd[1, 2], xbnd[2, 2], step = xres[2]))\n x3g = collect(range(xbnd[1, 3], xbnd[2, 3], step = xres[3]))\n intrp_brck = ClimateMachine.InterpolationBrick(thegrid, xbnd, x1g, x2g, x3g)\n return intrp_brck\nend\n\n\"\"\"\n function interpolate_variables(objects, brick)\nCreate an interpolation function from data and evaluate that\nfunction on the interpolation brick passed.\n\"\"\"\nfunction interpolate_variables(objects, brick)\n i_objects = []\n for object in objects\n nvars = size(object.data, 2)\n i_object = Array{FT}(undef, brick.Npl, nvars)\n ClimateMachine.interpolate_local!(brick, object.data, i_object)\n push!(i_objects, i_object)\n end\n\n return i_objects\nend\n" }, { "path": "tutorials/Numerics/DGMethods/Box1D.jl", "content": "# A box advection test to visualise how different filters work\n\nusing MPI\nusing OrderedCollections\nusing Plots\nusing StaticArrays\nusing Printf\n\nusing CLIMAParameters\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nusing ClimateMachine\nusing ClimateMachine.Mesh.Topologies\nusing ClimateMachine.Mesh.Grids\nusing ClimateMachine.DGMethods\nusing ClimateMachine.DGMethods.NumericalFluxes\nusing ClimateMachine.BalanceLaws:\n BalanceLaw, Prognostic, Auxiliary, Gradient, GradientFlux\nusing ClimateMachine.Mesh.Geometry: LocalGeometry\nusing ClimateMachine.Mesh.Filters\nusing ClimateMachine.MPIStateArrays\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.SingleStackUtils\n\nimport ClimateMachine.BalanceLaws:\n vars_state,\n source!,\n flux_second_order!,\n flux_first_order!,\n compute_gradient_argument!,\n compute_gradient_flux!,\n update_auxiliary_state!,\n nodal_init_state_auxiliary!,\n init_state_prognostic!,\n boundary_conditions,\n wavespeed\n\nClimateMachine.init(; disable_gpu = true, log_level = \"warn\");\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"docs\", \"plothelpers.jl\"));\n\nBase.@kwdef struct Box1D{FT, _init_q, _amplitude, _velo} <: BalanceLaw\n param_set::AbstractParameterSet = param_set\n init_q::FT = _init_q\n amplitude::FT = _amplitude\n velo::FT = _velo\nend\n\nvars_state(::Box1D, ::Auxiliary, FT) = @vars(z_dim::FT);\nvars_state(::Box1D, ::Prognostic, FT) = @vars(q::FT);\nvars_state(::Box1D, ::Gradient, FT) = @vars();\nvars_state(::Box1D, ::GradientFlux, FT) = @vars();\n\nfunction wavespeed(\n ::Box1D{FT, _init_q, _amplitude, _velo},\n _...,\n) where {FT, _init_q, _amplitude, _velo}\n return _velo\nend\n\nfunction nodal_init_state_auxiliary!(\n m::Box1D,\n aux::Vars,\n tmp::Vars,\n geom::LocalGeometry,\n)\n aux.z_dim = geom.coord[3]\nend;\n\nfunction init_state_prognostic!(\n m::Box1D,\n state::Vars,\n aux::Vars,\n localgeo,\n t::Real,\n)\n if aux.z_dim >= 75 && aux.z_dim <= 125\n state.q = m.init_q + m.amplitude\n else\n state.q = m.init_q\n end\nend;\n\nfunction update_auxiliary_state!(\n dg::DGModel,\n m::Box1D,\n Q::MPIStateArray,\n t::Real,\n elems::UnitRange,\n)\n return true\nend;\n\nfunction source!(m::Box1D, _...) end;\n\n@inline function flux_first_order!(\n m::Box1D,\n flux::Grad,\n state::Vars,\n aux::Vars,\n t::Real,\n _...,\n)\n FT = eltype(state)\n @inbounds begin\n flux.q = SVector(FT(0), FT(0), state.q * m.velo)\n end\nend\n\n@inline function flux_second_order!(\n m::Box1D,\n flux::Grad,\n state::Vars,\n diffusive::Vars,\n hyperdiffusive::Vars,\n aux::Vars,\n t::Real,\n) end\n\n@inline function flux_second_order!(\n m::Box1D,\n flux::Grad,\n state::Vars,\n τ,\n d_h_tot,\n) end\n\nboundary_condtions(m::Box1D) = ()\n\nfunction run_box1D(\n N_poly::Int,\n init_q::FT,\n amplitude::FT,\n velo::FT,\n plot_name::String;\n tmar_filter::Bool = false,\n cutoff_filter::Bool = false,\n exp_filter::Bool = false,\n boyd_filter::Bool = false,\n cutoff_param::Int = 1,\n exp_param_1::Int = 0,\n exp_param_2::Int = 32,\n boyd_param_1::Int = 0,\n boyd_param_2::Int = 32,\n numerical_flux_first_order = CentralNumericalFluxFirstOrder(),\n) where {FT}\n N_poly = N_poly\n nelem = 128\n zmax = FT(350)\n\n m = Box1D{FT, init_q, amplitude, velo}()\n\n driver_config = ClimateMachine.SingleStackConfiguration(\n \"Box1D\",\n N_poly,\n nelem,\n zmax,\n param_set,\n m,\n numerical_flux_first_order = numerical_flux_first_order,\n boundary = ((0, 0), (0, 0), (0, 0)),\n periodicity = (true, true, true),\n )\n\n t0 = FT(0)\n timeend = FT(450)\n\n Δ = min_node_distance(driver_config.grid, VerticalDirection())\n max_vel = m.velo\n dt = Δ / max_vel\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n ode_dt = dt,\n )\n grid = solver_config.dg.grid\n Q = solver_config.Q\n aux = solver_config.dg.state_auxiliary\n\n output_dir = @__DIR__\n mkpath(output_dir)\n\n z_label = \"z\"\n z = get_z(grid)\n\n # store initial condition at ``t=0``\n dons_arr = Dict[dict_of_nodal_states(solver_config)]\n time_data = FT[0] # store time data\n\n # output\n output_freq = floor(Int, timeend / dt) + 10\n\n cb_output = GenericCallbacks.EveryXSimulationSteps(output_freq) do\n push!(dons_arr, dict_of_nodal_states(solver_config))\n push!(time_data, gettime(solver_config.solver))\n nothing\n end\n\n filter_freq = 1\n # tmar filter\n cb_tmar =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:q,),\n solver_config.dg.grid,\n TMARFilter(),\n )\n nothing\n end\n # cutoff filter\n cb_cutoff =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:q,),\n solver_config.dg.grid,\n CutoffFilter(solver_config.dg.grid, cutoff_param),\n )\n nothing\n end\n # exponential filter\n cb_exp =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:q,),\n solver_config.dg.grid,\n ExponentialFilter(\n solver_config.dg.grid,\n exp_param_1,\n exp_param_2,\n ),\n )\n nothing\n end\n # Boyd Vandeven filter\n cb_boyd =\n GenericCallbacks.EveryXSimulationSteps(filter_freq) do (init = false)\n Filters.apply!(\n solver_config.Q,\n (:q,),\n solver_config.dg.grid,\n BoydVandevenFilter(\n solver_config.dg.grid,\n boyd_param_1,\n boyd_param_2,\n ),\n )\n nothing\n end\n\n user_cb_arr = [cb_output]\n if tmar_filter\n push!(user_cb_arr, cb_tmar)\n end\n if cutoff_filter\n push!(user_cb_arr, cb_cutoff)\n end\n if exp_filter\n push!(user_cb_arr, cb_exp)\n end\n if boyd_filter\n push!(user_cb_arr, cb_boyd)\n end\n user_cb = (user_cb_arr...,)\n\n initial_mass = weightedsum(solver_config.Q)\n ClimateMachine.invoke!(solver_config; user_callbacks = (user_cb))\n final_mass = weightedsum(solver_config.Q)\n @info @sprintf(\n \"\"\"\nMass Conservation:\n initial mass = %.16e\n final mass = %.16e\n difference = %.16e\n normalized difference = %.16e\"\"\",\n initial_mass,\n final_mass,\n final_mass - initial_mass,\n (final_mass - initial_mass) / initial_mass\n )\n\n push!(dons_arr, dict_of_nodal_states(solver_config))\n push!(time_data, gettime(solver_config.solver))\n\n export_plot(\n z,\n time_data,\n dons_arr,\n (\"q\",),\n joinpath(output_dir, plot_name);\n xlabel = \"x\",\n ylabel = \"q\",\n horiz_layout = true,\n )\n\nend\n" }, { "path": "tutorials/Numerics/DGMethods/showcase_filters.jl", "content": "# # Filters\n# In this tutorial we show the result of applying filters\n# available in the CliMA codebase in a 1 dimensional box advection setup.\n# See [Filters API](https://clima.github.io/ClimateMachine.jl/latest/APIs/Numerics/Meshes/Mesh/#Filters-1) for filters interface details.\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(clima_dir, \"tutorials\", \"Numerics\", \"DGMethods\", \"Box1D.jl\"))\n\nconst FT = Float64\n\noutput_dir = @__DIR__;\nmkpath(output_dir);\n\n# The unfiltered result of the box advection test for order 4 polynomial with\n# central flux is\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_no_filter.svg\"),\n)\n# ![](box_1D_4_no_filter.svg)\n\n# The unfiltered result of the box advection test for order 4 polynomial with\n# Rusanov flux (aka upwinding for advection) is\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_no_filter_upwind.svg\"),\n numerical_flux_first_order = RusanovNumericalFlux(),\n)\n# ![](box_1D_4_no_filter_upwind.svg)\n\n\n# Below we show results for the same box advection test\n# but using different filters.\n#\n# As seen in the results, when the TMAR filter is used mass is not necessarily\n# conserved (mass increases are possible).\n\n# `TMARFilter()` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_tmar.svg\");\n tmar_filter = true,\n)\n# ![](box_1D_4_tmar.svg)\n\n# Running the TMAR filter with Rusanov the mass conservation since some of the\n# are reduced, but mass is still not conserved.\n# `TMARFilter()` with Rusanov numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_tmar_upwind.svg\");\n tmar_filter = true,\n numerical_flux_first_order = RusanovNumericalFlux(),\n)\n# ![](box_1D_4_tmar_upwind.svg)\n\n# `CutoffFilter(grid, Nc=1)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_cutoff_1.svg\");\n cutoff_filter = true,\n cutoff_param = 1,\n)\n# ![](box_1D_4_cutoff_1.svg)\n\n# `CutoffFilter(grid, Nc=3)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_cutoff_3.svg\");\n cutoff_filter = true,\n cutoff_param = 3,\n)\n# ![](box_1D_4_cutoff_3.svg)\n\n# `ExponentialFilter(grid, Nc=1, s=4)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_exp_1_4.svg\");\n exp_filter = true,\n exp_param_1 = 1,\n exp_param_2 = 4,\n)\n# ![](box_1D_4_exp_1_4.svg)\n\n# `ExponentialFilter(grid, Nc=1, s=8)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_exp_1_8.svg\");\n exp_filter = true,\n exp_param_1 = 1,\n exp_param_2 = 8,\n)\n# ![](box_1D_4_exp_1_8.svg)\n\n# `ExponentialFilter(grid, Nc=1, s=32)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_exp_1_32.svg\");\n exp_filter = true,\n exp_param_1 = 1,\n exp_param_2 = 32,\n)\n# ![](box_1D_4_exp_1_32.svg)\n\n# `BoydVandevenFilter(grid, Nc=1, s=4)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_boyd_1_4.svg\");\n boyd_filter = true,\n boyd_param_1 = 1,\n boyd_param_2 = 4,\n)\n# ![](box_1D_4_boyd_1_4.svg)\n\n# `BoydVandevenFilter(grid, Nc=1, s=8)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_boyd_1_8.svg\");\n boyd_filter = true,\n boyd_param_1 = 1,\n boyd_param_2 = 8,\n)\n# ![](box_1D_4_boyd_1_8.svg)\n\n# `BoydVandevenFilter(grid, Nc=1, s=32)` with central numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_boyd_1_32.svg\");\n boyd_filter = true,\n boyd_param_1 = 1,\n boyd_param_2 = 32,\n)\n# ![](box_1D_4_boyd_1_32.svg)\n\n# `ExponentialFilter(grid, Nc=1, s=8)` and `TMARFilter()` with central numerical\n# flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_tmar_exp_1_8.svg\");\n exp_filter = true,\n tmar_filter = true,\n exp_param_1 = 1,\n exp_param_2 = 8,\n)\n# ![](box_1D_4_tmar_exp_1_8.svg)\n\n# `BoydVandevenFilter(grid, Nc=1, s=8)` and `TMARFilter()` with central\n# numerical flux:\nrun_box1D(\n 4,\n FT(0.0),\n FT(1.0),\n FT(1.0),\n joinpath(output_dir, \"box_1D_4_tmar_boyd_1_8.svg\");\n boyd_filter = true,\n tmar_filter = true,\n boyd_param_1 = 1,\n boyd_param_2 = 8,\n)\n# ![](box_1D_4_tmar_boyd_1_8.svg)\n" }, { "path": "tutorials/Numerics/SystemSolvers/bgmres.jl", "content": "# # Batched Generalized Minimal Residual\n# In this tutorial we describe the basics of using the batched gmres iterative solver.\n# At the end you should be able to\n# 1. Use BatchedGeneralizedMinimalResidual to solve batches of linear systems\n# 2. Construct a columnwise linear solver with BatchedGeneralizedMinimalResidual\n\n# ## What is the Generalized Minimal Residual Method?\n# The Generalized Minimal Residual Method (GMRES) is a [Krylov subspace](https://en.wikipedia.org/wiki/Krylov_subspace) method for solving linear systems:\n# ```math\n# Ax = b\n# ```\n# See the [wikipedia](https://en.wikipedia.org/wiki/Generalized_minimal_residual_method) for more details.\n\n# ## What is the Batched Generalized Minimal Residual Method?\n# As the name suggests it solves a whole bunch of independent GMRES problems\n\n# ## Basic Example\n# First we must load a few things\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing LinearAlgebra, Random, Plots\n\n# Next we define two linear systems that we would like to solve simultaneously.\n# The matrix for the first linear system is\nA1 = [\n 2.0 -1.0 0.0\n -1.0 2.0 -1.0\n 0.0 -1.0 2.0\n];\n# And the right hand side is\nb1 = ones(typeof(1.0), 3);\n# The exact solution to the first linear system is\nx1_exact = [1.5, 2.0, 1.5];\n# The matrix for the first linear system is\nA2 = [\n 2.0 -1.0 0.0\n 0.0 2.0 -1.0\n 0.0 0.0 2.0\n];\n# And the right hand side is\nb2 = ones(typeof(1.0), 3);\n# The exact solution to second linear system is\nx2_exact = [0.875, 0.75, 0.5];\n\n# We now define a function that performs the action of each linear operator independently.\nfunction closure_linear_operator(A1, A2)\n function linear_operator!(x, y)\n mul!(view(x, :, 1), A1, view(y, :, 1))\n mul!(view(x, :, 2), A2, view(y, :, 2))\n return nothing\n end\n return linear_operator!\nend;\n\n# To understand how this works let us construct an instance\n# of the linear operator and apply it to a vector\nlinear_operator! = closure_linear_operator(A1, A2);\n# Let us see what the action of this linear operator is\ny1 = ones(typeof(1.0), 3);\ny2 = ones(typeof(1.0), 3) * 2.0;\ny = [y1 y2];\nx = copy(y);\nlinear_operator!(x, y);\nx\n# We see that the first column is `A1 * [1 1 1]'`\n# and the second column is `A2 * [2 2 2]'`\n# that is,\n[A1 * y1 A2 * y2]\n\n# We are now ready to set up our Batched Generalized Minimal Residual solver\n# We must now set up the right hand side of the linear system\nb = [b1 b2];\n# as well as the exact solution, (to verify convergence)\nx_exact = [x1_exact x2_exact];\n# !!! warning\n# For BatchedGeneralizedMinimalResidual the assumption is that each column of b is independent and corresponds to a batch. This will come back later.\n\n# We now use an instance of the solver\nlinearsolver = BatchedGeneralizedMinimalResidual(b, size(A1, 1), 2);\n# As well as an initial guess, denoted by the variable x\nx1 = ones(typeof(1.0), 3);\nx2 = ones(typeof(1.0), 3);\nx = [x1 x2];\n# To solve the linear system, we just need to pass to the linearsolve! function\niters = linearsolve!(linear_operator!, nothing, linearsolver, x, b)\n# which is guaranteed to converge in 3 iterations since `length(b1)=length(b2)=3`\n# We can now check that the solution that we computed, x\nx\n# has converged to the exact solution\nx_exact\n# Which indeed it has.\n# ## Advanced Example\n\n# We now go through a more advanced application of the Batched Generalized Minimal Residual solver\n# !!! warning\n# Iterative methods should be used with preconditioners!\n# The first thing we do is define a linear operator that mimics\n# the behavior of a columnwise operator in ClimateMachine\nfunction closure_linear_operator!(A, tup)\n function linear_operator!(y, x)\n alias_x = reshape(x, tup)\n alias_y = reshape(y, tup)\n for i6 in 1:tup[6]\n for i4 in 1:tup[4]\n for i2 in 1:tup[2]\n for i1 in 1:tup[1]\n tmp = alias_x[i1, i2, :, i4, :, i6][:]\n tmp2 = A[i1, i2, i4, i6] * tmp\n alias_y[i1, i2, :, i4, :, i6] .=\n reshape(tmp2, (tup[3], tup[5]))\n end\n end\n end\n end\n end\nend;\n# Next we define the array structure of an MPIStateArray\n# in its true high dimensional form\ntup = (2, 2, 5, 2, 10, 2);\n# We define our linear operator as a random matrix\nRandom.seed!(1234);\nB = [\n randn(tup[3] * tup[5], tup[3] * tup[5])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\ncolumnwise_A = [\n B[i1, i2, i4, i6] + 3 * (i1 + i2 + i4 + i6) * I\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\n# as well as its inverse\ncolumnwise_inv_A = [\n inv(columnwise_A[i1, i2, i4, i6])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\ncolumnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup);\ncolumnwise_inverse_linear_operator! =\n closure_linear_operator!(columnwise_inv_A, tup);\n# The structure of an MPIStateArray is related to its true\n# higher dimensional form as follows:\nmpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6]);\n# We now define the right hand side of our Linear system\nb = randn(mpi_tup);\n# As well as the initial guess\nx = copy(b);\nx += randn(mpi_tup) * 0.1;\n# In the previous tutorial we mentioned that it is assumed that\n# the right hand side is an array whose column vectors all independent linear\n# systems. But right now the array structure of ``x`` and ``b`` do not follow\n# this requirement.\n# To handle this case we must pass in additional arguments that tell the\n# linear solver how to reconcile these differences.\n# The first thing that the linear solver must know of is the higher tensor\n# form of the MPIStateArray, which is just the `tup` from before\nreshape_tuple_f = tup;\n# The second thing it needs to know is which indices correspond to a column\n# and we want to make sure that these are the first set of indices that appear\n# in the permutation tuple (which can be thought of as enacting\n# a Tensor Transpose).\npermute_tuple_f = (5, 3, 4, 6, 1, 2);\n# It has this format since the 3 and 5 index slots\n# are the ones associated with traversing a column. And the 4 index\n# slot corresponds to a state.\n# We also need to tell our solver which kind of Array struct to use\nArrayType = Array;\n\n# We are now ready to finally define our linear solver, which uses a number\n# of keyword arguments\ngmres = BatchedGeneralizedMinimalResidual(\n b,\n tup[3] * tup[5] * tup[4],\n tup[1] * tup[2] * tup[6];\n atol = eps(Float64) * 100,\n rtol = eps(Float64) * 100,\n forward_reshape = reshape_tuple_f,\n forward_permute = permute_tuple_f,\n);\n# `m` is the number of gridpoints along a column. As mentioned previously,\n# this is `tup[3]*tup[5]*tup[4]`. The `n` term corresponds to the batch size\n# or the number of columns in this case. `atol` and `rtol` are relative and\n# absolute tolerances\n\n# All the hard work is done, now we just call our linear solver\niters = linearsolve!(\n columnwise_linear_operator!,\n nothing,\n gmres,\n x,\n b,\n max_iters = tup[3] * tup[5] * tup[4],\n)\n# We see that it converged in less than `tup[3]*tup[5] = 50` iterations.\n# Let us verify that it is indeed correct by computing the exact answer\n# numerically and comparing it against the iterative solver.\nx_exact = copy(x);\ncolumnwise_inverse_linear_operator!(x_exact, b);\n# Now we can compare with some norms\nnorm(x - x_exact) / norm(x_exact)\ncolumnwise_linear_operator!(x_exact, x);\nnorm(x_exact - b) / norm(b)\n# Which we see are small, given our choice of `atol` and `rtol`.\n# The struct also keeps a record of its convergence rate\n# in the residual member (max over all batched solvers).\n# The can be visualized via\nplot(log.(gmres.resnorms[:]) / log(10));\nplot!(legend = false, xlims = (1, iters), ylims = (-15, 2));\nplot!(ylabel = \"log10 residual\", xlabel = \"iterations\")\n" }, { "path": "tutorials/Numerics/SystemSolvers/cg.jl", "content": "# # Conjugate Gradient\n# In this tutorial we describe the basics of using the conjugate gradient iterative solvers\n# At the end you should be able to\n# 1. Use Conjugate Gradient to solve a linear system\n# 2. Know when to not use it\n# 3. Contruct a column-wise linear solver with Conjugate Gradient\n\n# ## What is it?\n# Conjugate Gradient is an iterative method for solving special kinds of linear systems:\n# ```math\n# Ax = b\n# ```\n# via iterative methods.\n# !!! warning\n# The linear operator need to be symmetric positive definite and the preconditioner must be symmetric.\n# See the [wikipedia](https://en.wikipedia.org/wiki/Conjugate_gradient_method) for more details.\n\n# ## Basic Example\n# First we must load a few things\nusing ClimateMachine\nusing ClimateMachine.SystemSolvers\nusing LinearAlgebra, Random\n\n# Next we define a 3x3 symmetric positive definite linear system. (In the ClimateMachine code a symmetric positive definite system could arise from treating diffusion implicitly.)\nA = [\n 2.0 -1.0 0.0\n -1.0 2.0 -1.0\n 0.0 -1.0 2.0\n];\n# We define the matrix `A` here as a global variable for convenience later.\n\n# We can see that it is symmetric. We can check that it is positive definite by checking the spectrum\neigvals(A)\n\n# The linear operators that are passed into the abstract iterative solvers need to be defined as functions that act on vectors. Let us do that with our matrix. We are using function closures for type stability.\nfunction closure_linear_operator!(A)\n function linear_operator!(x, y)\n mul!(x, A, y)\n end\n return linear_operator!\nend;\n\n# We now define our linear operator using the function closure\n\nlinear_operator! = closure_linear_operator!(A)\n\n# We now define our `b` in the linear system\nb = ones(typeof(1.0), 3);\n\n# The exact solution to the system `Ax = b` is\nx_exact = [1.5, 2.0, 1.5];\n\n# Now we can set up the ConjugateGradient struct\nlinearsolver = ConjugateGradient(b);\n# and an initial guess for the iterative solver.\nx = ones(typeof(1.0), 3);\n# To solve the linear system we just need to pass to the linearsolve! function\niters = linearsolve!(linear_operator!, nothing, linearsolver, x, b)\n# The variable `x` gets overwritten during the linear solve\n# The norm of the error is\nnorm(x - x_exact) / norm(x_exact)\n# The relative norm of the residual is\nnorm(A * x - b) / norm(b)\n# The number of iterations is\niters\n# Conjugate Gradient is guaranteed to converge in 3 iterations with perfect arithmetic in this case.\n\n# ## Non-Example\n\n# Conjugate Gradient is not guaranteed to converge with nonsymmetric matrices. Consider\nA = [\n 2.0 -1.0 0.0\n 0.0 2.0 -1.0\n 0.0 0.0 2.0\n];\n# We define the matrix `A` here as a global variable for convenience later.\n\n# We can see that it is not symmetric, but it does have all positive eigenvalues\neigvals(A)\n\n# The linear operators that are passed into the abstract iterative solvers need to be defined as functions that act on vectors. Let us do that with our matrix. We are using function closures for type stability.\nfunction closure_linear_operator!(A)\n function linear_operator!(x, y)\n mul!(x, A, y)\n end\n return linear_operator!\nend;\n\n# We define the linear operator using our function closure\n\nlinear_operator! = closure_linear_operator!(A)\n\n# We now define our `b` in the linear system\nb = ones(typeof(1.0), 3);\n\n# The exact solution to the system `Ax = b` is\nx_exact = [0.875, 0.75, 0.5];\n\n# Now we can set up the ConjugateGradient struct\nlinearsolver = ConjugateGradient(b, max_iter = 100);\n# We also passed in the keyword argument \"max_iter\" for the maximum number of iterations of the iterative solver. By default it is assumed to be the size of the vector.\n# As before we need to define an initial guess\nx = ones(typeof(1.0), 3);\n# To (not) solve the linear system we just need to pass to the linearsolve! function\niters = linearsolve!(linear_operator!, nothing, linearsolver, x, b)\n# The variable `x` gets overwitten during the linear solve\n# The norm of the error is\nnorm(x - x_exact) / norm(x_exact)\n# The relative norm of the residual is\nnorm(A * x - b) / norm(b)\n# The number of iterations is\niters\n# Conjugate Gradient is guaranteed to converge in 3 iterations with perfect arithmetic for a symmetric positive definite matrix. Here we see that the matrix is not symmetric and it didn't converge even after 100 iterations.\n\n\n# ## More Complex Example\n# Here we show how to construct a column-wise iterative solver similar to what is is in the ClimateMachine code. The following is not for the faint of heart.\n# We must first define a linear operator that acts like one in the ClimateMachine\nfunction closure_linear_operator!(A, tup)\n function linear_operator!(y, x)\n alias_x = reshape(x, tup)\n alias_y = reshape(y, tup)\n for i6 in 1:tup[6]\n for i4 in 1:tup[4]\n for i2 in 1:tup[2]\n for i1 in 1:tup[1]\n tmp = alias_x[i1, i2, :, i4, :, i6][:]\n tmp2 = A[i1, i2, i4, i6] * tmp\n alias_y[i1, i2, :, i4, :, i6] .=\n reshape(tmp2, (tup[3], tup[5]))\n end\n end\n end\n end\n end\nend;\n\n# Now that we have this function, we can define a linear system that we will solve columnwise\n# First we define the structure of our array as `tup` in a manner that is similar to a stacked brick topology\ntup = (3, 4, 7, 2, 20, 2);\n# where\n# 1. tup[1] is the number of Gauss–Lobatto points in the x-direction\n# 2. tup[2] is the number of Gauss–Lobatto points in the y-direction\n# 3. tup[3] is the number of Gauss–Lobatto points in the z-direction\n# 4. tup[4] is the number of states\n# 6. tup[5] is the number of elements in the vertical direction\n# 7. tup[6] is the number of elements in the other directions\n\n# Now we define our linear operator as a random matrix.\nRandom.seed!(1235);\nB = [\n randn(tup[3] * tup[5], tup[3] * tup[5])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\ncolumnwise_A = [\n B[i1, i2, i4, i6] * B[i1, i2, i4, i6]' + 10I\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\ncolumnwise_inv_A = [\n inv(columnwise_A[i1, i2, i4, i6])\n for i1 in 1:tup[1], i2 in 1:tup[2], i4 in 1:tup[4], i6 in 1:tup[6]\n];\ncolumnwise_linear_operator! = closure_linear_operator!(columnwise_A, tup);\ncolumnwise_inverse_linear_operator! =\n closure_linear_operator!(columnwise_inv_A, tup);\n\n# We define our `x` and `b` with matrix structures similar to an MPIStateArray\nmpi_tup = (tup[1] * tup[2] * tup[3], tup[4], tup[5] * tup[6]);\nb = randn(mpi_tup);\nx = randn(mpi_tup);\n\n# Now we solve the linear system columnwise\nlinearsolver = ConjugateGradient(\n x,\n max_iter = tup[3] * tup[5],\n dims = (3, 5),\n reshape_tuple = tup,\n);\n# The keyword arguments dims is the reduction dimension for the linear solver. In this case dims = (3,5) are the ones associated with a column. The reshape_tuple argument is to convert the shapes of the array `x` in the a form that is more easily usable for reductions in the linear solver\n\n# Now we can solve it\niters = linearsolve!(columnwise_linear_operator!, nothing, linearsolver, x, b);\nx_exact = copy(x);\ncolumnwise_inverse_linear_operator!(x_exact, b);\n# The norm of the error is\nnorm(x - x_exact) / norm(x_exact)\n# The number of iterations is\niters\n# The algorithm converges within `tup[3]*tup[5] = 140` iterations\n\n# ## Tips\n# 1. The convergence criteria should be changed, machine precision is too small and the maximum iterations is often too large\n# 2. Use a preconditioner if possible\n# 3. Make sure that the linear system really is symmetric and positive-definite\n" }, { "path": "tutorials/Numerics/TimeStepping/explicit_lsrk.jl", "content": "# # [Single-rate Explicit Timestepping](@id Single-rate-Explicit-Timestepping)\n\n# In this tutorial, we shall explore the use of explicit Runge-Kutta\n# methods for the solution of nonautonomous (or non time-invariant) equations.\n# For our model problem, we shall reuse the rising thermal bubble\n# tutorial. See its [tutorial page](@ref Rising-Thermal-Bubble-Configuration)\n# for details on the model and parameters. For the purposes of this tutorial,\n# we will only run the experiment for a total of 100 simulation seconds.\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"tutorials\",\n \"Numerics\",\n \"TimeStepping\",\n \"tutorial_risingbubble_config.jl\",\n))\n\nFT = Float64;\n\n# After discretizing the spatial terms in the equation, the semi-discretization\n# of the governing equations have the form:\n\n# ``\n# \\begin{aligned}\n# \\frac{\\mathrm{d} \\boldsymbol{q}}{ \\mathrm{d} t} &= M^{-1}\\left(M S +\n# D^{T} M (F^{adv} + F^{visc}) + \\sum_{f=1}^{N_f} L^T M_f(\\widehat{F}^{adv} + \\widehat{F}^{visc})\n# \\right) \\equiv \\mathcal{T}(\\boldsymbol{q}).\n# \\end{aligned}\n# ``\n\n# Referencing the canonical form introduced in [Time integration](@ref\n# Time-integration) we have that in any explicit\n# formulation ``\\mathcal{F}(t, \\boldsymbol{q}) \\equiv 0`` and, in this particular\n# forumlation, ``\\mathcal{T}(t, \\boldsymbol{q}) \\equiv \\mathcal{G}(t, \\boldsymbol{q})``.\n\n# The time step restriction for an explicit method must satisfy the stable\n# [Courant number](https://en.wikipedia.org/wiki/Courant%E2%80%93Friedrichs%E2%80%93Lewy_condition)\n# for the specific time-integrator and must be selected from the following\n# constraints\n#\n# ``\n# \\Delta t_{\\mathrm{explicit}} = min \\left( \\frac{C \\Delta x_i}{u_i + a}, \\frac{C \\Delta x_i^2}{\\nu} \\right)\n# ``\n#\n# where ``C`` is the stable Courant number, ``u_i`` denotes the velocity components,\n# ``a`` the speed of sound, ``\\Delta x_i`` the grid spacing (non-uniform in case of\n# spectral element methods) along the direction ``(x_1,x_2,x_3)``, and ``\\nu`` the\n# kinematic viscosity. The first term on the right is the time step condition\n# due to the non-dissipative components, while the second term to the dissipation.\n# For explicit time-integrators, we have to find the minimum time step that\n# satisfies this condition along all three spatial directions.\n#\n# ## Runge-Kutta methods\n#\n# A single step of an ``s``-stage Runge-Kutta (RK) method for\n# solving the resulting ODE problem presented above and can be\n# expressed as the following:\n#\n# ```math\n# \\begin{align}\n# \t\\boldsymbol{q}^{n+1} = \\boldsymbol{q}^n + \\Delta t \\sum_{i=1}^{s} b_i \\mathcal{T}(\\boldsymbol{Q}^i),\n# \\end{align}\n# ```\n#\n# where ``\\boldsymbol{\\mathcal{T}}(\\boldsymbol{Q}^i)`` is the evaluation of the\n# right-hand side tendency at the stage value ``\\boldsymbol{Q}^i``, defined at\n# each stage of the RK method:\n#\n# ```math\n# \\begin{align}\n# \t\\boldsymbol{Q}^i = \\boldsymbol{q}^{n} +\n# \\Delta t \\sum_{j=1}^{s} a_{i,j}\n# \\mathcal{T}(\\boldsymbol{Q}^j).\n# \\end{align}\n# ```\n#\n# The first stage is initialized using the field at the previous time step:\n# ``\\boldsymbol{Q}^{1} \\leftarrow \\boldsymbol{q}^n``.\n#\n# In the above expressions, we define\n# ``\\boldsymbol{A} = \\lbrace a_{i,j} \\rbrace \\in \\mathbb{R}^{s\\times s}``,\n# ``\\boldsymbol{b} = \\lbrace b_i \\rbrace \\in \\mathbb{R}^s``, and\n# ``\\boldsymbol{c} = \\lbrace c_i \\rbrace \\in \\mathbb{R}^s`` as the\n# characteristic coefficients of a given RK method. This means we can\n# associate any RK method with its so-called *Butcher tableau*:\n#\n# ```math\n# \\begin{align}\n# \\begin{array}{c|c}\n# \\boldsymbol{c} &\\boldsymbol{A}\\\\\n# \\hline\n# & \\boldsymbol{b}^T\n# \\end{array} =\n# \\begin{array}{c|c c c c}\n# c_1 & a_{1,1} & a_{1,2} & \\cdots & a_{1,s}\\\\\n# c_2 & a_{2,1} & a_{2,2} & \\cdots & a_{2,s}\\\\\n# \\vdots & \\vdots & \\vdots & \\ddots & \\vdots\\\\\n# c_s & a_{s,1} & a_{s,2} & \\cdots & a_{s,s}\\\\\n# \\hline\n# & b_1 & b_2 & \\cdots & b_s\n# \\end{array}.\n# \\end{align}\n# ```\n#\n# The vector ``\\boldsymbol{c}`` is often called the *consistency vector*,\n# and is typically subjected to the row-sum condition:\n#\n# ``\n# c_i = \\sum_{j=1}^{s} a_{i,j}, \\quad \\forall i = 1, \\cdots, s.\n# ``\n#\n# This simplifies the order conditions for higher-order RK methods.\n# For more information on general RK methods, we refer the interested reader\n# to Ch. 5.2 of [Atkinson2011](@cite).\n#\n# ### [Low-storage Runge-Kutta (LSRK) methods](@id lsrk)\n# `ClimateMachine.jl` contains the following low-storage methods:\n# - Forward Euler [`LowStorageRungeKutta2N`](@ref ClimateMachine.ODESolvers.LowStorageRungeKutta2N),\n# - A 5-stage 4th-order Runge-Kutta method of Carpenter and Kennedy [`LSRK54CarpenterKennedy`](@ref ClimateMachine.ODESolvers.LSRK54CarpenterKennedy)\n# - A 14-stage 4th-order Runge-Kutta method developed by Niegemann, Diehl, and Busch [`LSRK144NiegemannDiehlBusch`](@ref ClimateMachine.ODESolvers.LSRK144NiegemannDiehlBusch).\n#\n# To start, let's try using the 5-stage method: `LSRK54CarpenterKennedy`.\n\n# As is the case for all explicit methods, we are limited by the fastest\n# propogating waves described by our governing equations. In our case,\n# these are the acoustic waves (with approximate wave speed given by the\n# speed of sound of 343 m/s).\n# For the rising bubble example used here, we use 4th order polynomials in\n# a discontinuous Galerkin approximation, with a domain resolution of\n# 125 meters in each spatial direction. This gives an effective\n# minimanl nodal distance (distance between LGL nodes) of 86 meters\n# over the entire mesh. Using the equation for the explcit time-step above,\n# we can determine the ``\\Delta t`` by specifying the desired Courant number\n# ``C`` (denoted `CFL` in the code below).\n# In our case, a heuristically determined value of 0.4 is used.\n\ntimeend = FT(100)\node_solver =\n ClimateMachine.ExplicitSolverType(solver_method = LSRK54CarpenterKennedy)\nCFL = FT(0.4)\nrun_simulation(ode_solver, CFL, timeend);\n\n# What if we wish to take a larger timestep size? We could\n# try to increase the target Courant number, say ``C = 1.7``, and\n# re-run the simulation. But this would break! In fact, in this case the numerical\n# scheme would fall outside of the stability region and the simulation would crash.\n# This occurs when the time step _exceeds_ the maximal stable time-step size\n# of the method. For the 5-stage method, one can typically get away with using\n# time-step sizes corresponding to a Courant number of ``C \\approx 0.4`` but\n# typically not much larger. In contrast, we can use an LSRK method with\n# a larger stability region. Let's illustrate this by using the 14-stage method\n# with a ``C = 1.7`` instead.\n\node_solver = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n)\nCFL = FT(1.7)\nrun_simulation(ode_solver, CFL, timeend);\n\n# And it successfully completes. Currently, the 14-stage LSRK method\n# `LSRK144NiegemannDiehlBusch` contains the largest stability region of the\n# low-storage methods available in `ClimateMachine.jl`.\n\n# ### [Strong Stability Preserving Runge--Kutta (SSPRK) methods](@id ssprk)\n\n# Just as with the LSRK methods, the SSPRK methods are self-starting,\n# with ``\\boldsymbol{Q}^{1} = \\boldsymbol{q}^n``, and stage-values are of the form\n\n# ```math\n# \\begin{align}\n# \\boldsymbol{Q}^{i+1} = a_{i,1} \\boldsymbol{q}^n\n# + a_{i,2} \\boldsymbol{Q}^{i}\n# + \\Delta t b_i\\mathcal{T}(\\boldsymbol{Q}^{i})\n# \\end{align}\n# ```\n#\n# with the value at the next step being the ``(N+1)``-th stage value\n# ``\\boldsymbol{q}^{n+1} = \\boldsymbol{Q}^{(N+1)}``. This allows the updates\n# to be performed with only three copies of the state vector\n# (storing ``\\boldsymbol{q}^n``, ``\\boldsymbol{Q}^{i}`` and ``\\mathcal{T}(\\boldsymbol{Q}^{i})``).\n# We illustrate here the use of a [`SSPRK33ShuOsher`](@ref ClimateMachine.ODESolvers.SSPRK33ShuOsher) method.\n\node_solver = ClimateMachine.ExplicitSolverType(solver_method = SSPRK33ShuOsher)\nCFL = FT(0.2)\nrun_simulation(ode_solver, CFL, timeend);\n\n# ## References\n# - [Shu1988](@cite)\n# - [Heun1900](@cite)\n" }, { "path": "tutorials/Numerics/TimeStepping/imex_ark.jl", "content": "# # [Implicit-Explicit (IMEX) Additively-Partitioned Runge-Kutta Timestepping](@id Single-rate-IMEXARK-Timestepping)\n\n# In this tutorial, we shall explore the use of IMplicit-EXplicit (IMEX) methods\n# for the solution of nonautonomous (or non time-invariant) equations.\n# For our model problem, we shall reuse the acoustic wave test in the GCM\n# configuration. See its [code](@ref Acoustic-Wave-Configuration)\n# for details on the model and parameters. For the purposes of this tutorial,\n# we will only run the experiment for a total of 3600 simulation seconds.\n# Details on this test case can be found in Sec. 4.3 of [Giraldo2013](@cite).\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"tutorials\",\n \"Numerics\",\n \"TimeStepping\",\n \"tutorial_acousticwave_config.jl\",\n));\n\n# The acoustic wave test case used in this tutorial represents a global-scale\n# problem with inertia-gravity waves traveling around the entire planet.\n# It has a hydrostatically balanced initial state that is given a pressure\n# perturbation.\n# This initial pressure perturbation causes an acoustic wave to travel to\n# the antipode, coalesce, and return to the initial position. The exact solution\n# of this test case is simple in that the (linear) acoustic theory allows one\n# to verify the analytic speed of sound based on the thermodynamics variables.\n# The initial condition is defined as a hydrostatically balanced atmosphere\n# with background (reference) potential temperature.\n\n# To fully demonstrate the advantages of using an IMEX scheme over fully explicit\n# schemes, we start here by going over a simple, fully explicit scheme. The\n# reader can refer to the [Single-rate Explicit Timestepping tutorial](@ref Single-rate-Explicit-Timestepping)\n# for detailes on such schemes. Here we use the the 14-stage LSRK method\n# [`LSRK144NiegemannDiehlBusch`]((@ref ClimateMachine.ODESolvers.LSRK144NiegemannDiehlBusch)), which contains the largest stability region of\n# the low-storage methods available in `ClimateMachine.jl`.\n\nFT = Float64\ntimeend = FT(100)\n\node_solver = ClimateMachine.ExplicitSolverType(\n solver_method = LSRK144NiegemannDiehlBusch,\n);\n\n# In the following example, the timestep calculation is based on the CFL condition\n# for horizontally-propogating acoustic waves. We use a Courant number ``C = 0.002``\n# (denoted by `CFL` in the code bellow) in the horizontal, which corresponds\n# to a timestep size of approximately ``1`` second.\n\nCFL = FT(0.002)\ncfl_direction = HorizontalDirection()\nrun_acousticwave(ode_solver, CFL, cfl_direction, timeend);\n\n# However, as it is imaginable, for real-world climate processes a time step\n# of 1 second would lead to extemely long time-to-solution simulations.\n# How can we do better? To be able to take larger time step, we can treat the\n# most restrictive wave speeds (vertical acoustic) implicitly rather than\n# explicitly. This motivates the use of an IMplicit-EXplicit (IMEX) methods.\n\n# In general, a single step of an ``s``-stage, ``N``-part additive RK method\n# (`ARK_N`) is defined by its generalized Butcher tableau:\n\n# ```math\n# \\begin{align}\n# \\begin{array}{c|c|c|c}\n# \\boldsymbol{c} &\\boldsymbol{A}_{1} & \\cdots & \\boldsymbol{A}_{N}\\\\\n# \\hline\n# & \\boldsymbol{b}_1^T & \\cdots & \\boldsymbol{b}_N^T\\\\\n# \\hline\n# & \\widehat{\\boldsymbol{b}}_1^T & \\cdots & \\widehat{\\boldsymbol{b}}_N^T\n# \\end{array} =\n# \\begin{array}{c|c c c | c | c c c }\n# c_1 & a^{[ 1 ]}_{1,1} & \\cdots & a^{[ 1 ]}_{1,s} & \\cdots\n# & a^{[ \\nu ]}_{1,1} & \\cdots & a^{[ \\nu ]}_{1,s}\\\\\n# \\vdots & \\vdots & \\ddots & \\vdots & \\cdots\n# & \\vdots & \\ddots & \\vdots \\\\\n# c_s & a^{[ 1 ]}_{s,1} & \\cdots & a^{[ 1 ]}_{s,s} & \\cdots\n# & a^{[ \\nu ]}_{s,1} & \\cdots & a^{[ \\nu ]}_{s,s}\\\\\n# \\hline\n# & b^{[ 1 ]}_1 & \\cdots & b^{[ 1 ]}_s & \\cdots\n# & b^{[ \\nu ]}_1 & \\cdots & b^{[ \\nu ]}_s\\\\\n# \\hline\n# & \\widehat{b}^{[ 1 ]}_1 & \\cdots & \\widehat{b}^{[ 1 ]}_s &\n# & \\widehat{b}^{[ \\nu ]}_1 & \\cdots & \\widehat{b}^{[ \\nu ]}_s\n# \\end{array}\n# \\end{align}\n# ```\n\n# and is given by\n\n# ``\n# \t\\boldsymbol{q}^{n+1} = \\boldsymbol{q}^n + \\Delta t \\left( \\underbrace{\\sum_{i=1}^{s}}_{\\textrm{Stages}} \\underbrace{\\sum_{\\nu=1}^{N}}_{\\textrm{Components}} b_i^{[ \\nu ]} {\\mathcal{T}}^{[ \\nu ]}(\\boldsymbol{Q}^i)) \\right)\n# ``\n\n# where ``s`` denotes the stages and ``N`` the components, and where the stage values are given by:\n#\n# ``\n# \t\\boldsymbol{Q}^i = \\boldsymbol{q}^n + \\Delta t \\sum_{j=1}^{s} \\sum_{\\nu = 1}^{N} a_{i,j}^{[ \\nu ]}\n# \t{\\mathcal{T}}^{[ \\nu]}(\\boldsymbol{Q}^j).\n# ``\n#\n# Similar to standard RK methods, the stage vectors are approximations to the state at each stage\n# of the ARK method. Moreover, the temporal coefficients ``c_i`` satisfy a similar\n# row-sum condition, holding for all ``\\nu = 1, \\cdots, N``:\n\n# ``\n# \tc_i = \\sum_{j=1}^{s} a_{i, j}^{[ \\nu ]}, \\quad \\forall \\nu = 1, \\cdots, N.\n# ``\n#\n# The Butcher coefficients ``\\boldsymbol{c}``, ``\\boldsymbol{b}_{\\nu}``, ``\\boldsymbol{A}_{\\nu}``, and ``\\widehat{\\boldsymbol{b}}_{\\nu}``\n# are constrained by certain accuracy and stability requirements, which are summarized in\n# [Kennedy2001](@cite).\n\n# A common setting is the case ``N = 2``. This gives the typical context for\n# Implicit-Explicit (IMEX) splitting methods, where the tendency ``{\\mathcal{T}}``\n# is assumed to have the decomposition:\n#\n# ``\n# \t\\dot{\\boldsymbol{q}} = \\mathcal{T}(\\boldsymbol{q}) \\equiv\n# \t{\\mathcal{T}}_{s}(\\boldsymbol{q}) + {\\mathcal{T}}_{ns}(\\boldsymbol{q}),\n# ``\n# where the right-hand side has been split into a \"stiff\" component ``{\\mathcal{T}}_{s}``,\n# to be treated implicitly, and a non-stiff part ``{\\mathcal{T}}_{ns}`` to be treated explicitly.\n\n# Referencing the canonical form introduced in [Time integration](@ref\n# Time-integration) we have that in this particular forumlation\n# ``\\mathcal{T}_{ns}(t, \\boldsymbol{q}) \\equiv \\mathcal{G}(t, \\boldsymbol{q})`` and\n# ``\\mathcal{T}_{s}(t, \\boldsymbol{q}) \\equiv \\mathcal{F}(t, \\boldsymbol{q})``.\n\n# Two different RK methods are applied to ``{\\mathcal{T}}_{s}`` and ``{\\mathcal{T}}_{ns}``\n# separately, which have been specifically designed and coupled. Examples can be found in\n# [Giraldo2013](@cite). The Butcher Tableau for an `ARK_2` method will have the\n# form\n#\n# ```math\n# \\begin{align}\n# \\begin{array}{c|c|c}\n# \\boldsymbol{c} &\\boldsymbol{A}_E &\\boldsymbol{A}_I\\\\\n# \\hline\n# & \\boldsymbol{b}_E^T & \\boldsymbol{b}_I^T \\\\\n# \\hline\n# & \\widehat{\\boldsymbol{b}}_E^T & \\widehat{\\boldsymbol{b}}_I^T\n# \\end{array},\n# \\end{align}\n# ```\n#\n# with\n#\n# ``\n# \t\\boldsymbol{A}_O = \\left\\lbrace a_{i, j}^O \\right\\rbrace, \\quad\n# \t\\boldsymbol{b}_O = \\left\\lbrace b_{i}^O \\right\\rbrace, \\quad\n# \t\\widehat{\\boldsymbol{b}}_O = \\left\\lbrace \\widehat{b}_{i}^O \\right\\rbrace,\n# ``\n#\n# where ``O`` denotes the label (either ``E`` for explicit or ``I`` for implicit).\n\n# For the acoustic wave example used here, we use 4th order polynomials in\n# our discontinuous Galerkin approximation, with 6 elements in each horizontal\n# direction and 4 elements in the vertical direction, on the cubed-sphere.\n# This gives an effective minimal node-distance (distance between LGL nodes)\n# of roughly 203000 m.\n# As in the [previous tutorial](@ref Single-rate-Explicit-Timestepping),\n# we can determine our ``\\Delta t`` by specifying our desired horizontal\n# Courant number ``C`` (the timestep calculation is based on the CFL condition\n# for horizontally-propogating acoustic waves). In this very simple test case,\n# we can use a value of 0.5, which corresponds to a time-step size of\n# around 257 seconds. But for this particular example, even higher values\n# might work.\n\ntimeend = FT(3600)\node_solver = ClimateMachine.IMEXSolverType(\n solver_method = ARK2GiraldoKellyConstantinescu,\n)\nCFL = FT(0.5)\ncfl_direction = HorizontalDirection()\nrun_acousticwave(ode_solver, CFL, cfl_direction, timeend);\n\n# ## References\n# - [Giraldo2013](@cite)\n# - [Kennedy2001](@cite)\n" }, { "path": "tutorials/Numerics/TimeStepping/mis.jl", "content": "# # [Multirate Infinitesimal Step (MIS) Timestepping](@id MIS-Timestepping)\n\n# In this tutorial, we shall explore the use of Multirate Infinitesimal Step\n# (MIS) methods for the solution of nonautonomous (or non time-invariant) equations.\n# For our model problem, we shall reuse the rising thermal bubble\n# tutorial. See its [tutorial page](@ref Rising-Thermal-Bubble-Configuration)\n# for details on the model and parameters. For the purposes of this tutorial,\n# we will only run the experiment for a total of 500 simulation seconds.\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"tutorials\",\n \"Numerics\",\n \"TimeStepping\",\n \"tutorial_risingbubble_config.jl\",\n))\n\nFT = Float64;\n\n# Referencing the formulation introduced in the previous\n# [Multirate RK methods tutorial](@ref Multirate-RK-Timestepping), we can\n# describe Multirate Infinitesimal Step (MIS) methods by\n\n# ```math\n# \\begin{align}\n# v_i (0)\n# &= q^n + \\sum_{j=1}^{i-1} \\alpha_{ij} (Q^{(j)} - q^n) \\\\\n# \\frac{dv_i}{d\\tau}\n# &= \\sum_{j=1}^{i-1} \\frac{\\gamma_{ij}}{d_i \\Delta t} (Q^{(j)} - q^n)\n# + \\sum_{j=1}^i \\frac{\\beta_{ij}}{d_i} \\mathcal{T}_S (Q^{(j)}, t + \\Delta t c_i)\n# + \\mathcal{T}_F(v_i, t^n + \\Delta t \\tilde c_i + \\frac{c_i - \\tilde c_i}{d_i} \\tau),\n# \\quad \\tau \\in [0, \\Delta t d_i] \\\\\n# Q^{(i)} &= v_i(\\Delta t d_i),\n# \\end{align}\n# ```\n#\n# where we have used the the stage values ``Q^{(i)} = v_i(\\tau_i)`` as the\n# solution to the _inner_ ODE problem, ``{\\mathcal{T}_{s}}``\n# for the slow component, and ``{\\mathcal{T}_{f}}` for the fast\n# one, as in the [Multirate RK methods tutorial](@ref Multirate-RK-Timestepping).\n\n# Referencing the canonical form introduced in [Time integration](@ref\n# Time-integration), both ``{\\mathcal{T}_{f}}`` and ``{\\mathcal{T}_{s}}``\n# could be discretized either explicitly or implicitly, hence, they could\n# belong to either ``\\mathcal{F}(t, \\boldsymbol{q})`` or ``\\mathcal{G}(t, \\boldsymbol{q})``\n# term.\n#\n# The method is defined in terms of the lower-triangular matrices ``\\alpha``,\n# ``\\beta`` and ``\\gamma``, with ``d_i = \\sum_j \\beta_{ij}``,\n# ``c_i = (I - \\alpha - \\gamma)^{-1} d`` and ``\\tilde c = \\alpha c``.\n# More details can be found in [WenschKnothGalant2009](@cite) and\n# [KnothWensch2014](@cite).\n\node_solver = ClimateMachine.MISSolverType(;\n mis_method = MIS2,\n fast_method = LSRK144NiegemannDiehlBusch,\n nsubsteps = (40,),\n)\n\ntimeend = FT(500)\nCFL = FT(20)\nrun_simulation(ode_solver, CFL, timeend);\n\n# The reader can compare the Courant number (denoted by `CFL` in the code snippet)\n# used in this example, with the adopted in the\n# [single-rate explicit timestepping tutorial page](@ref Single-rate-Explicit-Timestepping)\n# in which we use the same scheme as the fast method employed in this case,\n# and notice that with this MIS method we are able to take a much larger\n# Courant number.\n\n# ## References\n# - [WenschKnothGalant2009](@cite)\n# - [KnothWensch2014](@cite)\n" }, { "path": "tutorials/Numerics/TimeStepping/multirate_rk.jl", "content": "# # [Multirate Runge-Kutta Timestepping](@id Multirate-RK-Timestepping)\n\n# In this tutorial, we shall explore the use of Multirate Runge-Kutta\n# methods for the solution of nonautonomous (or non time-invariant) equations.\n# For our model problem, we shall reuse the acoustic wave test in the GCM\n# configuration. See its [code](@ref Acoustic-Wave-Configuration)\n# for details on the model and parameters. For the purposes of this tutorial,\n# we will only run the experiment for a total of 3600 simulation seconds.\n# Details on this test case can be found in Sec. 4.3 of [Giraldo2013](@cite).\n\nusing ClimateMachine\nconst clima_dir = dirname(dirname(pathof(ClimateMachine)));\ninclude(joinpath(\n clima_dir,\n \"tutorials\",\n \"Numerics\",\n \"TimeStepping\",\n \"tutorial_acousticwave_config.jl\",\n))\n\nFT = Float64;\n\n# The typical context for Multirate splitting methods is given by problems\n# in which the tendency ``\\mathcal{T}`` is assumed to have single parts that\n# operate on different time rates (such as a slow time scale and a fast time scale).\n# A general form is given by\n#\n# ``\n# \t\\dot{\\boldsymbol{q}} = \\mathcal{T}(\\boldsymbol{q}) \\equiv\n# \t{\\mathcal{T}}_{f}(\\boldsymbol{q}) + {\\mathcal{T}}_{s}(\\boldsymbol{q}),\n# ``\n#\n# where the right-hand side has been split into a \"fast\" component ``{\\mathcal{T}_{f}}``,\n# and a \"slow\" component ``{\\mathcal{T}_{s}}``.\n\n# Referencing the canonical form introduced in [Time integration](@ref\n# Time-integration), both ``{\\mathcal{T}_{f}}`` and ``{\\mathcal{T}_{s}}``\n# could be discretized either explicitly or implicitly, hence, they could\n# belong to either ``\\mathcal{F}(t, \\boldsymbol{q})`` or ``\\mathcal{G}(t, \\boldsymbol{q})``\n# term.\n#\n# For a given time-step size ``\\Delta t``, the two-rate method in [Schlegel2009](@cite)\n# is summarized as the following:\n#\n# ```math\n# \\begin{align}\n# \\boldsymbol{Q}_1 &= \\boldsymbol{q}(t_n), \\\\\n# \\boldsymbol{r}_{i} &= \\sum_{j=1}^{i-1}\\tilde{a}^O_{ij} {\\mathcal{T}_{s}}(\\boldsymbol{Q}_{j}), \\\\\n# \\boldsymbol{w}_{i,1} &= \\boldsymbol{Q}_{i-1},\\\\\n# \\boldsymbol{w}_{i,k} &= \\boldsymbol{w}_{i,k-1} + \\Delta t \\tilde{c}_i^O \\sum_{j=1}^{k-1}\\tilde{a}^I_{k,j}\n# \\left(\\frac{1}{\\tilde{c}_i^O}\\boldsymbol{r}_i + {\\mathcal{T}_{f}}(\\boldsymbol{w}_{i,j})\\right),\\\\\n# & \\quad\\quad i = 2, \\cdots, s^O + 1 \\text{ and } k = 2, \\cdots, s^I + 1,\\nonumber \\\\\n# \\boldsymbol{Q}_i &= \\boldsymbol{w}_{i,s^I + 1}\n# \\end{align}\n# ```\n#\n# where the tilde parameters denote increments per RK stage:\n#\n# ```math\n# \\begin{align}\n# \\tilde{a}_{ij} &= \\begin{cases}\n# a_{i,j} - a_{i-1, j} & \\text{if } i < s + 1 \\\\\n# b_j - a_{s,j} & \\text{if } i = s + 1\n# \\end{cases},\\\\\n# \\tilde{c}_{i} &= \\begin{cases}\n# c_{i} - c_{i-1} & \\text{if } i < s + 1 \\\\\n# 1 - c_{s} & \\text{if } i = s + 1\n# \\end{cases},\n# \\end{align}\n# ```\n#\n# where the coefficients ``a``, ``b``, and ``c`` correspond to the Butcher\n# tableau for a given RK method. The superscripts ``O`` and ``I`` denote the\n# *outer* (slow) and *inner* (fast) components of the multirate method\n# respectively. Thus, tilde coefficients should be associated with the RK\n# method indicated by the superscripts. In other words, the RK methods\n# for the slow ``{\\mathcal{T}_{s}}`` and fast\n# ``{\\mathcal{T}_{f}}`` components have Butcher tables given by:\n#\n# ```math\n# \\begin{align}\n# \\begin{array}{c|c}\n# \\boldsymbol{c}_{O} &\\boldsymbol{A}_{O} \\\\\n# \\hline\n# & \\boldsymbol{b}_O^T\n# \\end{array}, \\quad\n# \\begin{array}{c|c}\n# \\boldsymbol{c}_{I} &\\boldsymbol{A}_{I} \\\\\n# \\hline\n# & \\boldsymbol{b}_I^T\n# \\end{array},\n# \\end{align}\n# ```\n#\n# where ``\\boldsymbol{A}_O = \\lbrace a_{i,j}^O\\rbrace``, ``\\boldsymbol{b}_O = \\lbrace b_i^O \\rbrace``, and\n# ``c_O = \\lbrace c_i^O \\rbrace`` (similarly for ``\\boldsymbol{A}_I``, ``\\boldsymbol{b}_I``, and ``\\boldsymbol{c}_I``).\n# The method described here is for an explicit RK outer method with ``s`` stages.\n# More details can be found in [Schlegel2012](@cite).\n\n# The acoustic wave test case used in this tutorial represents a global-scale\n# problem with inertia-gravity waves traveling around the entire planet.\n# It has a hydrostatically balanced initial state that is given a pressure\n# perturbation.\n# This initial pressure perturbation causes an acoustic wave to travel to\n# the antipode, coalesce, and return to the initial position. The exact solution\n# of this test case is simple in that the (linear) acoustic theory allows one\n# to verify the analytic speed of sound based on the thermodynamics variables.\n# The initial condition is defined as a hydrostatically balanced atmosphere\n# with background (reference) potential temperature.\n\node_solver = ClimateMachine.MultirateSolverType(\n splitting_type = ClimateMachine.HEVISplitting(),\n slow_method = LSRK54CarpenterKennedy,\n fast_method = ARK2GiraldoKellyConstantinescu,\n implicit_solver_adjustable = true,\n timestep_ratio = 100,\n)\n\ntimeend = FT(3600)\nCFL = FT(5)\ncfl_direction = HorizontalDirection()\nrun_acousticwave(ode_solver, CFL, cfl_direction, timeend);\n\n# The interested reader can explore the combination of different slow and\n# fast methods for Multirate solvers, consulting the ones available in\n# `ClimateMachine.jl`, such as the\n# [`Low-Storage-Runge-Kutta-methods`](@ref ClimateMachine.ODESolvers.LowStorageRungeKutta2N),\n# [`Strong-Stability-Preserving-RungeKutta-methods`](@ref ClimateMachine.ODESolvers.StrongStabilityPreservingRungeKutta),\n# and [`Additive-Runge-Kutta-methods`](@ref ClimateMachine.ODESolvers.AdditiveRungeKutta).\n\n# ## References\n# - [Giraldo2013](@cite)\n# - [Schlegel2009](@cite)\n# - [Schlegel2012](@cite)\n" }, { "path": "tutorials/Numerics/TimeStepping/ts_intro.jl", "content": "# # [Time integration](@id Time-integration)\n\n# Time integration methods for the numerical solution of Ordinary Differential\n# Equations (ODEs), also called timesteppers, can be of different nature and\n# flavor (e.g., explicit, semi-implicit, single-stage, multi-stage, single-step,\n# multi-step, single-rate, multi-rate, etc). ClimateMachine supports several\n# of them. Before showing the different nature of some of these methods, let us\n# introduce some common notation.\n\n# A commonly used notation for Initial Value Problems (IVPs) is:\n\n# ```math\n# \\begin{align}\n# \\frac{\\mathrm{d} \\boldsymbol{q}}{ \\mathrm{d} t} &= \\mathcal{T}(t, \\boldsymbol{q}),\\\\\n# \\boldsymbol{q}(t_0) &= \\boldsymbol{q_0},\n# \\end{align}\n# ```\n\n# where ``\\boldsymbol{q}`` is an unknown function (vector in most of our cases)\n# of time ``t``, which we would like to approximate, and at the initial time ``t_0``\n# the corresponding initial value ``\\boldsymbol{q}_0`` is given.\n\n# The given general formulation, is suitable for single-step explicit schemes.\n# Generally, the equation can be represented in the following canonical form:\n\n# ```math\n# \\begin{align}\n# \\dot {\\boldsymbol{q}} + \\mathcal{F}(t, \\boldsymbol{q}) &= \\mathcal{G}(t, \\boldsymbol{q}),\n# \\end{align}\n# ```\n\n# where we have used ``\\dot {\\boldsymbol{q}} = d \\boldsymbol{q} / dt``.\n# We refer to the term ``\\mathcal{G}``\n# as the right-hand-side (RHS) or explicit term, and to the spatial terms of\n# ``\\mathcal{F}`` as the left-hand-side (LHS) or implicit term.\n" }, { "path": "tutorials/Numerics/TimeStepping/tutorial_acousticwave_config.jl", "content": "# # [Acoustic Wave Configuration](@id Acoustic-Wave-Configuration)\n\n#\n# In this example, we demonstrate the usage of the `ClimateMachine`\n# [AtmosModel](@ref AtmosModel-docs) machinery to solve the fluid\n# dynamics of an acoustic wave.\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.Checkpoint\nusing ClimateMachine.ConfigTypes\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\nusing ClimateMachine.Grids\nusing ClimateMachine.ODESolvers\n\nusing CLIMAParameters\n\nusing StaticArrays\n\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet()\n\nBase.@kwdef struct AcousticWaveSetup{FT}\n domain_height::FT = 10e3\n T_ref::FT = 300\n α::FT = 3\n γ::FT = 100\n nv::Int = 1\nend\n\nfunction (setup::AcousticWaveSetup)(problem, bl, state, aux, localgeo, t)\n ## callable to set initial conditions\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n λ = longitude(bl, aux)\n φ = latitude(bl, aux)\n z = altitude(bl, aux)\n\n β = min(FT(1), setup.α * acos(cos(φ) * cos(λ)))\n f = (1 + cos(FT(π) * β)) / 2\n g = sin(setup.nv * FT(π) * z / setup.domain_height)\n Δp = setup.γ * f * g\n p = aux.ref_state.p + Δp\n\n ts = PhaseDry_pT(param_set, p, setup.T_ref)\n q_pt = PhasePartition(ts)\n e_pot = gravitational_potential(bl.orientation, aux)\n e_int = internal_energy(ts)\n\n state.ρ = air_density(ts)\n state.ρu = SVector{3, FT}(0, 0, 0)\n state.energy.ρe = state.ρ * (e_int + e_pot)\n return nothing\nend\n\nfunction run_acousticwave(\n ode_solver_type,\n CFL::FT,\n CFL_direction,\n timeend::FT,\n) where {FT}\n\n ## DG polynomial orders\n N = (4, 4)\n\n ## Domain resolution\n nelem_horz = 6\n nelem_vert = 4\n resolution = (nelem_horz, nelem_vert)\n\n t0 = FT(0)\n\n setup = AcousticWaveSetup{FT}()\n T_profile = IsothermalProfile(param_set, setup.T_ref)\n ref_state = HydrostaticState(T_profile)\n turbulence = ConstantDynamicViscosity(FT(0))\n physics = AtmosPhysics{FT}(\n param_set;\n ref_state = ref_state,\n turbulence = turbulence,\n moisture = DryModel(),\n )\n model = AtmosModel{FT}(\n AtmosGCMConfigType,\n physics;\n init_state_prognostic = setup,\n source = (Gravity(),),\n )\n\n driver_config = ClimateMachine.AtmosGCMConfiguration(\n \"GCM Driver: Acoustic wave test\",\n N,\n resolution,\n setup.domain_height,\n param_set,\n setup;\n model = model,\n )\n\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n Courant_number = CFL,\n init_on_cpu = true,\n ode_solver_type = ode_solver_type,\n CFL_direction = CFL_direction,\n )\n\n ClimateMachine.invoke!(solver_config)\nend\n" }, { "path": "tutorials/Numerics/TimeStepping/tutorial_risingbubble_config.jl", "content": "# # [Rising Thermal Bubble Configuration](@id Rising-Thermal-Bubble-Configuration)\n#\n# In this example, we demonstrate the usage of the `ClimateMachine`\n# [AtmosModel](@ref AtmosModel-docs) machinery to solve the fluid\n# dynamics of a thermal perturbation in a neutrally stratified background state\n# defined by its uniform potential temperature. We solve a flow in a box configuration -\n# this is representative of a large-eddy simulation. Several versions of the problem\n# setup may be found in literature, but the general idea is to examine the\n# vertical ascent of a thermal bubble (we can interpret these as simple\n# representation of convective updrafts).\n#\n# ## Description of experiment\n# 1) Dry Rising Bubble (circular potential temperature perturbation)\n# 2) Boundaries\n# Top and Bottom boundaries:\n# - `Impenetrable(FreeSlip())` - Top and bottom: no momentum flux, no mass flux through\n# walls.\n# - `Impermeable()` - non-porous walls, i.e. no diffusive fluxes through\n# walls.\n# Lateral boundaries\n# - Laterally periodic\n# 3) Domain - 2500m (horizontal) x 2500m (horizontal) x 2500m (vertical)\n# 4) Resolution - 50m effective resolution\n# 5) Total simulation time - 1000s\n# 6) Mesh Aspect Ratio (Effective resolution) 1:1\n# 7) Overrides defaults for\n# - CPU Initialisation\n# - Time integrator\n# - Sources\n# - Smagorinsky Coefficient\n\n#md # !!! note\n#md # This experiment setup assumes that you have installed the\n#md # `ClimateMachine` according to the instructions on the landing page.\n#md # We assume the users' familiarity with the conservative form of the\n#md # equations of motion for a compressible fluid (see the\n#md # [AtmosModel](@ref AtmosModel-docs) page).\n#md #\n#md # The following topics are covered in this example\n#md # - Package requirements\n#md # - Defining a `model` subtype for the set of conservation equations\n#md # - Defining the initial conditions\n#md # - Applying source terms\n#md # - Choosing a turbulence model\n#md # - Adding tracers to the model\n#md # - Choosing a time-integrator\n#md # - Choosing diagnostics (output) configurations\n#md #\n#md # The following topics are not covered in this example\n#md # - Defining new boundary conditions\n#md # - Defining new turbulence models\n#md # - Building new time-integrators\n#md # - Adding diagnostic variables (beyond a standard pre-defined list of\n#md # variables)\n\n# ## [Loading code](@id Loading-code-rtb)\n\n# Before setting up our experiment, we recognize that we need to import some\n# pre-defined functions from other packages. Julia allows us to use existing\n# modules (variable workspaces), or write our own to do so. Complete\n# documentation for the Julia module system can be found\n# [here](https://docs.julialang.org/en/v1/manual/modules/#).\n\n# We need to use the `ClimateMachine` module! This imports all functions\n# specific to atmospheric and ocean flow modeling.\n\nusing ClimateMachine\nClimateMachine.init()\nusing ClimateMachine.Atmos\nusing ClimateMachine.Orientations\nusing ClimateMachine.ConfigTypes\nusing ClimateMachine.Diagnostics\nusing ClimateMachine.GenericCallbacks\nusing ClimateMachine.ODESolvers\nusing Thermodynamics.TemperatureProfiles\nusing Thermodynamics\nusing ClimateMachine.TurbulenceClosures\nusing ClimateMachine.VariableTemplates\n\n# In ClimateMachine we use `StaticArrays` for our variable arrays.\n# We also use the `Test` package to help with unit tests and continuous\n# integration systems to design sensible tests for our experiment to ensure new\n# / modified blocks of code don't damage the fidelity of the physics. The test\n# defined within this experiment is not a unit test for a specific\n# subcomponent, but ensures time-integration of the defined problem conditions\n# within a reasonable tolerance. Immediately useful macros and functions from\n# this include `@test` and `@testset` which will allow us to define the testing\n# parameter sets.\nusing StaticArrays\nusing Test\nusing CLIMAParameters\nusing CLIMAParameters.Atmos.SubgridScale: C_smag\nusing CLIMAParameters.Planet: R_d, cp_d, cv_d, MSLP, grav\nstruct EarthParameterSet <: AbstractEarthParameterSet end\nconst param_set = EarthParameterSet();\n\n# ## [Initial Conditions](@id init-rtb)\n# This example demonstrates the use of functions defined\n# in the [`Thermodynamics`](@ref Thermodynamics) module to\n# generate the appropriate initial state for our problem.\n\n#md # !!! note\n#md # The following variables are assigned in the initial condition\n#md # - `state.ρ` = Scalar quantity for initial density profile\n#md # - `state.ρu`= 3-component vector for initial momentum profile\n#md # - `state.energy.ρe`= Scalar quantity for initial total-energy profile\n#md # humidity\n#md # - `state.tracers.ρχ` = Vector of four tracers (here, for demonstration\n#md # only; we can interpret these as dye injections for visualization\n#md # purposes)\nfunction init_risingbubble!(problem, bl, state, aux, localgeo, t)\n (x, y, z) = localgeo.coord\n\n ## Problem float-type\n FT = eltype(state)\n param_set = parameter_set(bl)\n\n ## Unpack constant parameters\n R_gas::FT = R_d(param_set)\n c_p::FT = cp_d(param_set)\n c_v::FT = cv_d(param_set)\n p0::FT = MSLP(param_set)\n _grav::FT = grav(param_set)\n γ::FT = c_p / c_v\n\n ## Define bubble center and background potential temperature\n xc::FT = 5000\n yc::FT = 1000\n zc::FT = 2000\n r = sqrt((x - xc)^2 + (z - zc)^2)\n rc::FT = 2000\n θamplitude::FT = 2\n\n ## This is configured in the reference hydrostatic state\n ref_state = reference_state(bl)\n θ_ref::FT = ref_state.virtual_temperature_profile.T_surface\n\n ## Add the thermal perturbation:\n Δθ::FT = 0\n if r <= rc\n Δθ = θamplitude * (1.0 - r / rc)\n end\n\n ## Compute perturbed thermodynamic state:\n θ = θ_ref + Δθ ## potential temperature\n π_exner = FT(1) - _grav / (c_p * θ) * z ## exner pressure\n ρ = p0 / (R_gas * θ) * (π_exner)^(c_v / R_gas) ## density\n T = θ * π_exner\n e_int = internal_energy(param_set, T)\n ts = PhaseDry(param_set, e_int, ρ)\n ρu = SVector(FT(0), FT(0), FT(0)) ## momentum\n ## State (prognostic) variable assignment\n e_kin = FT(0) ## kinetic energy\n e_pot = gravitational_potential(bl, aux) ## potential energy\n ρe_tot = ρ * total_energy(e_kin, e_pot, ts) ## total energy\n\n ρχ = FT(0) ## tracer\n\n ## We inject tracers at the initial condition at some specified z coordinates\n if 500 < z <= 550\n ρχ += FT(0.05)\n end\n\n ## We want 4 tracers\n ntracers = 4\n\n ## Define 4 tracers, (arbitrary scaling for this demo problem)\n ρχ = SVector{ntracers, FT}(ρχ, ρχ / 2, ρχ / 3, ρχ / 4)\n\n ## Assign State Variables\n state.ρ = ρ\n state.ρu = ρu\n state.energy.ρe = ρe_tot\n state.tracers.ρχ = ρχ\nend\n\n# ## [Model Configuration](@id config-helper)\n# We define a configuration function to assist in prescribing the physical\n# model. The purpose of this is to populate the\n# `ClimateMachine.AtmosLESConfiguration` with arguments\n# appropriate to the problem being considered.\nfunction config_risingbubble(\n ::Type{FT},\n N,\n resolution,\n xmax,\n ymax,\n zmax,\n) where {FT}\n ## Since we want four tracers, we specify this and include the appropriate\n ## diffusivity scaling coefficients (normally these would be physically\n ## informed but for this demonstration we use integers corresponding to the\n ## tracer index identifier)\n ntracers = 4\n δ_χ = SVector{ntracers, FT}(1, 2, 3, 4)\n ## To assemble `AtmosModel` with no tracers, set `tracers = NoTracers()`.\n\n ## The model coefficient for the turbulence closure is defined via the\n ## [CLIMAParameters\n ## package](https://CliMA.github.io/CLIMAParameters.jl/latest/) A reference\n ## state for the linearisation step is also defined.\n T_surface = FT(300)\n T_min_ref = FT(0)\n T_profile = DryAdiabaticProfile{FT}(param_set, T_surface, T_min_ref)\n ref_state = HydrostaticState(T_profile)\n\n ## Here we assemble the `AtmosModel`.\n _C_smag = FT(C_smag(param_set))\n physics = AtmosPhysics{FT}(\n param_set; ## Parameter set corresponding to earth parameters\n ref_state = ref_state, ## Reference state\n turbulence = SmagorinskyLilly(_C_smag), ## Turbulence closure model\n moisture = DryModel(), ## Exclude moisture variables\n tracers = NTracers{ntracers, FT}(δ_χ), ## Tracer model with diffusivity coefficients\n )\n\n model = AtmosModel{FT}(\n AtmosLESConfigType, ## Flow in a box, requires the AtmosLESConfigType\n physics; ## Atmos physics\n init_state_prognostic = init_risingbubble!, ## Apply the initial condition\n source = (Gravity(),), ## Gravity is the only source term here\n )\n\n ## Finally, we pass a `Problem Name` string, the mesh information, and the\n ## model type to the [`AtmosLESConfiguration`] object.\n config = ClimateMachine.AtmosLESConfiguration(\n \"DryRisingBubble\", ## Problem title [String]\n N, ## Polynomial order [Int]\n resolution, ## (Δx, Δy, Δz) effective resolution [m]\n xmax, ## Domain maximum size [m]\n ymax, ## Domain maximum size [m]\n zmax, ## Domain maximum size [m]\n param_set, ## Parameter set.\n init_risingbubble!, ## Function specifying initial condition\n model = model, ## Model type\n )\n return config\nend\n\n#md # !!! note\n#md # `Keywords` are used to specify some arguments (see appropriate source\n#md # files).\n\n# ## Diagnostics\n# Here we define the diagnostic configuration specific to this problem.\nfunction config_diagnostics(driver_config)\n interval = \"10000steps\"\n dgngrp = setup_atmos_default_diagnostics(\n AtmosLESConfigType(),\n interval,\n driver_config.name,\n )\n return ClimateMachine.DiagnosticsConfiguration([dgngrp])\nend\n\nfunction run_simulation(ode_solver_type, CFL::FT, timeend::FT) where {FT}\n\n ## We need to specify the polynomial order for the DG discretization,\n ## effective resolution, simulation end-time, the domain bounds, and the\n ## courant-number for the time-integrator. Note how the time-integration\n ## components `solver_config` are distinct from the spatial / model\n ## components in `driver_config`. `init_on_cpu` is a helper keyword argument\n ## that forces problem initialization on CPU (thereby allowing the use of\n ## random seeds, spline interpolants and other special functions at the\n ## initialization step.)\n N = 4\n Δh = FT(125)\n Δv = FT(125)\n resolution = (Δh, Δh, Δv)\n xmax = FT(10000)\n ymax = FT(500)\n zmax = FT(10000)\n t0 = FT(0)\n ## timeend = FT(100)\n ## For full simulation set `timeend = 1000`\n\n driver_config = config_risingbubble(FT, N, resolution, xmax, ymax, zmax)\n solver_config = ClimateMachine.SolverConfiguration(\n t0,\n timeend,\n driver_config,\n init_on_cpu = true,\n Courant_number = CFL,\n ode_solver_type = ode_solver_type,\n )\n dgn_config = config_diagnostics(driver_config)\n @show solver_config.ode_solver_type\n\n ## Invoke solver (calls `solve!` function for time-integrator), pass the driver,\n ## solver and diagnostic config information.\n result = ClimateMachine.invoke!(\n solver_config;\n diagnostics_config = dgn_config,\n user_callbacks = (),\n check_euclidean_distance = true,\n )\n return result\nend\n" }, { "path": "tutorials/Ocean/geostrophic_adjustment.jl", "content": "# # Geostrophic adjustment in the hydrostatic Boussinesq equations\n#\n# This example simulates a one-dimensional geostrophic adjustement problem\n# using the `ClimateMachine.Ocean` subcomponent to solve the hydrostatic\n# Boussinesq equations.\n#\n# First we `ClimateMachine.init()`.\n\nusing ClimateMachine\n\nClimateMachine.init()\n\n# # Domain setup\n#\n# We formulate our problem in a Cartesian domain 100 km in ``x, y`` and 400 m\n# deep, and discretized on a grid with 100 fourth-order elements in ``x``, and 1\n# fourth-order element in ``y, z``,\n\nusing ClimateMachine.CartesianDomains\n\ndomain = RectangularDomain(\n Ne = (25, 1, 1),\n Np = 4,\n x = (0, 1e6),\n y = (0, 1e6),\n z = (-400, 0),\n periodicity = (false, true, false),\n)\n\n# # Physical parameters\n#\n# We use a Coriolis parameter appropriate for mid-latitudes,\n\nf = 1e-4 # s⁻¹, Coriolis parameter\nnothing # hide\n\n# and Earth's gravitational acceleration,\n\nusing CLIMAParameters: AbstractEarthParameterSet, Planet\nstruct EarthParameters <: AbstractEarthParameterSet end\n\ng = Planet.grav(EarthParameters()) # m s⁻²\n\n# # An unbalanced initial state\n#\n# We use a Gaussian, partially-balanced initial condition with parameters\n\nU = 0.1 # geostrophic velocity (m s⁻¹)\nL = domain.L.x / 40 # Gaussian width (m)\na = f * U * L / g # amplitude of the geostrophic surface displacement (m)\nx₀ = domain.L.x / 4 # Gaussian origin (m, recall that x ∈ [0, Lx])\n\n# and functional form\n\nGaussian(x, L) = exp(-x^2 / (2 * L^2))\n\n## Geostrophic ``y``-velocity: f V = g ∂_x η\nvᵍ(x, y, z) = -U * (x - x₀) / L * Gaussian(x - x₀, L)\n\n## Geostrophic surface displacement\nηᵍ(x, y, z) = a * Gaussian(x - x₀, L)\n\n# We double the initial surface displacement so that the surface is half-balanced,\n# half unbalanced,\n\nηⁱ(x, y, z) = 2 * ηᵍ(x, y, z)\n\n# In summary,\n\nusing ClimateMachine.Ocean.OceanProblems: InitialConditions\n\ninitial_conditions = InitialConditions(v = vᵍ, η = ηⁱ)\n\n@info \"\"\" Parameters for the Geostrophic adjustment problem are...\n\n Coriolis parameter: $f s⁻¹\n Gravitational acceleration: $g m s⁻²\n Geostrophic velocity: $U m s⁻¹\n Width of the initial geostrophic perturbation: $L m\n Amplitude of the initial surface perturbation: $a m\n Rossby number (U / f L): $(U / (f * L))\n\n\"\"\"\n\n# # Boundary conditions, Driver configuration, and Solver configuration\n#\n# Next, we configure the `HydrostaticBoussinesqModel` and build the `DriverConfiguration`.\n# We configure our model in a domain which is bounded in the ``x`` direction.\n# Both the boundary conditions in ``x`` and in ``z`` require boundary conditions,\n# which we define:\n\nusing ClimateMachine.Ocean:\n Impenetrable, Penetrable, FreeSlip, Insulating, OceanBC\n\nsolid_surface_boundary_conditions = OceanBC(\n Impenetrable(FreeSlip()), # Velocity boundary conditions\n Insulating(), # Temperature boundary conditions\n)\n\nfree_surface_boundary_conditions = OceanBC(\n Penetrable(FreeSlip()), # Velocity boundary conditions\n Insulating(), # Temperature boundary conditions\n)\n\nboundary_conditions =\n (solid_surface_boundary_conditions, free_surface_boundary_conditions)\n\n# We refer to these boundary conditions by their indices in the `boundary_tags` tuple\n# when specifying the boundary conditions for the `state`; in other words, \"1\" corresponds to\n# `solid_surface_boundary_conditions`, while `2` corresponds to `free_surface_boundary_conditions`,\n\nboundary_tags = (\n (1, 1), # (west, east) boundary conditions\n (0, 0), # (south, north) boundary conditions\n (1, 2), # (bottom, top) boundary conditions\n)\n\n# We're now ready to build the model.\n\nusing ClimateMachine.Ocean\n\nmodel = Ocean.HydrostaticBoussinesqSuperModel(\n domain = domain,\n time_step = 2.0,\n initial_conditions = initial_conditions,\n parameters = EarthParameters(),\n turbulence_closure = (νʰ = 0, κʰ = 0, νᶻ = 0, κᶻ = 0),\n coriolis = (f₀ = f, β = 0),\n boundary_tags = boundary_tags,\n boundary_conditions = boundary_conditions,\n);\nnothing\n\n# !!! info \"Horizontallly-periodic boundary conditions\"\n# To set horizontally-periodic boundary conditions with\n# `(solid_surface_boundary_conditions, free_surface_boundary_conditions)`\n# in the vertical direction use `periodicity = (true, true, false)` in\n# the `domain` constructor and `boundary_tags = ((0, 0), (0, 0), (1, 2))`\n# in the constructor for `HydrostaticBoussinesqSuperModel`.\n\n# # Animating the solution\n#\n# To animate the `ClimateMachine.Ocean` solution, we'll create a callback\n# that draws a plot and stores it in an array. When the simulation is finished,\n# we'll string together the plotted frames into an animation.\n\nusing Printf\nusing Plots\nusing ClimateMachine.GenericCallbacks: EveryXSimulationSteps\nusing ClimateMachine.CartesianFields: assemble\nusing ClimateMachine.Ocean: current_step, current_time\n\nu, v, η, θ = model.fields\n\n## Container to hold the plotted frames\nmovie_plots = []\n\nplot_every = 200 # iterations\n\nplot_maker = EveryXSimulationSteps(plot_every) do\n\n @info \"Steps: $(current_step(model)), time: $(current_time(model))\"\n\n assembled_u = assemble(u.elements)\n assembled_v = assemble(v.elements)\n assembled_η = assemble(η.elements)\n\n umax = 0.5 * max(maximum(abs, u), maximum(abs, v))\n ulim = (-umax, umax)\n\n u_plot = plot(\n assembled_u.x[:, 1, 1],\n [assembled_u.data[:, 1, 1] assembled_v.data[:, 1, 1]],\n xlim = domain.x,\n ylim = (-0.7U, 0.7U),\n label = [\"u\" \"v\"],\n linewidth = 2,\n xlabel = \"x (m)\",\n ylabel = \"Velocities (m s⁻¹)\",\n )\n\n η_plot = plot(\n assembled_η.x[:, 1, 1],\n assembled_η.data[:, 1, 1],\n xlim = domain.x,\n ylim = (-0.01a, 1.2a),\n linewidth = 2,\n label = nothing,\n xlabel = \"x (m)\",\n ylabel = \"η (m)\",\n )\n\n push!(movie_plots, (u = u_plot, η = η_plot, time = current_time(model)))\n\n return nothing\nend\n\n# # Running the simulation and animating the results\n#\n# Finally, we run the simulation,\n\nhours = 3600.0\n\nmodel.solver_configuration.timeend = 2hours\n\nresult = ClimateMachine.invoke!(\n model.solver_configuration;\n user_callbacks = [plot_maker],\n)\n\n# and animate the results,\n\nanimation = @animate for p in movie_plots\n title = @sprintf(\"Geostrophic adjustment at t = %.2f hours\", p.time / hours)\n frame = plot(\n p.u,\n p.η,\n layout = (2, 1),\n size = (800, 600),\n title = [title \"\"],\n )\nend\n\ngif(animation, \"geostrophic_adjustment.mp4\", fps = 5) # hide\n" }, { "path": "tutorials/Ocean/internal_wave.jl", "content": "# # Mode-1 internal wave reflection\n#\n# This example simulates the propagation of a mode-1 internal wave\n# using the `ClimateMachine.Ocean` subcomponent to solve the hydrostatic\n# Boussinesq equations.\n#\n# First we `ClimateMachine.init()`.\n\nusing ClimateMachine\n\nClimateMachine.init()\n\n# # Domain setup\n#\n# We formulate a non-dimension problem in a Cartesian domain with oceanic anisotropy,\n\nusing ClimateMachine.CartesianDomains\n\ndomain = RectangularDomain(\n Ne = (32, 1, 4),\n Np = 4,\n x = (-128, 128),\n y = (-128, 128),\n z = (-1, 0),\n periodicity = (false, false, false),\n)\n\n# # Parameters\n#\n# We choose parameters appropriate for a hydrostatic internal wave,\n\n# Non-dimensional internal wave parameters\nf = 1 # Coriolis\nN = 10 # Buoyancy frequency\n\n# Note that the validity of the hydrostatic approximation requires\n# small aspect ratio motions with ``k / m \\\\ll 1``.\n# The hydrostatic dispersion relation for inertia-gravity waves then implies that\n\nλ = 8 # horizontal wave-length\nk = 2π / λ # horizontal wavenumber\nm = π\n\nω² = f^2 + N^2 * k^2 / m^2 # and\n\nω = √(ω²)\n\n# # Internal wave initial condition\n# \n# We impose modest gravitational acceleration to render time-stepping feasible,\n\nusing CLIMAParameters: AbstractEarthParameterSet, Planet\nstruct NonDimensionalParameters <: AbstractEarthParameterSet end\n\nPlanet.grav(::NonDimensionalParameters) = 256.0\n\n# we'd like to use `θ` as a buoyancy variable, which requires\n# setting the thermal expansion coefficient ``αᵀ`` to\n\ng = Planet.grav(NonDimensionalParameters())\n\nαᵀ = 1 / g\n\n# We then use the \"polarization relations\" for vertically-standing, horizontally-\n# propagating hydrostatic internal waves to initialze two wave packets.\n# The hydrostatic polarization relations require\n#\n# ```math\n# \\begin{gather}\n# (∂_t^2 + f^2) u = - ∂_x ∂_t p\n# ∂_t v = - f u\n# b = ∂_z p\n# \\end{gather}\n# ```\n#\n# Thus given ``p = \\cos (k x - ω t) \\cos (m z)``, we find\n\nδ = domain.L.x / 15\na(x) = 1e-6 * exp(-x^2 / 2 * δ^2)\n\nũ(x, z, t) = +a(x) * ω * sin(k * x - ω * t) * cos(m * z)\nṽ(x, z, t) = -a(x) * f * cos(k * x - ω * t) * cos(m * z)\nθ̃(x, z, t) = -a(x) * m / k * (ω^2 - f^2) * sin(k * x - ω * t) * sin(m * z)\n\nuᵢ(x, y, z) = ũ(x, z, 0)\nvᵢ(x, y, z) = ṽ(x, z, 0)\nθᵢ(x, y, z) = θ̃(x, z, 0) + N^2 * z\n\nusing ClimateMachine.Ocean.OceanProblems: InitialConditions\n\ninitial_conditions = InitialConditions(u = uᵢ, v = vᵢ, θ = θᵢ)\n\n# # Model configuration\n# \n# We choose a time-step that resolves the gravity wave phase speed,\n\ntime_step = 0.005 # close to Δx / c = 0.5 * 1/16, where Δx is nominal resolution\n\n# and build a model with a smidgeon of viscosity and diffusion,\n\nusing ClimateMachine.Ocean: HydrostaticBoussinesqSuperModel\n\nmodel = HydrostaticBoussinesqSuperModel(\n domain = domain,\n time_step = time_step,\n initial_conditions = initial_conditions,\n parameters = NonDimensionalParameters(),\n turbulence_closure = (νʰ = 1e-6, νᶻ = 1e-6, κʰ = 1e-6, κᶻ = 1e-6),\n coriolis = (f₀ = f, β = 0),\n buoyancy = (αᵀ = αᵀ,),\n boundary_tags = ((1, 1), (1, 1), (1, 2)),\n);\nnothing\n\n# # Fetching data for an animation\n#\n# To animate the `ClimateMachine.Ocean` solution, we assemble and\n# cache the horizontal velocity ``u`` at periodic intervals:\n\nusing ClimateMachine.Ocean: current_time\nusing ClimateMachine.CartesianFields: assemble\nusing ClimateMachine.GenericCallbacks: EveryXSimulationTime\n\nfetched_states = []\nfetch_every = 0.2 * 2π / ω # time\n\ndata_fetcher = EveryXSimulationTime(fetch_every) do\n push!(\n fetched_states,\n (\n u = assemble(model.fields.u.elements),\n θ = assemble(model.fields.θ.elements),\n η = assemble(model.fields.η.elements),\n time = current_time(model),\n ),\n )\n return nothing\nend\n\n# We also build a callback to log the progress of our simulation,\n\nusing Printf\nusing ClimateMachine.GenericCallbacks: EveryXSimulationSteps\nusing ClimateMachine.Ocean: current_time, current_step, Δt\n\nprint_every = 100 # iterations\nwall_clock = [time_ns()]\n\ntiny_progress_printer = EveryXSimulationSteps(print_every) do\n\n @info(@sprintf(\n \"Steps: %d, time: %.2f, Δt: %.2f, max(|u|): %.2e, elapsed time: %.2f secs\",\n current_step(model),\n current_time(model),\n Δt(model),\n maximum(abs, model.fields.u),\n 1e-9 * (time_ns() - wall_clock[1])\n ))\n\n wall_clock[1] = time_ns()\nend\n\n# # Running the simulation and animating the results\n#\n# We're ready to launch.\n\nmodel.solver_configuration.timeend = 6 * 2π / ω\n## model.solver.dt = 0.05 # make this work\n\n@info \"\"\" Simulating a hydrostatic Gaussian wave packet with parameters\n\n f (Coriolis parameter): $f\n N (buoyancy frequency): $N\n Internal wave frequency: $(abs(ω))\n Surface wave frequency: $(k * sqrt(g * domain.L.z))\n Surface wave group velocity: $(sqrt(g * domain.L.z))\n Internal wave group velocity: $(N^2 * k / (ω * m))\n Domain width: $(domain.L.x) \n Domain height: $(domain.L.z) \n\n\"\"\"\n\nresult = ClimateMachine.invoke!(\n model.solver_configuration;\n user_callbacks = [tiny_progress_printer, data_fetcher],\n)\n\n# # Animating the result\n#\n# We first analye the results to generate plotting limits and contour levels\n\nηmax = maximum([maximum(abs, state.η.data) for state in fetched_states])\numax = maximum([maximum(abs, state.u.data) for state in fetched_states])\n\nηlim = (-ηmax, ηmax)\nulim = (-umax, umax)\nulevels = range(ulim[1], ulim[2], length = 31)\n\n# and then animate both fields in a loop,\n\nusing Plots\n\nanimation = @animate for (i, state) in enumerate(fetched_states)\n @info \"Plotting frame $i of $(length(fetched_states))...\"\n\n η_plot = plot(\n state.u.x[:, 1, 1],\n state.η.data[:, 1, 1],\n ylim = ηlim,\n label = nothing,\n title = @sprintf(\"η at t = %.2f\", state.time),\n )\n\n u_plot = contourf(\n state.u.x[:, 1, 1],\n state.u.z[1, 1, :],\n clamp.(state.u.data[:, 1, :], ulim[1], ulim[2])';\n aspectratio = 64,\n linewidth = 0,\n xlim = domain.x,\n ylim = domain.z,\n xlabel = \"x\",\n ylabel = \"z\",\n color = :balance,\n colorbar = false,\n clim = ulim,\n levels = ulevels,\n title = @sprintf(\"u at t = %.2f\", state.time),\n )\n\n plot(\n η_plot,\n u_plot,\n layout = Plots.grid(2, 1, heights = (0.3, 0.7)),\n link = :x,\n size = (600, 300),\n )\nend\n\ngif(animation, \"internal_wave.mp4\", fps = 5) # hide\n" }, { "path": "tutorials/Ocean/shear_instability.jl", "content": "# # Shear instability of a free-surface flow\n#\n# This script simulates the instability of a sheared, free-surface\n# flow using `ClimateMachine.Ocean.HydrostaticBoussinesqSuperModel`.\n\nusing Printf\nusing Plots\nusing ClimateMachine\n\nClimateMachine.init()\n\nClimateMachine.Settings.array_type = Array\n\nusing ClimateMachine.Ocean\nusing ClimateMachine.CartesianDomains\nusing ClimateMachine.CartesianFields\n\nusing ClimateMachine.GenericCallbacks: EveryXSimulationTime\nusing ClimateMachine.Ocean: current_step, Δt, current_time\nusing CLIMAParameters: AbstractEarthParameterSet, Planet\n\n# We begin by specifying the domain and mesh,\n\ndomain = RectangularDomain(\n Ne = (24, 24, 1),\n Np = 4,\n x = (-3π, 3π),\n y = (-3π, 3π),\n z = (0, 1),\n periodicity = (true, false, false),\n)\n\n# Note that the default solid-wall boundary conditions are free-slip and\n# insulating on tracers. Next, we specify model parameters and the sheared\n# initial conditions\n\nstruct NonDimensionalParameters <: AbstractEarthParameterSet end\nPlanet.grav(::NonDimensionalParameters) = 1\n\ninitial_conditions = InitialConditions(\n u = (x, y, z) -> tanh(y) + 0.1 * cos(x / 3) + 0.01 * randn(),\n v = (x, y, z) -> 0.1 * sin(y / 3),\n θ = (x, y, z) -> x,\n)\n\nmodel = Ocean.HydrostaticBoussinesqSuperModel(\n domain = domain,\n time_step = 0.05,\n initial_conditions = initial_conditions,\n parameters = NonDimensionalParameters(),\n turbulence_closure = (νʰ = 1e-2, κʰ = 1e-2, νᶻ = 1e-2, κᶻ = 1e-2),\n rusanov_wave_speeds = (cʰ = 0.1, cᶻ = 1),\n boundary_tags = ((0, 0), (1, 1), (1, 2)),\n boundary_conditions = (\n OceanBC(Impenetrable(FreeSlip()), Insulating()),\n OceanBC(Penetrable(FreeSlip()), Insulating()),\n ),\n);\nnothing\n\n# We prepare a callback that periodically fetches the horizontal velocity and\n# tracer concentration for later animation,\n\nu, v, η, θ = model.fields\nfetched_states = []\n\nstart_time = time_ns()\n\ndata_fetcher = EveryXSimulationTime(1) do\n step = @sprintf(\"Step: %d\", current_step(model))\n time = @sprintf(\"time: %.2f min\", current_time(model) / 60)\n max_u = @sprintf(\"max|u|: %.6f\", maximum(abs, u))\n\n elapsed = (time_ns() - start_time) * 1e-9\n wall_time = @sprintf(\"elapsed wall time: %.2f min\", elapsed / 60)\n\n @info \"$step, $time, $max_u, $wall_time\"\n\n push!(\n fetched_states,\n (u = assemble(u), θ = assemble(θ), time = current_time(model)),\n )\nend\n\n# and then run the simulation.\n\nmodel.solver_configuration.timeend = 100.0\n\nresult = ClimateMachine.invoke!(\n model.solver_configuration;\n user_callbacks = [data_fetcher],\n)\n\n# Finally, we make an animation of the evolving shear instability.\n\nanimation = @animate for (i, state) in enumerate(fetched_states)\n local u\n local θ\n\n @info \"Plotting frame $i of $(length(fetched_states))...\"\n\n kwargs =\n (xlim = domain.x, ylim = domain.y, linewidth = 0, aspectratio = 1)\n\n x, y = state.u.x[:, 1, 1], state.u.y[1, :, 1]\n\n u = state.u.data[:, :, 1]\n θ = state.θ.data[:, :, 1]\n\n ulim = 1\n θlim = 8\n\n ulevels = range(-ulim, ulim, length = 31)\n θlevels = range(-θlim, θlim, length = 31)\n\n u_plot = contourf(\n x,\n y,\n clamp.(u, -ulim, ulim)';\n levels = ulevels,\n color = :balance,\n kwargs...,\n )\n θ_plot = contourf(\n x,\n y,\n clamp.(θ, -θlim, θlim)';\n levels = θlevels,\n color = :thermal,\n kwargs...,\n )\n\n u_title = @sprintf(\"u at t = %.2f\", state.time)\n θ_title = @sprintf(\"θ at t = %.2f\", state.time)\n\n plot(u_plot, θ_plot, title = [u_title θ_title], size = (600, 250))\nend\n\ngif(animation, \"shear_instability.mp4\", fps = 5)\n" }, { "path": "tutorials/TutorialList.jl", "content": "# # Tutorials\n# A suite of concrete examples are provided here as a guidance for constructing experiments.\n# ## Balance Law\n# An introduction on components within a balance law is provided.\n# ## Atmos\n# Showcase drivers for atmospheric modelling in GCM, single stack, and LES simulations are provided.\n# - Dry Idealzed GCM: The Held-Suarez configuration is used as a guidance to create a driver that runs a simple GCM simulation.\n# - Single Element Stack: The Burgers Equations with a passive tracer is used as a guidance to run the simulation on a single element stack.\n# - LES Experiment: The dry rising bubble case is used as a quigance in creating an LES driver.\n# - Topography: Experiments of dry flow over prescirbe topography (Agnesi mountain) are provided for:\n# * Linear Hydrostatic Mountain\n# * Linear Non-Hydrostatic Mountain\n# ## Ocean\n# A showcase for Ocean model is still under construction.\n# ## Land\n# Examples are provided in constructing balance law and solving for fundemental equations in land modelling.\n# - Heat: A tutorial shows how to create a HeatModel to solve the heat equation and visualize the outputs.\n# - Soil: Examples of solving fundemental equations in the soil model are provided.\n# * Hydraulic Functions: a tutorial to specify the hydraulic function in the Richard's equation.\n# * Soil Heat Equations: a tutorial for solving the heat equation in the soil.\n# * Coupled Water and Heat: a tutorial for solving interactive heat and wateri in the soil model.\n# ## Numerics (need to be moved to How-to-Guide)\n# - System Solvers: Two numerical methods to solve the linear system Ax=b are provided.\n# * Conjugate Gradient\n# * Batched Generalized Minimal Residual\n# - DG Methods\n# * Filters\n# ## Diagnostics\n# A diagnostic tool that can\n# - generate statistics for MPIStateArrays\n# - validate with reference values\n# for debugging purposes.\n" }, { "path": "tutorials/literate_markdown.jl", "content": "# # How to generate a literate tutorial file\n\n# To create an tutorial using ClimateMachine, please use [Literate.jl](https://github.com/fredrikekre/Literate.jl), and consult the [Literate documentation](https://fredrikekre.github.io/Literate.jl/stable/) for questions.\n#\n# For now, all literate tutorials are held in the `tutorials` directory\n\n# With Literate, all comments turn into markdown text and any Julia code is read and run *as if it is in the Julia REPL*.\n# As a small caveat to this, you might need to suppress the output of certain commands.\n# For example, if you define and run the following function\n\nfunction f()\n return x = [i * i for i in 1:10]\nend\n\nx = f()\n\n# The entire list will be output, while\n\nf();\n\n# does not (because of the `;`).\n#\n# To show plots, you may do something like the following:\nusing Plots\nplot(x)\n\n# Please consider writing the comments in your tutorial as if they are meant to be read as an *article explaining the topic the tutorial is meant to explain.*\n# If there are any specific nuances to writing Literate documentation for ClimateMachine, please let us know!\n" } ]