diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 62ad092398a..39c70da0c39 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -246,6 +246,20 @@ steps: --job_id sphere_baroclinic_wave_rhoe_equilmoist --out_dir sphere_baroclinic_wave_rhoe_equilmoist artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist/*" + + - label: ":computer: no lim ARS baroclinic wave (ρe) equilmoist explicit vertdiff" + command: > + julia --color=yes --project=examples examples/hybrid/driver.jl + --config_file $CONFIG_PATH/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml + + julia --color=yes --project=examples post_processing/remap/remap_pipeline.jl + --data_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff + --out_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff + + julia --color=yes --project=examples post_processing/plot/plot_pipeline.jl + --nc_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff + --fig_dir sphere_baroclinic_wave_rhoe_equilmoist_expvdiff --case_name aquaplanet + artifact_paths: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff/*" - label: ":computer: SSP zalesak tracer & energy upwind baroclinic wave (ρe_tot) equilmoist" command: > @@ -1023,6 +1037,14 @@ steps: agents: slurm_mem: 20GB + - label: ":fire: Flame graph: perf target (diagnostics)" + command: > + julia --color=yes --project=perf perf/flame.jl + --config_file $PERF_CONFIG_PATH/flame/diagnostics.yml + artifact_paths: "flame_perf_diagnostics/*" + agents: + slurm_mem: 20GB + # Inference - label: ":rocket: JET n-failures (inference)" command: > diff --git a/.buildkite/scaling/pipeline.sh b/.buildkite/scaling/pipeline.sh index 28af1daa0e0..d579b6cf90f 100755 --- a/.buildkite/scaling/pipeline.sh +++ b/.buildkite/scaling/pipeline.sh @@ -65,32 +65,28 @@ for i in "${!resolutions[@]}"; do done # set up environment and agents -cat << EOM -env: - JULIA_VERSION: "1.9.3" - MPICH_VERSION: "4.0.0" - OPENMPI_VERSION: "4.1.1" - MPI_IMPL: "$mpi_impl" - CUDA_VERSION: "11.3" - OPENBLAS_NUM_THREADS: 1 - CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: "true" - +cat << 'EOM' agents: - config: cpu queue: central + modules: julia/1.9.3 cuda/11.8 ucx/1.14.1_cuda-11.8 openmpi/4.1.5_cuda-11.8 hdf5/1.12.2-ompi415 nsight-systems/2023.2.1 + +env: + JULIA_LOAD_PATH: "${JULIA_LOAD_PATH}:${BUILDKITE_BUILD_CHECKOUT_PATH}/.buildkite" + OPENBLAS_NUM_THREADS: 1 + JULIA_NVTX_CALLBACKS: gc + OMPI_MCA_opal_warn_on_missing_libcuda: 0 + JULIA_MAX_NUM_PRECOMPILE_FILES: 100 + JULIA_CPU_TARGET: 'broadwell;skylake' + SLURM_KILL_BAD_EXIT: 1 steps: - label: "init :computer:" key: "init_cpu_env" command: - - echo "--- Configure MPI" - - julia -e 'using Pkg; Pkg.add("MPIPreferences"); using MPIPreferences; use_system_binary()' - - echo "--- Instantiate" - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'" - "julia --project=examples -e 'using Pkg; Pkg.precompile()'" - "julia --project=examples -e 'using Pkg; Pkg.status()'" - agents: slurm_cpus_per_task: 8 env: @@ -129,11 +125,7 @@ if [[ "$profiling" == "enable" ]]; then else cpus_per_proc=1 fi -if [[ "$mpi_impl" == "mpich" ]]; then - launcher="srun --cpu-bind=cores" -else - launcher="mpiexec --map-by node:PE=$cpus_per_proc --bind-to core" -fi +launcher="srun --cpu-bind=cores" if [[ "$res" == "low" ]]; then time="04:00:00" @@ -160,7 +152,6 @@ cat << EOM - label: "$nprocs" key: "$job_id" command: - - "module load cuda/11.3 nsight-systems/2022.2.1" - "$launcher $command" - "find ${job_id} -iname '*.nsys-rep' -printf '%f\\\\n' | sort -V | jq --raw-input --slurp 'split(\"\n\") | .[0:-1] | {files: .} + {\"extension\": \"nsys-view\", \"version\": \"1.0\"}' > ${job_id}/${job_id}.nsys-view" - "find ${job_id} -iname '*.nsys-*' | sort -V | tar cvzf ${job_id}-nsys.tar.gz -T -" @@ -170,8 +161,6 @@ cat << EOM env: CLIMACORE_DISTRIBUTED: "MPI" agents: - config: cpu - queue: central slurm_time: $time EOM @@ -209,8 +198,6 @@ cat << EOM - "${res}-*.png" - "${res}-*.pdf" agents: - config: cpu - queue: central slurm_nodes: 1 slurm_tasks_per_node: 1 diff --git a/.gitignore b/.gitignore index f1c4fd2abd4..8f3dc90f697 100644 --- a/.gitignore +++ b/.gitignore @@ -22,6 +22,9 @@ deps/src/ docs/build/ docs/site/ +# File generated by make_available_diagnostics.jl +docs/src/available_diagnostics.md + # File generated by Pkg, the package manager, based on a corresponding Project.toml # It records a fixed state of all packages used by the project. As such, it should not be # committed for packages, but should be committed for applications that require a static diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml index 46ed8faadff..20ec70dc9b5 100644 --- a/config/default_configs/default_config.yml +++ b/config/default_configs/default_config.yml @@ -249,3 +249,6 @@ override_τ_precip: log_params: help: "Log parameters to file [`false` (default), `true`]" value: false +output_default_diagnostics: + help: "Output the default diagnostics associated to the selected atmospheric model" + value: false diff --git a/config/model_configs/edmfx_adv_test_box.yml b/config/model_configs/edmfx_adv_test_box.yml index 9a54b0105be..5ec3535702d 100644 --- a/config/model_configs/edmfx_adv_test_box.yml +++ b/config/model_configs/edmfx_adv_test_box.yml @@ -10,12 +10,14 @@ hyperdiff: "true" kappa_4: 1e8 x_max: 1e4 y_max: 1e4 -z_max: 3e4 +z_max: 5.5e4 x_elem: 2 y_elem: 2 -z_elem: 45 +z_elem: 63 dz_bottom: 30.0 +dz_top: 3000.0 dt: "10secs" t_end: "3600secs" dt_save_to_disk: "100secs" +FLOAT_TYPE: "Float64" toml: [toml/edmfx_box_advection.toml] diff --git a/config/model_configs/edmfx_trmm_box.yml b/config/model_configs/edmfx_trmm_box.yml index e9f1012fb32..62bbd4cfdf5 100644 --- a/config/model_configs/edmfx_trmm_box.yml +++ b/config/model_configs/edmfx_trmm_box.yml @@ -23,7 +23,7 @@ y_elem: 2 z_elem: 82 z_stretch: false dt: 1secs -t_end: 6hours +t_end: 2hours dt_save_to_disk: 10mins FLOAT_TYPE: "Float64" toml: [toml/edmfx_box.toml] diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe.yml b/config/model_configs/sphere_baroclinic_wave_rhoe.yml index 1a0176c8df7..8febe78d17e 100644 --- a/config/model_configs/sphere_baroclinic_wave_rhoe.yml +++ b/config/model_configs/sphere_baroclinic_wave_rhoe.yml @@ -1,6 +1,6 @@ dt_save_to_disk: "2days" regression_test: true initial_condition: "DryBaroclinicWave" -dt: "580secs" +dt: "400secs" t_end: "10days" job_id: "sphere_baroclinic_wave_rhoe" diff --git a/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml new file mode 100644 index 00000000000..45ca42e3b95 --- /dev/null +++ b/config/model_configs/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.yml @@ -0,0 +1,11 @@ +precip_model: "0M" +vert_diff: "true" +z_elem: 20 +dz_bottom: 100 +dt_save_to_disk: "12hours" +initial_condition: "MoistBaroclinicWave" +dt: "40secs" +t_end: "12hours" +job_id: "sphere_baroclinic_wave_rhoe_equilmoist_expvdiff" +moist: "equil" +toml: [toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml] diff --git a/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml b/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml index 466b9944dfb..d758db4bee2 100644 --- a/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml +++ b/config/model_configs/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.yml @@ -13,3 +13,4 @@ viscous_sponge: true job_id: "sphere_held_suarez_rhoe_equilmoist_hightop_sponge" moist: "equil" toml: [toml/sphere_held_suarez_rhoe_equilmoist_hightop_sponge.toml] +output_default_diagnostics: true diff --git a/config/model_configs/sphere_held_suarez_rhoe_hightop.yml b/config/model_configs/sphere_held_suarez_rhoe_hightop.yml index 45f060b6f2f..89e86b61cdb 100644 --- a/config/model_configs/sphere_held_suarez_rhoe_hightop.yml +++ b/config/model_configs/sphere_held_suarez_rhoe_hightop.yml @@ -3,7 +3,7 @@ dt_save_to_disk: "4days" regression_test: true t_end: "8days" forcing: "held_suarez" -dt: "500secs" +dt: "400secs" z_elem: 25 job_id: "sphere_held_suarez_rhoe_hightop" z_max: 45000.0 diff --git a/config/model_configs/sphere_held_suarez_rhotheta.yml b/config/model_configs/sphere_held_suarez_rhotheta.yml index e80b5c63cc5..2223e1ee8cd 100644 --- a/config/model_configs/sphere_held_suarez_rhotheta.yml +++ b/config/model_configs/sphere_held_suarez_rhotheta.yml @@ -3,5 +3,5 @@ dt_save_to_disk: "10days" regression_test: true t_end: "20days" forcing: "held_suarez" -dt: "500secs" +dt: "400secs" job_id: "sphere_held_suarez_rhotheta" diff --git a/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml b/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml index df7f22f94eb..aa077d17487 100644 --- a/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml +++ b/config/model_configs/sphere_zalesak_upwind_tracer_energy_ssp_baroclinic_wave_rhoe_equilmoist.yml @@ -1,7 +1,7 @@ dt_save_to_disk: "5days" initial_condition: "MoistBaroclinicWave" max_newton_iters_ode: 4 -dt: "500secs" +dt: "400secs" tracer_upwinding: zalesak t_end: "5days" ode_algo: "SSP333" diff --git a/config/perf_configs/flame/diagnostics.yml b/config/perf_configs/flame/diagnostics.yml new file mode 100644 index 00000000000..5a80e59bdfb --- /dev/null +++ b/config/perf_configs/flame/diagnostics.yml @@ -0,0 +1,11 @@ +job_id: "flame_perf_diagnostics" +diagnostics: + - short_name: ua + period: 1secs + reduction_time: average + - short_name: va + period: 1secs + reduction_time: max + - short_name: ta + period: 1secs + reduction_time: max diff --git a/config/perf_configs/gpu_baroclinic_wave_rhoe.yml b/config/perf_configs/gpu_baroclinic_wave_rhoe.yml index 90fdecc5e93..99349ca7bb8 100644 --- a/config/perf_configs/gpu_baroclinic_wave_rhoe.yml +++ b/config/perf_configs/gpu_baroclinic_wave_rhoe.yml @@ -1,5 +1,5 @@ job_id: "gpu_baroclinic_wave_rhoe" -dt: "580secs" +dt: "400secs" t_end: "10days" dt_save_to_disk: "2days" initial_condition: "DryBaroclinicWave" diff --git a/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml b/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml index 5eac11e6c55..b785b5d01fd 100644 --- a/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml +++ b/config/perf_configs/gpu_held_suarez_rhoe_hightop.yml @@ -3,6 +3,6 @@ z_elem: 25 dz_bottom: 300 forcing: "held_suarez" job_id: "gpu_held_suarez_rhoe_hightop" -dt: "500secs" +dt: "400secs" t_end: "8days" dt_save_to_disk: "4days" diff --git a/docs/make.jl b/docs/make.jl index 3f067559883..a2d06ebd901 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -8,6 +8,8 @@ disable_logging(Base.CoreLogging.Info) # Hide doctest's `@info` printing doctest(ClimaAtmos) disable_logging(Base.CoreLogging.BelowMinLevel) # Re-enable all logging +include("make_diagnostic_table.jl") + makedocs( CitationBibliography(joinpath(@__DIR__, "bibliography.bib")), modules = [ClimaAtmos], @@ -26,6 +28,8 @@ makedocs( "Contributor Guide" => "contributor_guide.md", "Equations" => "equations.md", "EDMF Equations" => "edmf_equations.md", + "Diagnostics" => "diagnostics.md", + "Available Diagnostics" => "available_diagnostics.md", "Diagnostic EDMF Equations" => "diagnostic_edmf_equations.md", "Gravity Wave Drag Parameterizations" => "gravity_wave.md", "Radiative Equilibrium" => "radiative_equilibrium.md", diff --git a/docs/make_diagnostic_table.jl b/docs/make_diagnostic_table.jl new file mode 100644 index 00000000000..ee6d007ae73 --- /dev/null +++ b/docs/make_diagnostic_table.jl @@ -0,0 +1,32 @@ +import ClimaAtmos as CA + +# Read all the diagnostics we know how to compute, and print them into a +# markdown table that is later compiled into the docs + +# basename(pwd()) if the code is run from inside the docs folder. If we don't +# have that, we will assume that we are running from the project root. If this +# code is run from anywhere but these two places, mkdocs will fail to find +# availbale_diagnostics.md +prefix = basename(pwd()) == "docs" ? "" : "docs/" + +out_path = "$(prefix)src/available_diagnostics.md" + +open(out_path, "w") do file + + write(file, "# Available diagnostic variables\n\n") + + write( + file, + "| Short name | Long name | Standard name | Units | Comments |\n", + ) + write(file, "|---|---|---|---|---|\n") + + for d in values(CA.Diagnostics.ALL_DIAGNOSTICS) + write(file, "| `$(d.short_name)` ") + write(file, "| $(d.long_name) ") + write(file, "| `$(d.standard_name)` ") + write(file, "| $(d.units) ") + write(file, "| $(d.comments)|\n") + end +end +@info "Written $out_path" diff --git a/docs/src/diagnostics.md b/docs/src/diagnostics.md new file mode 100644 index 00000000000..b55ad74b690 --- /dev/null +++ b/docs/src/diagnostics.md @@ -0,0 +1,316 @@ +# Computing and saving diagnostics + +## I want to compute and output a diagnostic variable + +### From a YAML file + +If you configure your simulation with YAML files, there are two options that are +important to know about. When `output_default_diagnostics` is set to `true`, the +default diagnostics for the given atmospheric model will be output. Note that +they might be incompatible with your simulation (e.g., you want to output +hourly maxima when the timestep is 4 hours). + +Second, you can specify the diagnostics you want to output directly in the +`diagnostics` section of your YAML file. For instance: +``` +diagnostics: + - short_name: rhoa + output_name: a_name + period: 3hours + writer: nc + - reduction_time: average + short_name: rhoa + period: 12hours + writer: h5 +``` +This adds two diagnostics (both for `rhoa`). The `period` keyword +identifies the period over which to compute the reduction and how often to save +to disk. `output_name` is optional, and if provided, it identifies the name of the +output file. + +The default `writer` is HDF5. If `writer` is `nc` or `netcdf`, the output is +remapped non-conservatively on a Cartesian grid and saved to a NetCDF file. +Currently, only 3D fields on cubed spheres are supported. + +### From a script + +The simplest way to get started with diagnostics is to use the defaults for your +atmospheric model. `ClimaAtmos` defines a function `default_diagnostic`. You +can execute this function on an `AtmosModel` or on any of its fields to obtain a +list of diagnostics ready to be passed to the simulation. So, for example + +```julia + +model = ClimaAtmos.AtmosModel(..., moisture_model = ClimaAtmos.DryModel(), ...) + +diagnostics = ClimaAtmos.default_diagnostics(model) +# => List of diagnostics that include the ones specified for the DryModel +``` + +Technically, the diagnostics are represented as `ScheduledDiagnostic` objects, +which contain information about what variable has to be computed, how often, +where to save it, and so on (read below for more information on this). You can +construct your own lists of `ScheduledDiagnostic`s starting from the variables +defined by `ClimaAtmos`. The `DiagnosticVariable`s in `ClimaAtmos` are +identified with by the short and unique name, so that you can access them +directly with the function `diagnostic_variable`. One way to do so is by +using the provided convenience functions for common operations, e.g., continuing +the previous example + +```julia + +push!(diagnostics, daily_max("air_density", "air_temperature")) +``` + +Now `diagnostics` will also contain the instructions to compute the daily +maximum of `air_density` and `air_temperature`. + +The diagnostics that are built-in `ClimaAtmos` are collected in [Available +diagnostic variables](@ref). + +If you are using `ClimaAtmos` with a script-based interface, you have access to +the complete flexibility in your diagnostics. Read the section about the +low-level interface to see how to implement custom diagnostics, reductions, or +writers. + +### The low-level interface + +Diagnostics are computed and output through callbacks to the main integrator. +`ClimaAtmos` produces the list of callbacks from a ordered list of +`ScheduledDiagnostic`s. These callbacks are orchestrated by a callback +`orchestrate_diagnostics` that runs at the end of every step and calls all the +diagnostic callbacks that are scheduled to be run at that step. + +A `ScheduledDiagnostic` is an instruction on how to compute and output a given +`DiagnosticVariable` (see below), along with specific choices regarding +reductions, compute/output frequencies, and so on. It can be point-wise in space +and time, or can be the result of a reduction in a period that is defined by +`output_every` (e.g., the daily average temperature). + +More specifically, a `ScheduledDiagnostic` contains the following pieces of data + +- `variable`: The diagnostic variable that has to be computed and output. +- `output_every`: Frequency of how often to save the results to disk. +- `output_writer`: Function that controls out to save the computed diagnostic + variable to disk. +- `reduction_time_func`: If not `nothing`, the `ScheduledDiagnostic` receives an + area of scratch space `acc` where to accumulate partial results. Then, at + every `compute_every`, `reduction_time_func` is computed between the + previously stored value in `acc` and the new value. This implements a running + reduction. For example, if `reduction_time_func = max`, the space `acc` will + hold the running maxima of the diagnostic. `acc` is reset after output. +- `reduction_space_func`: NOT IMPLEMENTED YET +- `compute_every`: Run the computations every `compute_every`. This is not + particularly useful for point-wise diagnostics, where we enforce that + `compute_every` = `output_every`. `compute_every` has to evenly divide + `output_every`. +- `pre_output_hook!`: Function that has to be run before saving to disk for + reductions (mostly used to implement averages). The function + `pre_output_hook!` is called with two arguments: the value accumulated during + the reduction, and the number of times the diagnostic was computed from the + last time it was output. +- `output_short_name`: A descriptive name that can be used by the + `output_writer`. If not provided, a default one is generated. It has to be + unique. +- `output_long_name`: A descriptive name that can be used by the `output_writer` + as attribute. If not provided, a default one is generated. + +To implement operations like the arithmetic average, the `reduction_time_func` +has to be chosen as `+`, and a `pre_output_hook!` that renormalize `acc` by the +number of samples has to be provided. `pre_output_hook!` should mutate the +accumulator in place. The return value of `pre_output_hook!` is discarded. An +example of `pre_output_hook!` to compute the arithmetic average is +`pre_output_hook!(acc, N) = @. acc = acc / N`. `ClimaAtmos` provides an alias to +the function needed to compute averages `ClimaAtmos.average_pre_output_hook!`. + +For custom reductions, it is necessary to also specify the identity of operation +by defining a new method to `identity_of_reduction`. `identity_of_reduction` is +a function that takes a `op` argument, where `op` is the operation for which we +want to define the identity. For instance, for the `max`, +`identity_of_reduction` would be `identity_of_reduction(::typeof{max}) = -Inf`. +The identities known to `ClimaAtmos` are defined in the +`diagnostics/reduction_identities.jl` file. The identity is needed to ensure +that we have a neutral state for the accumulators that are used in the +reductions. + +A key entry in a `ScheduledDiagnostic` is the `output_writer`, the function +responsible for saving the output to disk. `output_writer` is called with three +arguments: the value that has to be output, the `ScheduledDiagnostic`, and the +integrator. Internally, the integrator contains extra information (such as the +current timestep), so the `output_writer` can implement arbitrarily complex +behaviors. It is responsibility of the `output_writer` to properly use the +provided information for meaningful output. `ClimaAtmos` provides functions that +return `output_writer`s for standard operations. The main one is currently +`HDF5Writer`, which should be enough for most use cases. To use it, just +initialize a `ClimaAtmos.HDF5Writer` object with your choice of configuration +and pass it as a `output_writer` argument to the `ScheduledDiagnostic`. More +information about the options supported by `ClimaAtmos.HDF5Writer` is available +in its constructor. + +There are two flavors of `ScheduledDiagnostic`s: +`ScheduledDiagnosticIterations`, and `ScheduledDiagnosticTime`. The main +difference between the two is the domain of definition of their frequencies, +which is measured in timesteps for the first one, and in seconds for the other +one. `ScheduledDiagnosticTime`s offer a more natural way to set up physically +meaningful reductions (e.g., we want a daily average). However, it is not clear +what to do when the period does not line up with the timestep. What is the +expected behavior when we want a daily average but our timestep is of 10 hours? +There are multiple possible answer to this question. In `ClimaAtmos`, we enforce +that all the periods are multiples of the timestep. With this restriction, a +`ScheduledDiagnosticTime` can be translated to a +`ScheduledDiagnosticIterations`, where the problem is perfectly represented (in +this sense, one can think of `ScheduledDiagnosticIterations` as being as +internal representation and as `ScheduledDiagnosticTime` being the user-facing +one). + +`ScheduledDiagnosticTime` behave like as `ScheduledDiagnosticIterations`, with +the exception that they can take a special value for `compute_every`, which can +be set to `:timestep` (the symbol) to ensure that the diagnostic is computed at +the end of every integration step. This is particularly convenient when defining +default diagnostics, as they should be largely independent on the choice of the +specific simulation being run (and its timestep). + +Given a timestep `dt`, a `ScheduledDiagnosticIterations` can be obtained from a +`ScheduledDiagnosticTime` `sd` simply by calling +``ScheduledDiagnosticIterations(sd, dt)`. + +## I want to add a new diagnostic variable + +Diagnostic variables are represented in `ClimaAtmos` with a `DiagnosticVariable` +`struct`. Fundamentally, a `DiagnosticVariable` contains metadata about the +variable, and a function that computes it from the state. + +### Metadata + +The metadata we currently support is `short_name`, `long_name`, `units`, +`comments`. This metadata is relevant mainly in the context of how the variable +is output. Therefore, it is responsibility of the `output_writer` (see +`ScheduledDiagnostic`) to handle the metadata properly. The `output_writer`s +provided by `ClimaAtmos` use this metadata. + +In `ClimaAtmos`, we follow the convention that: + +- `short_name` is the name used to identify the variable in the output files and + in the file names. It is short, but descriptive. We identify + diagnostics by their short name, so the diagnostics defined by + `ClimaAtmos` have to have unique `short_name`s. + +- `long_name`: Name used to describe the variable in the output file as attribute. + +- `standard_name`: Standard name, as in + [CF + conventions](http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html) + +- `units`: Physical units of the variable. + +- `comments`: More verbose explanation of what the variable is, or comments related to how + it is defined or computed. + +In `ClimaAtmos`, we follow the [CMIP6 MIP table](https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb) +for short names and long names where available. Standard names in the table are not used. + +### Compute function + +The other piece of information needed to specify a `DiagnosticVariable` is a +function `compute!`. Schematically, a `compute!` has to look like +```julia +function compute!(out, state, cache, time) + if isnothing(out) + return ... # Calculations with the state and the cache + else + out .= ... # Calculations with the state and the cache + end +end +``` +The first time `compute!` is called, the function has to allocate memory and +return its output. All the subsequent times, `out` will be the pre-allocated +area of memory, so the function has to write the new value in place. + +Diagnostics are implemented as callbacks functions which pass the `state`, +`cache`, and `time` from the integrator to `compute!`. + +`compute!` also takes a second argument, `out`, which is used to +avoid extra memory allocations (which hurt performance). If `out` is `nothing`, +and new area of memory is allocated. If `out` is a `ClimaCore.Field`, the +operation is done in-place without additional memory allocations. + +If your diagnostic depends on the details of the model, we recommend using +additional functions so that the correct one can be found through dispatching. +For instance, if you want to compute relative humidity, which does not make +sense for dry simulations, you should define the functions + +```julia +function compute_relative_humidity!( + out, + state, + cache, + time, + moisture_model::T, +) where {T} + error("Cannot compute relative_humidity with moisture_model = $T") +end + +function compute_relative_humidity!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + if isnothing(out) + return TD.relative_humidity.(thermo_params, cache.ᶜts) + else + out .= TD.relative_humidity.(thermo_params, cache.ᶜts) + end +end + +compute_relative_humidity!(out, state, cache, time) = + compute_relative_humidity!( + out, + state, + cache, + time, + cache.atmos.moisture_model, + ) +``` + +This will return the correct relative humidity and throw informative errors when +it cannot be computed. We could specialize +`compute_relative_humidity` further if the relative humidity +were computed differently for `EquilMoistModel` and `NonEquilMoistModel`. + +In `ClimaAtmos`, we define some helper functions to produce error messages, so +the above code can be written as +```julia +compute_relative_humidity!(out, state, cache, time) = + compute_relative_humidity!(out, state, cache, time, cache.atmos.moisture_model) +compute_relative_humidity!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("relative_humidity", model) + +function compute_relative_humidity!( + out, + state, + cache, + time, + ::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.relative_humidity.(thermo_params, cache.ᶜts) + else + out .= TD.relative_humidity.(thermo_params, cache.ᶜts) + end +end +``` + +### The `ClimaAtmos` `DiagnosticVariable`s + +`ClimaAtmos` comes with a collection of pre-defined `DiagnosticVariable`, index +with their `short_name`s. If you are extending `ClimaAtmos` and want to add a +new diagnostic variable, go ahead and look at the files we `include` in +`diagnostics/Diagnostics.jl`. You can add more diagnostics in those files or add +a new one. We provide a convenience function, `add_diagnostic_variable!` to add +new `DiagnosticVariable`s. `add_diagnostic_variable!` take the same arguments as +the constructor for `DiagnosticVariable`, but also performs additional checks. +Similarly, if you want to retrieve a diagnostic from `ALL_DIAGNOSTICS`, use the +`get_diagnostic_variable`function. diff --git a/perf/flame.jl b/perf/flame.jl index a2490677a3a..80ab8d826be 100644 --- a/perf/flame.jl +++ b/perf/flame.jl @@ -54,13 +54,14 @@ allocs = @allocated SciMLBase.step!(integrator) @info "`allocs ($job_id)`: $(allocs)" allocs_limit = Dict() -allocs_limit["flame_perf_target"] = 7968 -allocs_limit["flame_perf_target_tracers"] = 207600 -allocs_limit["flame_perf_target_edmfx"] = 274832 -allocs_limit["flame_perf_target_diagnostic_edmfx"] = 748176 -allocs_limit["flame_perf_target_edmf"] = 12031972048 +allocs_limit["flame_perf_target"] = 12864 +allocs_limit["flame_perf_target_tracers"] = 212496 +allocs_limit["flame_perf_target_edmfx"] = 304064 +allocs_limit["flame_perf_diagnostics"] = 3024344 +allocs_limit["flame_perf_target_diagnostic_edmfx"] = 762784 +allocs_limit["flame_perf_target_edmf"] = 12459299664 allocs_limit["flame_perf_target_threaded"] = 6175664 -allocs_limit["flame_perf_target_callbacks"] = 45111592 +allocs_limit["flame_perf_target_callbacks"] = 46413904 allocs_limit["flame_perf_gw"] = 4911463328 if allocs < allocs_limit[job_id] * buffer diff --git a/post_processing/contours_and_profiles.jl b/post_processing/contours_and_profiles.jl index 080c6c59f3f..7b01ddec4a8 100644 --- a/post_processing/contours_and_profiles.jl +++ b/post_processing/contours_and_profiles.jl @@ -429,6 +429,7 @@ function contours_and_profiles(output_dir, ref_job_id = nothing) (all_categories, :w_velocity), (all_categories, :buoyancy), (all_categories, :specific_enthalpy), + (all_categories, :density), ) if has_moisture profile_variables = ( diff --git a/post_processing/plot_scaling_results.jl b/post_processing/plot_scaling_results.jl index d330e5854b6..0a979f12a0d 100644 --- a/post_processing/plot_scaling_results.jl +++ b/post_processing/plot_scaling_results.jl @@ -89,7 +89,8 @@ ax1 = Axis( xgridvisible = true, ygridvisible = false, ) -scatterlines!(ax1, nprocs_clima_atmos, sypd_clima_atmos) +scatterlines!(nprocs_clima_atmos, sypd_clima_atmos) +# Plot a second axis to display tick labels clearly ax1 = Axis( fig[1, 1], yaxisposition = :right, @@ -103,15 +104,7 @@ ax1 = Axis( yscale = log10, ytickformat = "{:.2f}", ) -hlines!( - ax1, - sypd_clima_atmos, - xscale = log10, - yscale = log10, - color = :gray, - alpha = 0.5, - linestyle = :dash, -) +scatterlines!(nprocs_clima_atmos, sypd_clima_atmos) ax2 = Axis( fig[2, 1], diff --git a/regression_tests/ref_counter.jl b/regression_tests/ref_counter.jl index a949a93dfcc..b0d73241cad 100644 --- a/regression_tests/ref_counter.jl +++ b/regression_tests/ref_counter.jl @@ -1 +1 @@ -128 +129 diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl index 07f6847372a..1f0c7fcac38 100644 --- a/src/ClimaAtmos.jl +++ b/src/ClimaAtmos.jl @@ -110,6 +110,9 @@ include(joinpath("prognostic_equations", "limited_tendencies.jl")) include(joinpath("callbacks", "callbacks.jl")) +include(joinpath("diagnostics", "Diagnostics.jl")) +import .Diagnostics as CAD + include(joinpath("solver", "model_getters.jl")) # high-level (using parsed_args) model getters include(joinpath("solver", "type_getters.jl")) include(joinpath("solver", "yaml_helper.jl")) diff --git a/src/cache/diagnostic_edmf_precomputed_quantities.jl b/src/cache/diagnostic_edmf_precomputed_quantities.jl index a496a80bc0f..80236a82b43 100644 --- a/src/cache/diagnostic_edmf_precomputed_quantities.jl +++ b/src/cache/diagnostic_edmf_precomputed_quantities.jl @@ -681,7 +681,7 @@ function set_diagnostic_edmf_precomputed_quantities!(Y, p, t) ᶠu⁰ = p.ᶠtemp_C123 @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰) ᶜstrain_rate = p.ᶜtemp_UVWxUVW - compute_strain_rate!(ᶜstrain_rate, ᶠu⁰) + compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰) @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate) ᶜprandtl_nvec = p.ᶜtemp_scalar diff --git a/src/cache/edmf_precomputed_quantities.jl b/src/cache/edmf_precomputed_quantities.jl index b4a517089a4..2276fff43bb 100644 --- a/src/cache/edmf_precomputed_quantities.jl +++ b/src/cache/edmf_precomputed_quantities.jl @@ -230,7 +230,7 @@ function set_edmf_precomputed_quantities!(Y, p, ᶠuₕ³, t) ᶠu⁰ = p.ᶠtemp_C123 @. ᶠu⁰ = C123(ᶠinterp(Y.c.uₕ)) + C123(ᶠu³⁰) ᶜstrain_rate = p.ᶜtemp_UVWxUVW - compute_strain_rate!(ᶜstrain_rate, ᶠu⁰) + compute_strain_rate_center!(ᶜstrain_rate, ᶠu⁰) @. ᶜstrain_rate_norm = norm_sqr(ᶜstrain_rate) ᶜprandtl_nvec = p.ᶜtemp_scalar diff --git a/src/cache/temporary_quantities.jl b/src/cache/temporary_quantities.jl index c666866b903..e43a05acda0 100644 --- a/src/cache/temporary_quantities.jl +++ b/src/cache/temporary_quantities.jl @@ -54,6 +54,10 @@ function temporary_quantities(atmos, center_space, face_space) typeof(UVW(FT(0), FT(0), FT(0)) * UVW(FT(0), FT(0), FT(0))'), center_space, ), # ᶜstrain_rate + ᶠtemp_UVWxUVW = Fields.Field( + typeof(UVW(FT(0), FT(0), FT(0)) * UVW(FT(0), FT(0), FT(0))'), + face_space, + ), # ᶠstrain_rate # TODO: Remove this hack sfc_temp_C3 = Fields.Field(C3{FT}, Spaces.level(face_space, half)), # ρ_flux_χ ) diff --git a/src/callbacks/callback_helpers.jl b/src/callbacks/callback_helpers.jl index 3c028e91ea1..7430bf95508 100644 --- a/src/callbacks/callback_helpers.jl +++ b/src/callbacks/callback_helpers.jl @@ -97,3 +97,6 @@ function n_steps_per_cycle_per_cb(cbs::SciMLBase.CallbackSet, dt) end end end + +n_steps_per_cycle_per_cb_diagnostic(cbs) = + [callback_frequency(cb).n for cb in cbs] diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl index f07039909a1..ff13a84e647 100644 --- a/src/callbacks/callbacks.jl +++ b/src/callbacks/callbacks.jl @@ -219,6 +219,7 @@ function common_diagnostics(p, ᶜu, ᶜts) potential_temperature = TD.dry_pottemp.(thermo_params, ᶜts), specific_enthalpy = TD.specific_enthalpy.(thermo_params, ᶜts), buoyancy = CAP.grav(p.params) .* (p.ᶜρ_ref .- ᶜρ) ./ ᶜρ, + density = TD.air_density.(thermo_params, ᶜts), ) if !(p.atmos.moisture_model isa DryModel) diagnostics = (; diff --git a/src/diagnostics/Diagnostics.jl b/src/diagnostics/Diagnostics.jl new file mode 100644 index 00000000000..17f3d4ff5f7 --- /dev/null +++ b/src/diagnostics/Diagnostics.jl @@ -0,0 +1,40 @@ +module Diagnostics + +import LinearAlgebra: dot + +import ClimaCore: Fields, Geometry, InputOutput, Meshes, Spaces, Operators +import ClimaCore.Utilities: half +import Thermodynamics as TD + +import ..AtmosModel +import ..AtmosCallback +import ..EveryNSteps + +import ..Parameters as CAP + +import ..unit_basis_vector_data + +# moisture_model +import ..DryModel +import ..EquilMoistModel +import ..NonEquilMoistModel + +# energy_form +import ..TotalEnergy + +# precip_model +import ..Microphysics0Moment + +# radiation +import ClimaAtmos.RRTMGPInterface as RRTMGPI + +# turbconv_model +import ..EDMFX +import ..DiagnosticEDMFX + +# We need the abbreviations for symbols like curl, grad, and so on +include(joinpath("..", "utils", "abbreviations.jl")) + +include("diagnostic.jl") +include("writers.jl") +end diff --git a/src/diagnostics/core_diagnostics.jl b/src/diagnostics/core_diagnostics.jl new file mode 100644 index 00000000000..45d070e732e --- /dev/null +++ b/src/diagnostics/core_diagnostics.jl @@ -0,0 +1,474 @@ +# This file is included in Diagnostics.jl +# +# README: Adding a new core diagnostic: +# +# In addition to the metadata (names, comments, ...), the most important step in adding a +# new DiagnosticVariable is defining its compute! function. `compute!` has to take four +# arguments: (out, state, cache, time), and as to write the diagnostic in place into the +# `out` variable. +# +# Often, it is possible to compute certain diagnostics only for specific models (e.g., +# humidity for moist models). For that, it is convenient to adopt the following pattern: +# +# 1. Define a catch base function that does the computation we want to do for the case we know +# how to handle, for example +# +# function compute_hur!( +# out, +# state, +# cache, +# time, +# ::T, +# ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} +# thermo_params = CAP.thermodynamics_params(cache.params) +# out .= TD.relative_humidity.(thermo_params, cache.ᶜts) +# end +# +# 2. Define a function that has the correct signature and calls this function +# +# compute_hur!(out, state, cache, time) = +# compute_hur!(out, state, cache, time, cache.atmos.moisture_model) +# +# 3. Define a function that returns an error when the model is incorrect +# +# compute_hur!(_, _, _, _, model::T) where {T} = +# error_diagnostic_variable("relative_humidity", model) +# +# We can also output a specific error message +# +# compute_hur!(_, _, _, _, model::T) where {T} = +# error_diagnostic_variable("relative humidity makes sense only for moist models") + +# General helper functions for undefined diagnostics for a particular model +error_diagnostic_variable( + message = "Cannot compute $variable with model = $T", +) = error(message) +error_diagnostic_variable(variable, model::T) where {T} = + error_diagnostic_variable("Cannot compute $variable with model = $T") + +### +# Rho (3d) +### +add_diagnostic_variable!( + short_name = "rhoa", + long_name = "Air Density", + standard_name = "air_density", + units = "kg m^-3", + comments = "Density of air", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return copy(state.c.ρ) + else + out .= state.c.ρ + end + end, +) + +### +# U velocity (3d) +### +add_diagnostic_variable!( + short_name = "ua", + long_name = "Eastward Wind", + standard_name = "eastward_wind", + units = "m s^-1", + comments = "Eastward (zonal) wind component", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return copy(Geometry.UVector.(cache.ᶜu).components.data.:1) + else + out .= Geometry.UVector.(cache.ᶜu).components.data.:1 + end + end, +) + +### +# V velocity (3d) +### +add_diagnostic_variable!( + short_name = "va", + long_name = "Northward Wind", + standard_name = "northward_wind", + units = "m s^-1", + comments = "Northward (meridional) wind component", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return copy(Geometry.VVector.(cache.ᶜu).components.data.:1) + else + out .= Geometry.VVector.(cache.ᶜu).components.data.:1 + end + end, +) + +### +# W velocity (3d) +### +# TODO: may want to convert to omega (Lagrangian pressure tendency) as standard output, +# but this is probably more useful for now +# +add_diagnostic_variable!( + short_name = "wa", + long_name = "Upward Air Velocity", + standard_name = "upward_air_velocity", + units = "m s^-1", + comments = "Vertical wind component", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return copy(Geometry.WVector.(cache.ᶜu).components.data.:1) + else + out .= Geometry.WVector.(cache.ᶜu).components.data.:1 + end + end, +) + +### +# Temperature (3d) +### +add_diagnostic_variable!( + short_name = "ta", + long_name = "Air Temperature", + standard_name = "air_temperature", + units = "K", + comments = "Temperature of air", + compute! = (out, state, cache, time) -> begin + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_temperature.(thermo_params, cache.ᶜts) + else + out .= TD.air_temperature.(thermo_params, cache.ᶜts) + end + end, +) + +### +# Potential temperature (3d) +### +add_diagnostic_variable!( + short_name = "thetaa", + long_name = "Air Potential Temperature", + standard_name = "air_potential_temperature", + units = "K", + comments = "Potential temperature of air", + compute! = (out, state, cache, time) -> begin + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.dry_pottemp.(thermo_params, cache.ᶜts) + else + out .= TD.dry_pottemp.(thermo_params, cache.ᶜts) + end + end, +) + +### +# Air pressure (3d) +### +add_diagnostic_variable!( + short_name = "pfull", + long_name = "Pressure at Model Full-Levels", + units = "Pa", + comments = "Pressure of air", + compute! = (out, state, cache, time) -> begin + if isnothing(out) + return copy(cache.ᶜp) + else + out .= cache.ᶜp + end + end, +) + +### +# Vorticity (3d) +### +add_diagnostic_variable!( + short_name = "rv", + long_name = "Relative Vorticity", + standard_name = "relative_vorticity", + units = "s^-1", + comments = "Vertical component of relative vorticity", + compute! = (out, state, cache, time) -> begin + vort = @. Geometry.WVector(curlₕ(state.c.uₕ)).components.data.:1 + # We need to ensure smoothness, so we call DSS + Spaces.weighted_dss!(vort) + if isnothing(out) + return copy(vort) + else + out .= vort + end + end, +) + + +### +# Relative humidity (3d) +### +compute_hur!(out, state, cache, time) = + compute_hur!(out, state, cache, time, cache.atmos.moisture_model) +compute_hur!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("hur", model) + +function compute_hur!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.relative_humidity.(thermo_params, cache.ᶜts) + else + out .= TD.relative_humidity.(thermo_params, cache.ᶜts) + end +end + +add_diagnostic_variable!( + short_name = "hur", + long_name = "Relative Humidity", + standard_name = "relative_humidity", + units = "", + comments = "Total amount of water vapor in the air relative to the amount achievable by saturation at the current temperature", + compute! = compute_hur!, +) + +### +# Total specific humidity (3d) +### +compute_hus!(out, state, cache, time) = + compute_hus!(out, state, cache, time, cache.atmos.moisture_model) +compute_hus!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("hus", model) + +function compute_hus!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + if isnothing(out) + return state.c.ρq_tot ./ state.c.ρ + else + out .= state.c.ρq_tot ./ state.c.ρ + end +end + +add_diagnostic_variable!( + short_name = "hus", + long_name = "Specific Humidity", + standard_name = "specific_humidity", + units = "kg kg^-1", + comments = "Mass of all water phases per mass of air", + compute! = compute_hus!, +) + +### +# Surface specific humidity (2d) +### +compute_hussfc!(out, state, cache, time) = + compute_hussfc!(out, state, cache, time, cache.atmos.moisture_model) +compute_hussfc!(_, _, _, _, model::T) where {T} = + error_diagnostic_variable("hussfc", model) + +function compute_hussfc!( + out, + state, + cache, + time, + moisture_model::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.total_specific_humidity.( + thermo_params, + cache.sfc_conditions.ts, + ) + else + out .= + TD.total_specific_humidity.(thermo_params, cache.sfc_conditions.ts) + end +end + +add_diagnostic_variable!( + short_name = "hussfc", + long_name = "Surface Specific Humidity", + standard_name = "specific_humidity", + units = "kg kg^-1", + comments = "Mass of all water phases per mass of air in the layer infinitely close to the surface", + compute! = compute_hussfc!, +) + +### +# Surface temperature (2d) +### +add_diagnostic_variable!( + short_name = "ts", + long_name = "Surface Temperature", + standard_name = "surface_temperature", + units = "K", + comments = "Temperature of the lower boundary of the atmosphere", + compute! = (out, state, cache, time) -> begin + thermo_params = CAP.thermodynamics_params(cache.params) + if isnothing(out) + return TD.air_temperature.(thermo_params, cache.sfc_conditions.ts) + else + out .= TD.air_temperature.(thermo_params, cache.sfc_conditions.ts) + end + end, +) + +### +# Eastward and northward surface drag component (2d) +### +compute_tau!(_, _, _, _, energy_form::T) where {T} = + error_diagnostic_variable("tau", energy_form) + +function compute_tau!(out, state, cache, component, energy_form::TotalEnergy) + sfc_local_geometry = + Fields.level(Fields.local_geometry_field(state.f), Fields.half) + surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry)) + (; ρ_flux_uₕ) = cache.sfc_conditions + + if isnothing(out) + return getproperty( + Geometry.UVVector.( + adjoint.(ρ_flux_uₕ) .* surface_ct3_unit + ).components.data, + component, + ) + else + out .= getproperty( + Geometry.UVVector.( + adjoint.(ρ_flux_uₕ) .* surface_ct3_unit + ).components.data, + component, + ) + end + + return +end + +compute_tauu!(out, state, cache, time) = + compute_tau!(out, state, cache, :1, cache.atmos.energy_form) +compute_tauv!(out, state, cache, time) = + compute_tau!(out, state, cache, :2, cache.atmos.energy_form) + +add_diagnostic_variable!( + short_name = "tauu", + long_name = "Surface Downward Eastward Wind Stress", + standard_name = "downward_eastward_stress", + units = "Pa", + comments = "Eastward component of the surface drag", + compute! = compute_tauu!, +) + +add_diagnostic_variable!( + short_name = "tauv", + long_name = "Surface Downward Northward Wind Stress", + standard_name = "downward_northward_stress", + units = "Pa", + comments = "Northward component of the surface drag", + compute! = compute_tauv!, +) + +### +# Surface energy flux (2d) - TODO: this may need to be split into sensible and latent heat fluxes +### +compute_hfes!(out, state, cache, time) = + compute_hfes!(out, state, cache, time, cache.atmos.energy_form) +compute_hfes!(_, _, _, _, energy_form::T) where {T} = + error_diagnostic_variable("hfes", energy_form) + +function compute_hfes!(out, state, cache, time, energy_form::TotalEnergy) + (; ρ_flux_h_tot) = cache.sfc_conditions + sfc_local_geometry = + Fields.level(Fields.local_geometry_field(state.f), Fields.half) + surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry)) + if isnothing(out) + return dot.(ρ_flux_h_tot, surface_ct3_unit) + else + out .= dot.(ρ_flux_h_tot, surface_ct3_unit) + end +end + +add_diagnostic_variable!( + short_name = "hfes", + long_name = "Surface Upward Energy Flux", + units = "W m^-2", + comments = "Energy flux at the surface", + compute! = compute_hfes!, +) + +### +# Surface evaporation (2d) +### +compute_evspsbl!(out, state, cache, time) = compute_evspsbl!( + out, + state, + cache, + time, + cache.atmos.moisture_model, + cache.atmos.energy_form, +) +compute_evspsbl!( + _, + _, + _, + _, + moisture_model::T1, + energy_form::T2, +) where {T1, T2} = error_diagnostic_variable( + "Can only compute surface_evaporation with energy_form = TotalEnergy() and with a moist model", +) + +function compute_evspsbl!( + out, + state, + cache, + time, + moisture_model::T, + energy_form::TotalEnergy, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + (; ρ_flux_q_tot) = cache.sfc_conditions + sfc_local_geometry = + Fields.level(Fields.local_geometry_field(state.f), Fields.half) + surface_ct3_unit = CT3.(unit_basis_vector_data.(CT3, sfc_local_geometry)) + + if isnothing(out) + return dot.(ρ_flux_q_tot, surface_ct3_unit) + else + out .= dot.(ρ_flux_q_tot, surface_ct3_unit) + end +end + +add_diagnostic_variable!( + short_name = "evspsbl", + long_name = "Evaporation Including Sublimation and Transpiration", + units = "kg m^-2 s^-1", + comments = "evaporation at the surface", + compute! = compute_evspsbl!, +) + +### +# Precipitation (2d) - TODO: change to kg m^-2 s^-1 +### +compute_pr!(out, state, cache, time) = + compute_pr!(out, state, cache, time, cache.atmos.precip_model) +compute_pr!(_, _, _, _, precip_model::T) where {T} = + error_diagnostic_variable("pr", precip_model) + +function compute_pr!(out, state, cache, time, precip_model::Microphysics0Moment) + if isnothing(out) + return cache.col_integrated_rain .+ cache.col_integrated_snow + else + out .= cache.col_integrated_rain .+ cache.col_integrated_snow + end +end + +add_diagnostic_variable!( + short_name = "pr", + long_name = "Precipitation", + standard_name = "precipitation", + units = "m s^-1", + comments = "Total precipitation including rain and snow", + compute! = compute_pr!, +) diff --git a/src/diagnostics/default_diagnostics.jl b/src/diagnostics/default_diagnostics.jl new file mode 100644 index 00000000000..8f2f6f54f31 --- /dev/null +++ b/src/diagnostics/default_diagnostics.jl @@ -0,0 +1,248 @@ +# This file is included by Diagnostics.jl and defines all the defaults for various models. A +# model here is either a global AtmosModel, or small (sub)models (e.g., DryModel()). +# +# If you are developing new models, add your defaults here. If you want to add more high +# level interfaces, add them here. Feel free to include extra files. + +""" + default_diagnostics(model) + + +Return a list of `ScheduledDiagnostic`s associated with the given `model`. + +""" +function default_diagnostics(model::AtmosModel) + # Unfortunately, [] is not treated nicely in a map (we would like it to be "excluded"), + # so we need to manually filter out the submodels that don't have defaults associated + # to + non_empty_fields = filter( + x -> default_diagnostics(getfield(model, x)) != [], + fieldnames(AtmosModel), + ) + + # We use a map because we want to ensure that diagnostics is a well defined type, not + # Any. This reduces latency. + return vcat( + core_default_diagnostics(), + map(non_empty_fields) do field + default_diagnostics(getfield(model, field)) + end..., + ) +end + +# Base case: if we call default_diagnostics on something that we don't have information +# about, we get nothing back (to be specific, we get an empty list, so that we can assume +# that all the default_diagnostics return the same type). This is used by +# default_diagnostics(model::AtmosModel), so that we can ignore defaults for submodels +# that have no given defaults. +default_diagnostics(_) = [] + + +""" + produce_common_diagnostic_function(period, reduction) + + +Helper function to define functions like `daily_max`. +""" +function common_diagnostics( + period, + reduction, + output_writer, + short_names...; + pre_output_hook! = nothing, +) + return [ + ScheduledDiagnosticTime( + variable = get_diagnostic_variable(short_name), + compute_every = :timestep, + output_every = period, # seconds + reduction_time_func = reduction, + output_writer = output_writer, + pre_output_hook! = pre_output_hook!, + ) for short_name in short_names + ] +end + +function average_pre_output_hook!(accum, counter) + @. accum = accum / counter + nothing +end + +""" + daily_maxs(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the daily max for the given variables. +""" +daily_maxs(short_names...; output_writer = HDF5Writer()) = + common_diagnostics(24 * 60 * 60, max, output_writer, short_names...) +""" + daily_max(short_names; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that computes the daily max for the given variable. +""" +daily_max(short_names; output_writer = HDF5Writer()) = + daily_maxs(short_names; output_writer)[1] + +""" + daily_mins(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the daily min for the given variables. +""" +daily_mins(short_names...; output_writer = HDF5Writer()) = + common_diagnostics(24 * 60 * 60, min, output_writer, short_names...) +""" + daily_min(short_names; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that computes the daily min for the given variable. +""" +daily_min(short_names; output_writer = HDF5Writer()) = + daily_mins(short_names; output_writer)[1] + +""" + daily_averages(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the daily average for the given variables. +""" +# An average is just a sum with a normalization before output +daily_averages(short_names...; output_writer = HDF5Writer()) = + common_diagnostics( + 24 * 60 * 60, + (+), + output_writer, + short_names...; + pre_output_hook! = average_pre_output_hook!, + ) +""" + daily_average(short_names; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that compute the daily average for the given variable. +""" +# An average is just a sum with a normalization before output +daily_average(short_names; output_writer = HDF5Writer()) = + daily_averages(short_names; output_writer)[1] +""" + hourly_maxs(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the hourly max for the given variables. +""" +hourly_maxs(short_names...; output_writer = HDF5Writer()) = + common_diagnostics(60 * 60, max, output_writer, short_names...) + +""" + hourly_max(short_names; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that computse the hourly max for the given variable. +""" +hourly_max(short_names...; output_writer = HDF5Writer()) = + hourly_maxs(short_names)[1] + +""" + hourly_mins(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the hourly min for the given variables. +""" +hourly_mins(short_names...; output_writer = HDF5Writer()) = + common_diagnostics(60 * 60, min, output_writer, short_names...) +""" + hourly_mins(short_names...; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that computes the hourly min for the given variable. +""" +hourly_min(short_names; output_writer = HDF5Writer()) = + hourly_mins(short_names; output_writer)[1] + +# An average is just a sum with a normalization before output +""" + hourly_averages(short_names...; output_writer = HDF5Writer()) + + +Return a list of `ScheduledDiagnostics` that compute the hourly average for the given variables. +""" +hourly_averages(short_names...; output_writer = HDF5Writer()) = + common_diagnostics( + 60 * 60, + (+), + output_writer, + short_names...; + pre_output_hook! = average_pre_output_hook!, + ) + +""" + hourly_average(short_names...; output_writer = HDF5Writer()) + + +Return a `ScheduledDiagnostics` that computes the hourly average for the given variable. +""" +hourly_average(short_names; output_writer = HDF5Writer()) = + hourly_averages(short_names; output_writer)[1] + +# Include all the subdefaults + +######## +# Core # +######## +function core_default_diagnostics() + core_diagnostics = ["ts", "ta", "thetaa", "pfull", "rhoa", "ua", "va", "wa"] + + return [ + daily_averages(core_diagnostics...)..., + daily_max("ts"), + daily_min("ts"), + ] +end + +############### +# Energy form # +############### +function default_diagnostics(::TotalEnergy) + total_energy_diagnostics = ["hfes"] + + return [daily_averages(total_energy_diagnostics...)...] +end + + +################## +# Moisture model # +################## +function default_diagnostics( + ::T, +) where {T <: Union{EquilMoistModel, NonEquilMoistModel}} + moist_diagnostics = ["hur", "hus", "hussfc", "evspsbl"] + + return [daily_averages(moist_diagnostics...)...] +end + +####################### +# Precipitation model # +####################### +function default_diagnostics(::Microphysics0Moment) + precip_diagnostics = ["pr"] + + return [daily_averages(precip_diagnostics...)...] +end + +################## +# Radiation mode # +################## +function default_diagnostics(::RRTMGPI.AbstractRRTMGPMode) + allsky_diagnostics = ["rsd", "rsu", "rld", "rlu"] + + return [daily_averages(allsky_diagnostics...)...] +end + + +function default_diagnostics(::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics) + clear_diagnostics = ["rsdcs", "rsucs", "rldcs", "rlucs"] + + return [daily_averages(clear_diagnostics...)...] +end diff --git a/src/diagnostics/diagnostic.jl b/src/diagnostics/diagnostic.jl new file mode 100644 index 00000000000..cd83ee7f7d9 --- /dev/null +++ b/src/diagnostics/diagnostic.jl @@ -0,0 +1,735 @@ +# diagnostic.jl +# +# This file contains: +# +# - The definition of what a DiagnosticVariable is. Conceptually, a DiagnosticVariable is a +# variable we know how to compute from the state. We attach more information to it for +# documentation and to reference to it with its short name. DiagnosticVariables can exist +# irrespective of the existence of an actual simulation that is being run. ClimaAtmos +# comes with several diagnostics already defined (in the `ALL_DIAGNOSTICS` dictionary). +# +# - A dictionary `ALL_DIAGNOSTICS` with all the diagnostics we know how to compute, keyed +# over their short name. If you want to add more diagnostics, look at the included files. +# You can add your own file if you want to define several new diagnostics that are +# conceptually related. The dictionary `ALL_DIAGNOSTICS` should be considered an +# implementation detail. +# +# - The definition of what a ScheduledDiagnostics is. Conceptually, a ScheduledDiagnostics is a +# DiagnosticVariable we want to compute in a given simulation. For example, it could be +# the temperature averaged over a day. We can have multiple ScheduledDiagnostics for the +# same DiagnosticVariable (e.g., daily and monthly average temperatures). +# +# We provide two types of ScheduledDiagnostics: ScheduledDiagnosticIterations and +# ScheduledDiagnosticTime, with the difference being only in what domain the recurrence +# time is defined (are we doing something at every N timesteps or every T seconds?). It is +# much cleaner and simpler to work with ScheduledDiagnosticIterations because iterations +# are well defined and consistent. On the other hand, working in the time domain requires +# dealing with what happens when the timestep is not lined up with the output period. +# Possible solutions to this problem include: uneven output, interpolation, or restricting +# the user from picking specific combinations of timestep/output period. In the current +# implementation, we choose the third option. So, ScheduledDiagnosticTime is provided +# because it is the physically interesting quantity. If we know what is the timestep, we +# can convert between the two and check if the diagnostics are well-posed in terms of the +# relationship between the periods and the timesteps. In some sense, you can think of +# ScheduledDiagnosticIterations as an internal representation and ScheduledDiagnosticTime +# as the external interface. +# +# - A function to convert a list of ScheduledDiagnosticIterations into a list of +# AtmosCallbacks. This function takes three arguments: the list of diagnostics and two +# dictionaries that map each scheduled diagnostic to an area of memory where to save the +# result and where to keep track of how many times the function was called (so that we +# can compute stuff like averages). +# +# - This file also also include several other files, including (but not limited to): +# - core_diagnostics.jl, radiation_diagnostics.jl +# - default_diagnostics.jl (which defines all the higher-level interfaces and defaults) +# - reduction_identities.jl +# - diagnostic_utils.jl + + +""" + DiagnosticVariable + + +A recipe to compute a diagnostic variable from the state, along with some useful metadata. + +The primary use for `DiagnosticVariable`s is to be embedded in a `ScheduledDiagnostic` to +compute diagnostics while the simulation is running. + +The metadata is used exclusively by the `output_writer` in the `ScheduledDiagnostic`. It is +responsibility of the `output_writer` to follow the conventions about the meaning of the +metadata and their use. + +In `ClimaAtmos`, we roughly follow the naming conventions listed in this file: +https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb + +Keyword arguments +================= + +- `short_name`: Name used to identify the variable in the output files and in the file + names. Short but descriptive. `ClimaAtmos` follows the CMIP conventions and + the diagnostics are identified by the short name. + +- `long_name`: Name used to describe the variable in the output files. + +- `standard_name`: Standard name, as in + http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html + +- `units`: Physical units of the variable. + +- `comments`: More verbose explanation of what the variable is, or comments related to how + it is defined or computed. + +- `compute!`: Function that compute the diagnostic variable from the state. + It has to take two arguments: the `integrator`, and a + pre-allocated area of memory where to write the result of the + computation. It the no pre-allocated area is available, a new + one will be allocated. To avoid extra allocations, this + function should perform the calculation in-place (i.e., using + `.=`). +""" +Base.@kwdef struct DiagnosticVariable{T} + short_name::String + long_name::String + standard_name::String + units::String + comments::String + compute!::T +end + +# ClimaAtmos diagnostics + +const ALL_DIAGNOSTICS = Dict{String, DiagnosticVariable}() + +""" + + add_diagnostic_variable!(; short_name, + long_name, + standard_name, + units, + description, + compute!) + + +Add a new variable to the `ALL_DIAGNOSTICS` dictionary (this function mutates the state of +`ClimaAtmos.ALL_DIAGNOSTICS`). + +If possible, please follow the naming scheme outline in +https://airtable.com/appYNLuWqAgzLbhSq/shrKcLEdssxb8Yvcp/tblL7dJkC3vl5zQLb + +Keyword arguments +================= + + +- `short_name`: Name used to identify the variable in the output files and in the file + names. Short but descriptive. `ClimaAtmos` diagnostics are identified by the + short name. We follow the Coupled Model Intercomparison Project conventions. + +- `long_name`: Name used to identify the variable in the output files. + +- `standard_name`: Standard name, as in + http://cfconventions.org/Data/cf-standard-names/71/build/cf-standard-name-table.html + +- `units`: Physical units of the variable. + +- `comments`: More verbose explanation of what the variable is, or comments related to how + it is defined or computed. + +- `compute!`: Function that compute the diagnostic variable from the state. + It has to take two arguments: the `integrator`, and a + pre-allocated area of memory where to write the result of the + computation. It the no pre-allocated area is available, a new + one will be allocated. To avoid extra allocations, this + function should perform the calculation in-place (i.e., using + `.=`). + +""" +function add_diagnostic_variable!(; + short_name, + long_name, + standard_name = "", + units, + comments = "", + compute!, +) + haskey(ALL_DIAGNOSTICS, short_name) && + error("diagnostic $short_name already defined") + + ALL_DIAGNOSTICS[short_name] = DiagnosticVariable(; + short_name, + long_name, + standard_name, + units, + comments, + compute!, + ) +end + +""" + + get_diagnostic_variable!(short_name) + + +Return a `DiagnosticVariable` from its `short_name`, if it exists. +""" +function get_diagnostic_variable(short_name) + haskey(ALL_DIAGNOSTICS, short_name) || + error("diagnostic $short_name does not exist") + + return ALL_DIAGNOSTICS[short_name] +end + +# Do you want to define more diagnostics? Add them here +include("core_diagnostics.jl") +include("radiation_diagnostics.jl") + +# Default diagnostics and higher level interfaces +include("default_diagnostics.jl") + +# Helper functions +include("diagnostics_utils.jl") + +# ScheduledDiagnostics + +# NOTE: The definitions of ScheduledDiagnosticTime and ScheduledDiagnosticIterations are +# nearly identical except for the fact that one is assumed to use units of seconds the other +# units of integration steps. However, we allow for this little repetition of code to avoid +# adding an extra layer of abstraction just to deal with these two objects (some people say +# that "duplication is better than over-abstraction"). Most users will only work with +# ScheduledDiagnosticTime. (It would be nice to have defaults fields in abstract types, as +# proposed in 2013 in https://github.com/JuliaLang/julia/issues/4935) Having two distinct +# types allow us to implement different checks and behaviors (e.g., we allow +# ScheduledDiagnosticTime to have placeholders values for {compute, output}_every so that we +# can plug the timestep in it). + +# Design note: pre_output_hook! +# +# One of our key requirements is to be able to compute arithmetic averages. Unfortunately, +# computing an arithmetic average requires keeping track of how many elements we are summing +# up. pre_output_hook! was introduced so that we can think of an average as a sum coupled +# with division, and perform the division (by the number of elements) before output. +# pre_output_hook! could be used for other operations, but we decided to keep it simple and +# target directly the most important use case for us. +# +# This choice restricts what reductions can be performed. For example, it is not possible to +# have a geometric average. If more complex reduction are needed, this mechanism has to be +# changed. + +struct ScheduledDiagnosticIterations{T1, T2, OW, F1, F2, PO} + variable::DiagnosticVariable + output_every::T1 + output_writer::OW + reduction_time_func::F1 + reduction_space_func::F2 + compute_every::T2 + pre_output_hook!::PO + output_short_name::String + output_long_name::String + + """ + ScheduledDiagnosticIterations(; variable::DiagnosticVariable, + output_every, + output_writer, + reduction_time_func = nothing, + reduction_space_func = nothing, + compute_every = isa_reduction ? 1 : output_every, + pre_output_hook! = nothing, + output_short_name = descriptive_short_name(self), + output_short_name = descriptive_long_name(self)) + + + A `DiagnosticVariable` that has to be computed and output during a simulation with a cadence + defined by the number of iterations, with an optional reduction applied to it (e.g., compute + the maximum temperature over the course of every 10 timesteps). This object is turned into + two callbacks (one for computing and the other for output) and executed by the integrator. + + Keyword arguments + ================= + + - `variable`: The diagnostic variable that has to be computed and output. + + - `output_every`: Save the results to disk every `output_every` iterations. + + - `output_writer`: Function that controls out to save the computed diagnostic variable to + disk. `output_writer` has to take three arguments: the value that has to + be output, the `ScheduledDiagnostic`, and the integrator. Internally, the + integrator contains extra information (such as the current timestep). It + is responsibility of the `output_writer` to properly use the provided + information for meaningful output. + + - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of + scratch space `acc` where to accumulate partial results. Then, at + every `compute_every`, `reduction_time_func` is computed between + the previously stored value in `acc` and the new value. This + implements a running reduction. For example, if + `reduction_time_func = max`, the space `acc` will hold the running + maxima of the diagnostic. To implement operations like the + arithmetic average, the `reduction_time_func` has to be chosen as + `sum`, and a `pre_output_hook!` that renormalize `acc` by the + number of samples has to be provided. For custom reductions, it is + necessary to also specify the identity of operation by defining a + new method to `identity_of_reduction`. + + - `reduction_space_func`: NOT IMPLEMENTED YET + + - `compute_every`: Run the computations every `compute_every` iterations. This is not + particularly useful for point-wise diagnostics, where we enforce that + `compute_every` = `output_every`. For time reductions, `compute_every` is + set to 1 (compute at every timestep) by default. `compute_every` has to + evenly divide `output_every`. + + - `pre_output_hook!`: Function that has to be run before saving to disk for reductions + (mostly used to implement averages). The function `pre_output_hook!` + is called with two arguments: the value accumulated during the + reduction, and the number of times the diagnostic was computed from + the last time it was output. `pre_output_hook!` should mutate the + accumulator in place. The return value of `pre_output_hook!` is + discarded. An example of `pre_output_hook!` to compute the arithmetic + average is `pre_output_hook!(acc, N) = @. acc = acc / N`. + + - `output_short_name`: A descriptive name for this particular diagnostic. If none is + provided, one will be generated mixing the short name of the + variable, the reduction, and the period of the reduction. + Normally, it has to be unique. In `ClimaAtmos`, we follow the CMIP + conventions for this. + + - `output_long_name`: A descriptive name for this particular diagnostic. If none is + provided, one will be generated mixing the short name of the + variable, the reduction, and the period of the reduction. + + """ + function ScheduledDiagnosticIterations(; + variable::DiagnosticVariable, + output_every, + output_writer, + reduction_time_func = nothing, + reduction_space_func = nothing, + compute_every = isnothing(reduction_time_func) ? output_every : 1, + pre_output_hook! = nothing, + output_short_name = descriptive_short_name( + variable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = false, + ), + output_long_name = descriptive_long_name( + variable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = false, + ), + ) + + # We provide an inner constructor to enforce some constraints + + output_every % compute_every == 0 || error( + "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_short_name)", + ) + + isa_reduction = !isnothing(reduction_time_func) + + # If it is not a reduction, we compute only when we output + if !isa_reduction && compute_every != output_every + @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_short_name), changing compute_every to match" + compute_every = output_every + end + + # pre_output_hook! has to be a function, but it is much more intuitive to specify + # `nothing` when we want nothing to happen. Here, we convert the nothing keyword + # into a function that does nothing + if isnothing(pre_output_hook!) + pre_output_hook! = (accum, count) -> nothing + end + + T1 = typeof(output_every) + T2 = typeof(compute_every) + OW = typeof(output_writer) + F1 = typeof(reduction_time_func) + F2 = typeof(reduction_space_func) + PO = typeof(pre_output_hook!) + + new{T1, T2, OW, F1, F2, PO}( + variable, + output_every, + output_writer, + reduction_time_func, + reduction_space_func, + compute_every, + pre_output_hook!, + output_short_name, + output_long_name, + ) + end +end + + +struct ScheduledDiagnosticTime{T1, T2, OW, F1, F2, PO} + variable::DiagnosticVariable + output_every::T1 + output_writer::OW + reduction_time_func::F1 + reduction_space_func::F2 + compute_every::T2 + pre_output_hook!::PO + output_short_name::String + output_long_name::String + + """ + ScheduledDiagnosticTime(; variable::DiagnosticVariable, + output_every, + output_writer, + reduction_time_func = nothing, + reduction_space_func = nothing, + compute_every = isa_reduction ? :timestep : output_every, + pre_output_hook! = nothing, + output_short_name = descriptive_short_name(self), + output_long_name = descriptive_long_name(self), + ) + + + A `DiagnosticVariable` that has to be computed and output during a simulation with a + cadence defined by how many seconds in simulation time, with an optional reduction + applied to it (e.g., compute the maximum temperature over the course of every day). This + object is turned into a `ScheduledDiagnosticIterations`, which is turned into two + callbacks (one for computing and the other for output) and executed by the integrator. + + Keyword arguments + ================= + + - `variable`: The diagnostic variable that has to be computed and output. + + - `output_every`: Save the results to disk every `output_every` seconds. + + - `output_writer`: Function that controls out to save the computed diagnostic variable to + disk. `output_writer` has to take three arguments: the value that has to + be output, the `ScheduledDiagnostic`, and the integrator. Internally, the + integrator contains extra information (such as the current timestep). It + is responsibility of the `output_writer` to properly use the provided + information for meaningful output. + + - `reduction_time_func`: If not `nothing`, this `ScheduledDiagnostic` receives an area of + scratch space `acc` where to accumulate partial results. Then, at + every `compute_every`, `reduction_time_func` is computed between + the previously stored value in `acc` and the new value. This + implements a running reduction. For example, if + `reduction_time_func = max`, the space `acc` will hold the running + maxima of the diagnostic. To implement operations like the + arithmetic average, the `reduction_time_func` has to be chosen as + `sum`, and a `pre_output_hook!` that renormalize `acc` by the + number of samples has to be provided. For custom reductions, it is + necessary to also specify the identity of operation by defining a + new method to `identity_of_reduction`. + + - `reduction_space_func`: NOT IMPLEMENTED YET + + - `compute_every`: Run the computations every `compute_every` seconds. This is not + particularly useful for point-wise diagnostics, where we enforce that + `compute_every` = `output_every`. For time reductions, + `compute_every` is set to `:timestep` (compute at every timestep) by + default. `compute_every` has to evenly divide `output_every`. + `compute_every` can take the special symbol `:timestep` which is a + placeholder for the timestep of the simulation to which this + `ScheduledDiagnostic` is attached. + + - `pre_output_hook!`: Function that has to be run before saving to disk for reductions + (mostly used to implement averages). The function `pre_output_hook!` + is called with two arguments: the value accumulated during the + reduction, and the number of times the diagnostic was computed from + the last time it was output. `pre_output_hook!` should mutate the + accumulator in place. The return value of `pre_output_hook!` is + discarded. An example of `pre_output_hook!` to compute the arithmetic + average is `pre_output_hook!(acc, N) = @. acc = acc / N`. + + - `output_short_name`: A descriptive name for this particular diagnostic. If none is + provided, one will be generated mixing the short name of the + variable, the reduction, and the period of the reduction. + Normally, it has to be unique. In `ClimaAtmos`, we follow the CMIP + conventions for this. + + - `output_long_name`: A descriptive name for this particular diagnostic. If none is + provided, one will be generated mixing the short name of the + variable, the reduction, and the period of the reduction. + """ + function ScheduledDiagnosticTime(; + variable::DiagnosticVariable, + output_every, + output_writer, + reduction_time_func = nothing, + reduction_space_func = nothing, + compute_every = isnothing(reduction_time_func) ? output_every : + :timestep, + pre_output_hook! = nothing, + output_short_name = descriptive_short_name( + variable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true, + ), + output_long_name = descriptive_long_name( + variable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true, + ), + ) + + # We provide an inner constructor to enforce some constraints + + # compute_every could be a Symbol (:timestep). We process this that when we process + # the list of diagnostics + if !isa(compute_every, Symbol) + output_every % compute_every == 0 || error( + "output_every ($output_every) should be multiple of compute_every ($compute_every) for diagnostic $(output_short_name)", + ) + end + + isa_reduction = !isnothing(reduction_time_func) + + # If it is not a reduction, we compute only when we output + if !isa_reduction && compute_every != output_every + @warn "output_every ($output_every) != compute_every ($compute_every) for $(output_short_name), changing compute_every to match" + compute_every = output_every + end + + # pre_output_hook! has to be a function, but it is much more intuitive to specify + # `nothing` when we want nothing to happen. Here, we convert the nothing keyword + # into a function that does nothing + if isnothing(pre_output_hook!) + pre_output_hook! = (accum, count) -> nothing + end + + T1 = typeof(output_every) + T2 = typeof(compute_every) + OW = typeof(output_writer) + F1 = typeof(reduction_time_func) + F2 = typeof(reduction_space_func) + PO = typeof(pre_output_hook!) + + new{T1, T2, OW, F1, F2, PO}( + variable, + output_every, + output_writer, + reduction_time_func, + reduction_space_func, + compute_every, + pre_output_hook!, + output_short_name, + output_long_name, + ) + end +end + +""" + ScheduledDiagnosticIterations(sd_time::ScheduledDiagnosticTime, Δt) + + +Create a `ScheduledDiagnosticIterations` given a `ScheduledDiagnosticTime` and a timestep +`Δt`. In this, ensure that `compute_every` and `output_every` are meaningful for the given +timestep. + +""" +function ScheduledDiagnosticIterations( + sd_time::ScheduledDiagnosticTime, + Δt::T, +) where {T} + + # If we have the timestep, we can convert time in seconds into iterations + + # if compute_every is :timestep, then we want to compute after every iterations + compute_every = + sd_time.compute_every == :timestep ? 1 : sd_time.compute_every / Δt + output_every = sd_time.output_every / Δt + + isinteger(output_every) || error( + "output_every ($(sd_time.output_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_short_name)", + ) + isinteger(compute_every) || error( + "compute_every ($(sd_time.compute_every)) should be multiple of the timestep ($Δt) for diagnostic $(sd_time.output_short_name)", + ) + + ScheduledDiagnosticIterations(; + sd_time.variable, + output_every = convert(Int, output_every), + sd_time.output_writer, + sd_time.reduction_time_func, + sd_time.reduction_space_func, + compute_every = convert(Int, compute_every), + sd_time.pre_output_hook!, + sd_time.output_short_name, + sd_time.output_long_name, + ) +end + +""" + ScheduledDiagnosticTime(sd_time::ScheduledDiagnosticIterations, Δt) + + +Create a `ScheduledDiagnosticTime` given a `ScheduledDiagnosticIterations` and a timestep +`Δt`. + +""" +function ScheduledDiagnosticTime( + sd_time::ScheduledDiagnosticIterations, + Δt::T, +) where {T} + + # If we have the timestep, we can convert time in iterations to seconds + + # if compute_every is :timestep, then we want to compute after every iterations + compute_every = + sd_time.compute_every == 1 ? :timestep : sd_time.compute_every * Δt + output_every = sd_time.output_every * Δt + + ScheduledDiagnosticTime(; + sd_time.variable, + output_every, + sd_time.output_writer, + sd_time.reduction_time_func, + sd_time.reduction_space_func, + compute_every, + sd_time.pre_output_hook!, + sd_time.output_short_name, + sd_time.output_long_name, + ) +end + +# We provide also a companion constructor for ScheduledDiagnosticIterations which returns +# itself (without copy) when called with a timestep. +# +# This is so that we can assume that +# ScheduledDiagnosticIterations(ScheduledDiagnostic{Time, Iterations}, Δt) +# always returns a valid ScheduledDiagnosticIterations +ScheduledDiagnosticIterations( + sd::ScheduledDiagnosticIterations, + _Δt::T, +) where {T} = sd + + +# We define all the known identities in reduction_identities.jl +include("reduction_identities.jl") + +# Helper functions for the callbacks: +# - reset_accumulator! +# - accumulate! + +# When the reduction is nothing, do nothing +reset_accumulator!(_, reduction_time_func::Nothing) = nothing + +# If we have a reduction, we have to reset the accumulator to its neutral state. (If we +# don't have a reduction, we don't have to do anything) +# +# ClimaAtmos defines methods for identity_of_reduction for standard reduction_time_func in +# reduction_identities.jl +function reset_accumulator!(diag_accumulator, reduction_time_func) + # identity_of_reduction works by dispatching over operation + identity = identity_of_reduction(reduction_time_func) + float_type = Spaces.undertype(axes((diag_accumulator))) + identity_ft = convert(float_type, identity) + fill!(parent(diag_accumulator), identity_ft) +end + +# When the reduction is nothing, we do not need to accumulate anything +accumulate!(_, _, reduction_time_func::Nothing) = nothing + +# When we have a reduction, apply it between the accumulated value one +function accumulate!(diag_accumulator, diag_storage, reduction_time_func) + diag_accumulator .= reduction_time_func.(diag_accumulator, diag_storage) + return nothing +end + +""" + get_callbacks_from_diagnostics(diagnostics, storage, counters) + + +Translate a list of diagnostics into a list of callbacks. + +Positional arguments +===================== + +- `diagnostics`: List of `ScheduledDiagnosticIterations` that have to be converted to + callbacks. We want to have `ScheduledDiagnosticIterations` here so that we + can define callbacks that occur at the end of every N integration steps. + +- `storage`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially + pre-allocated area of memory where to save the newly computed results. + +- `accumulator`: Dictionary that maps a given `ScheduledDiagnosticIterations` to a potentially + pre-allocated area of memory where to accumulate results. + +- `counters`: Dictionary that maps a given `ScheduledDiagnosticIterations` to the counter + that tracks how many times the given diagnostics was computed from the last + time it was output to disk. + +""" +function get_callbacks_from_diagnostics( + diagnostics, + storage, + accumulators, + counters, +) + # We have two types of callbacks: to compute and accumulate diagnostics, and to dump + # them to disk. Note that our callbacks do not contain any branching + + # storage is used to pre-allocate memory and to accumulate partial results for those + # diagnostics that perform reductions. + + callback_arrays = map(diagnostics) do diag + variable = diag.variable + compute_callback = + integrator -> begin + variable.compute!( + storage[diag], + integrator.u, + integrator.p, + integrator.t, + ) + + # accumulator[diag] is not defined for non-reductions + diag_accumulator = get(accumulators, diag, nothing) + + accumulate!( + diag_accumulator, + storage[diag], + diag.reduction_time_func, + ) + counters[diag] += 1 + return nothing + end + + output_callback = + integrator -> begin + # Move accumulated value to storage so that we can output it (for + # reductions). This provides a unified interface to pre_output_hook! and + # output, at the cost of an additional copy. If this copy turns out to be + # too expensive, we can move the if statement below. + isnothing(diag.reduction_time_func) || + (storage[diag] .= accumulators[diag]) + + # Any operations we have to perform before writing to output? + # Here is where we would divide by N to obtain an arithmetic average + diag.pre_output_hook!(storage[diag], counters[diag]) + + # Write to disk + diag.output_writer(storage[diag], diag, integrator) + + # accumulator[diag] is not defined for non-reductions + diag_accumulator = get(accumulators, diag, nothing) + + reset_accumulator!(diag_accumulator, diag.reduction_time_func) + counters[diag] = 0 + return nothing + end + + return [ + AtmosCallback(compute_callback, EveryNSteps(diag.compute_every)), + AtmosCallback(output_callback, EveryNSteps(diag.output_every)), + ] + end + + # We need to flatten to tuples + return vcat(callback_arrays...) +end diff --git a/src/diagnostics/diagnostics_utils.jl b/src/diagnostics/diagnostics_utils.jl new file mode 100644 index 00000000000..1580c66cbf9 --- /dev/null +++ b/src/diagnostics/diagnostics_utils.jl @@ -0,0 +1,128 @@ +# diagnostic_utils.jl +# +# This file contains: +# - descriptive_short_name: to condense ScheduledDiagnostic information into few characters. +# - descriptive_long_name: to produce full names that are clearly human-understandable + + +""" + descriptive_short_name(variable::DiagnosticVariable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true) + + +Return a compact, unique-ish, identifier generated from the given information. + +`output_every` is interpreted as in seconds if `units_are_seconds` is `true`. Otherwise, it +is interpreted as in units of number of iterations. + +This function is useful for filenames and error messages. + +""" +function descriptive_short_name( + variable::DiagnosticVariable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true, +) + var = "$(variable.short_name)" + isa_reduction = !isnothing(reduction_time_func) + + + if isa_reduction + red = "$(reduction_time_func)" + + # Let's check if we are computing the average. Note that this might slip under the + # radar if the user passes their own pre_output_hook!. + if reduction_time_func == (+) && + pre_output_hook! == average_pre_output_hook! + red = "average" + end + + if units_are_seconds + + # Convert period from seconds to days, hours, minutes, seconds + period = "" + + days, rem_seconds = divrem(output_every, 24 * 60 * 60) + hours, rem_seconds = divrem(rem_seconds, 60 * 60) + minutes, seconds = divrem(rem_seconds, 60) + + days > 0 && (period *= "$(days)d_") + hours > 0 && (period *= "$(hours)h_") + minutes > 0 && (period *= "$(minutes)m_") + seconds > 0 && (period *= "$(seconds)s_") + + suffix = period * red + else + suffix = "$(output_every)it_$(red)" + end + else + suffix = "inst" + end + return "$(var)_$(suffix)" +end + +""" + descriptive_long_name(variable::DiagnosticVariable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true) + + +Return a verbose description of the given output variable. + +`output_every` is interpreted as in seconds if `units_are_seconds` is `true`. Otherwise, it +is interpreted as in units of number of iterations. + +This function is useful for attributes in output files. + +""" +function descriptive_long_name( + variable::DiagnosticVariable, + output_every, + reduction_time_func, + pre_output_hook!; + units_are_seconds = true, +) + var = "$(variable.long_name)" + isa_reduction = !isnothing(reduction_time_func) + + if isa_reduction + red = "$(reduction_time_func)" + + # Let's check if we are computing the average. Note that this might slip under the + # radar if the user passes their own pre_output_hook!. + if reduction_time_func == (+) && + pre_output_hook! == average_pre_output_hook! + red = "average" + end + + if units_are_seconds + + # Convert period from seconds to days, hours, minutes, seconds + period = "" + + days, rem_seconds = divrem(output_every, 24 * 60 * 60) + hours, rem_seconds = divrem(rem_seconds, 60 * 60) + minutes, seconds = divrem(rem_seconds, 60) + + days > 0 && (period *= "$(days) Day(s)") + hours > 0 && (period *= "$(hours) Hour(s)") + minutes > 0 && (period *= "$(minutes) Minute(s)") + seconds > 0 && (period *= "$(seconds) Second(s)") + + period_str = period * red + else + period_str = "$(output_every) Iterations" + end + suffix = "$(red) within $(period_str)" + else + suffix = "Instantaneous" + end + return "$(var), $(suffix)" +end diff --git a/src/diagnostics/radiation_diagnostics.jl b/src/diagnostics/radiation_diagnostics.jl new file mode 100644 index 00000000000..fcded7d982c --- /dev/null +++ b/src/diagnostics/radiation_diagnostics.jl @@ -0,0 +1,307 @@ +# This file is included in Diagnostics.jl + +# Radiative fluxes + +### +# Downwelling shortwave radiation (3d) +### +compute_rsd!(out, state, cache, time) = + compute_rsd!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rsd!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rsd", radiation_mode) + +function compute_rsd!( + out, + state, + cache, + time, + radiation_mode::T, +) where {T <: RRTMGPI.AbstractRRTMGPMode} + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_sw_flux_dn), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_sw_flux_dn), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rsd", + long_name = "Downwelling Shortwave Radiation", + standard_name = "surface_downwelling_shortwave_flux_in_air", + units = "W m^-2", + comments = "Downwelling shortwave radiation", + compute! = compute_rsd!, +) + +### +# Upwelling shortwave radiation (3d) +### +compute_rsu!(out, state, cache, time) = + compute_rsu!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rsu!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rsu", radiation_mode) + +function compute_rsu!( + out, + state, + cache, + time, + radiation_mode::T, +) where {T <: RRTMGPI.AbstractRRTMGPMode} + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_sw_flux_up), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_sw_flux_up), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rsu", + long_name = "Upwelling Shortwave Radiation", + standard_name = "surface_upwelling_shortwave_flux_in_air", + units = "W m^-2", + comments = "Upwelling shortwave radiation", + compute! = compute_rsu!, +) + +### +# Downwelling longwave radiation (3d) +### +compute_rld!(out, state, cache, time) = + compute_rld!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rld!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rld", radiation_mode) + +function compute_rld!( + out, + state, + cache, + time, + radiation_mode::T, +) where {T <: RRTMGPI.AbstractRRTMGPMode} + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_lw_flux_dn), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_lw_flux_dn), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rld", + long_name = "Downwelling Longwave Radiation", + standard_name = "surface_downwelling_longwave_flux_in_air", + units = "W m^-2", + comments = "Downwelling longwave radiation", + compute! = compute_rld!, +) + +### +# Upwelling longwave radiation (3d) +### +compute_rlu!(out, state, cache, time) = + compute_rlu!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rlu!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rlu", radiation_mode) + +function compute_rlu!( + out, + state, + cache, + time, + radiation_mode::T, +) where {T <: RRTMGPI.AbstractRRTMGPMode} + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_lw_flux_up), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_lw_flux_up), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rlu", + long_name = "Upwelling Longwave Radiation", + standard_name = "surface_upwelling_longwave_flux_in_air", + units = "W m^-2", + comments = "Upwelling longwave radiation", + compute! = compute_rlu!, +) + +### +# Downelling clear sky shortwave radiation (3d) +### +compute_rsdcs!(out, state, cache, time) = + compute_rsdcs!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rsdcs!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rsdcs", radiation_mode) + +function compute_rsdcs!( + out, + state, + cache, + time, + radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics, +) + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_sw_flux_dn), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_sw_flux_dn), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rsdcs", + long_name = "Downwelling Clear-Sky Shortwave Radiation", + standard_name = "surface_downwelling_shortwave_flux_in_air_assuming_clear_sky", + units = "W m^-2", + comments = "Downwelling clear sky shortwave radiation", + compute! = compute_rsdcs!, +) + +### +# Upwelling clear sky shortwave radiation (3d) +### +compute_rsucs!(out, state, cache, time) = + compute_rsucs!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rsucs!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rsucs", radiation_mode) + +function compute_rsucs!( + out, + state, + cache, + time, + radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics, +) + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_sw_flux_up), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_sw_flux_up), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rsucs", + long_name = "Upwelling Clear-Sky Shortwave Radiation", + standard_name = "surface_upwelling_shortwave_flux_in_air_assuming_clear_sky", + units = "W m^-2", + comments = "Upwelling clear sky shortwave radiation", + compute! = compute_rsucs!, +) + +### +# Downwelling clear sky longwave radiation (3d) +### +compute_rldcs!(out, state, cache, time) = + compute_rldcs!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rldcs!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rldcs", radiation_mode) + +function compute_rldcs!( + out, + state, + cache, + time, + radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics, +) + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_lw_flux_dn), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_lw_flux_dn), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rldcs", + long_name = "Downwelling Clear-Sky Longwave Radiation", + standard_name = "surface_downwelling_longwave_flux_in_air_assuming_clear_sky", + units = "W m^-2", + comments = "Downwelling clear sky longwave radiation", + compute! = compute_rldcs!, +) + +### +# Upwelling clear sky longwave radiation (3d) +### +compute_rlucs!(out, state, cache, time) = + compute_rlucs!(out, state, cache, time, cache.atmos.radiation_mode) +compute_rlucs!(_, _, _, _, radiation_mode::T) where {T} = + error_diagnostic_variable("rlucs", radiation_mode) + +function compute_rlucs!( + out, + state, + cache, + time, + radiation_mode::RRTMGPI.AllSkyRadiationWithClearSkyDiagnostics, +) + FT = eltype(state) + if isnothing(out) + return RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_lw_flux_up), + axes(state.f), + ) + else + out .= RRTMGPI.array2field( + FT.(cache.radiation_model.face_clear_lw_flux_up), + axes(state.f), + ) + end +end + +add_diagnostic_variable!( + short_name = "rlucs", + long_name = "Upwelling Clear-Sky Longwave Radiation", + standard_name = "surface_upwelling_longwave_flux_in_air_assuming_clear_sky", + units = "W m^-2", + comments = "Upwelling clear sky longwave radiation", + compute! = compute_rlucs!, +) diff --git a/src/diagnostics/reduction_identities.jl b/src/diagnostics/reduction_identities.jl new file mode 100644 index 00000000000..0fce5e58845 --- /dev/null +++ b/src/diagnostics/reduction_identities.jl @@ -0,0 +1,19 @@ +# This file defines group identities (technically _monoid identities_) and is included in +# Diagnostics.jl. +# +# Several diagnostics require performing reductions, such as taking the maximum or the +# average. Since it is not feasible to store all the lists of all the intermediate values, +# we aggregate the results in specific storage areas (e.g., we take +# max(max(max(max(t1, t2), t3), t4), t5) instead of max(t1, t2, t3, t4, t5) +# In this, it is convenient to preallocate the space where we want to accumulate the +# intermediate. However, we have to fill the space with something that does not affect the +# reduction. This, by definition, is the identity of the operation. The identity of the operation +# + is 0 because x + 0 = x for every x. +# +# We have to know the identity for every operation we want to support. Of course, users are +# welcome to define their own by adding new methods to identity_of_reduction. + +identity_of_reduction(::typeof(max)) = -Inf +identity_of_reduction(::typeof(min)) = +Inf +identity_of_reduction(::typeof(+)) = 0 +identity_of_reduction(::typeof(*)) = 1 diff --git a/src/diagnostics/writers.jl b/src/diagnostics/writers.jl new file mode 100644 index 00000000000..43de0c78257 --- /dev/null +++ b/src/diagnostics/writers.jl @@ -0,0 +1,164 @@ +# Writers.jl +# +# This file defines function-generating functions for output_writers for diagnostics. The +# writers come with opinionated defaults. + +import ClimaCore.Remapping: interpolate_array +import NCDatasets + +""" + HDF5Writer() + + +Save a `ScheduledDiagnostic` to a HDF5 file inside the `output_dir` of the simulation. + + +TODO: This is a very barebone HDF5Writer. Do not consider this implementation as the "final +word". + +We need to implement the following features/options: +- Toggle for write new files/append +- Checks for existing files +- Check for new subfolders that have to be created +- More meaningful naming conventions (keeping in mind that we can have multiple variables + with different reductions) +- All variables in one file/each variable in its own file +- All timesteps in one file/each timestep in its own file +- Writing the correct attributes +- Overriding simulation.output_dir (e.g., if the path starts with /) +- ...more features/options + +""" +function HDF5Writer() + # output_drivers are called with the three arguments: the value, the ScheduledDiagnostic, + # and the integrator + function write_to_hdf5(value, diagnostic, integrator) + var = diagnostic.variable + time = integrator.t + + # diagnostic here is a ScheduledDiagnosticIteration. If we want to obtain a + # descriptive name (e.g., something with "daily"), we have to pass the timestep as + # well + + output_path = joinpath( + integrator.p.simulation.output_dir, + "$(diagnostic.output_short_name)_$(time).h5", + ) + + hdfwriter = InputOutput.HDF5Writer(output_path, integrator.p.comms_ctx) + InputOutput.write!(hdfwriter, value, "$(diagnostic.output_short_name)") + attributes = Dict( + "time" => time, + "long_name" => diagnostic.output_long_name, + "variable_units" => var.units, + "standard_variable_name" => var.standard_name, + ) + + # TODO: Use directly InputOutput functions + InputOutput.HDF5.h5writeattr( + hdfwriter.file.filename, + "fields/$(diagnostic.output_short_name)", + attributes, + ) + + Base.close(hdfwriter) + return nothing + end +end + + +""" + NetCDFWriter() + + +Save a `ScheduledDiagnostic` to a NetCDF file inside the `output_dir` of the simulation by +performing a pointwise (non-conservative) remapping first. This writer does not work on +distributed simulations. + + +TODO: This is a very barebone NetCDFWriter. This writer only supports 3D fields on cubed +spheres at the moment. Do not consider this implementation as the "final word". + +We need to implement the following features/options: +- Toggle for write new files/append +- Checks for existing files +- Check for new subfolders that have to be created +- More meaningful naming conventions (keeping in mind that we can have multiple variables + with different reductions) +- All variables in one file/each variable in its own file +- All timesteps in one file/each timestep in its own file +- Writing the correct attributes +- Overriding simulation.output_dir (e.g., if the path starts with /) +- ...more features/options + +""" +function NetCDFWriter(; + num_points_longitude = 100, + num_points_latitude = 100, + num_points_altitude = 100, + compression_level = 9, +) + function write_to_netcdf(field, diagnostic, integrator) + var = diagnostic.variable + time = integrator.t + + # TODO: Automatically figure out details on the remapping depending on the given + # field + + # Let's support only cubed spheres for the moment + typeof(axes(field).horizontal_space.topology.mesh) <: + Meshes.AbstractCubedSphere || + error("NetCDF writer supports only cubed sphere at the moment") + + # diagnostic here is a ScheduledDiagnosticIteration. If we want to obtain a + # descriptive name (e.g., something with "daily"), we have to pass the timestep as + # well + output_path = joinpath( + integrator.p.simulation.output_dir, + "$(diagnostic.output_short_name)_$time.nc", + ) + + vert_domain = axes(field).vertical_topology.mesh.domain + z_min, z_max = vert_domain.coord_min.z, vert_domain.coord_max.z + + FT = Spaces.undertype(axes(field)) + + longpts = range( + Geometry.LongPoint(-FT(180.0)), + Geometry.LongPoint(FT(180.0)), + length = num_points_longitude, + ) + latpts = range( + Geometry.LatPoint(-FT(80.0)), + Geometry.LatPoint(FT(80.0)), + length = num_points_latitude, + ) + zpts = range( + Geometry.ZPoint(FT(z_min)), + Geometry.ZPoint(FT(z_max)), + length = num_points_altitude, + ) + + nc = NCDatasets.Dataset(output_path, "c") + + NCDatasets.defDim(nc, "lon", num_points_longitude) + NCDatasets.defDim(nc, "lat", num_points_latitude) + NCDatasets.defDim(nc, "z", num_points_altitude) + + nc.attrib["long_name"] = diagnostic.output_long_name + nc.attrib["units"] = var.units + nc.attrib["comments"] = var.comments + + v = NCDatasets.defVar( + nc, + "$(var.short_name)", + FT, + ("lon", "lat", "z"), + deflatelevel = compression_level, + ) + + v[:, :, :] = interpolate_array(field, longpts, latpts, zpts) + + close(nc) + end +end diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl index 38762a61a9d..ca71174e567 100644 --- a/src/initial_conditions/initial_conditions.jl +++ b/src/initial_conditions/initial_conditions.jl @@ -470,10 +470,10 @@ Base.@kwdef struct MoistAdiabaticProfileEDMFX <: InitialCondition end draft_area(::Type{FT}) where {FT} = - z -> FT(0.5) * exp(-(z - FT(4000.0))^2 / 2 / FT(1000.0)^2) + z -> z < 0.7e4 ? FT(0.5) * exp(-(z - FT(4e3))^2 / 2 / FT(1e3)^2) : FT(0) edmfx_q_tot(::Type{FT}) where {FT} = - z -> FT(0.001) * exp(-(z - FT(4000.0))^2 / 2 / FT(1000.0)^2) + z -> z < 0.7e4 ? FT(1e-3) * exp(-(z - FT(4e3))^2 / 2 / FT(1e3)^2) : FT(0) function (initial_condition::MoistAdiabaticProfileEDMFX)(params) (; perturb) = initial_condition diff --git a/src/parameters/create_parameters.jl b/src/parameters/create_parameters.jl index f8c790536f6..b06cf796eca 100644 --- a/src/parameters/create_parameters.jl +++ b/src/parameters/create_parameters.jl @@ -7,9 +7,6 @@ import SurfaceFluxes.UniversalFunctions as UF import Insolation.Parameters as IP import Thermodynamics as TD import CloudMicrophysics as CM -# TODO: Remove these imports? -import ClimaCore -import ClimaCore as CC function create_parameter_set(config::AtmosConfig) # Helper function that creates a parameter struct. If a struct has nested @@ -70,7 +67,7 @@ function create_parameter_set(config::AtmosConfig) TCP.TurbulenceConvectionParameters; subparam_structs = (; microphys_params, surf_flux_params), ) - param_set = create_parameter_struct( + return create_parameter_struct( CAP.ClimaAtmosParameters; subparam_structs = (; thermodynamics_params = thermo_params, @@ -81,9 +78,4 @@ function create_parameter_set(config::AtmosConfig) turbconv_params, ), ) - if parsed_args["log_params"] - logfilepath = joinpath(@__DIR__, "$(parsed_args["job_id"])_$FT.toml") - CP.log_parameter_information(toml_dict, logfilepath) - end - return param_set end diff --git a/src/prognostic_equations/edmfx_sgs_flux.jl b/src/prognostic_equations/edmfx_sgs_flux.jl index 5a9709ef9e1..355205f52f8 100644 --- a/src/prognostic_equations/edmfx_sgs_flux.jl +++ b/src/prognostic_equations/edmfx_sgs_flux.jl @@ -11,14 +11,14 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX) (; edmfx_upwinding, sfc_conditions) = p (; ᶠu³, ᶜh_tot, ᶜspecific) = p (; ᶠu³ʲs, ᶜh_totʲs, ᶜspecificʲs) = p - (; ᶜρa⁰, ᶠu³⁰, ᶜh_tot⁰, ᶜspecific⁰) = p + (; ᶜρa⁰, ᶠu³⁰, ᶜu⁰, ᶜh_tot⁰, ᶜspecific⁰) = p (; ᶜK_u, ᶜK_h) = p (; dt) = p.simulation ᶜJ = Fields.local_geometry_field(Y.c).J ᶠgradᵥ = Operators.GradientC2F() if p.atmos.edmfx_sgs_flux - # mass flux + # energy mass flux ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx] ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx] for j in 1:n @@ -46,7 +46,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX) edmfx_upwinding, ) - # diffusive flux + # energy diffusive flux ᶠρaK_h = p.ᶠtemp_scalar @. ᶠρaK_h[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_h[colidx]) @@ -58,7 +58,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX) ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot⁰[colidx]))) if !(p.atmos.moisture_model isa DryModel) - # mass flux + # specific humidity mass flux for j in 1:n @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] @. ᶜq_tot_diff_colidx = @@ -86,7 +86,7 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX) edmfx_upwinding, ) - # diffusive flux + # specific humidity diffusive flux ᶜdivᵥ_ρq_tot = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue( @@ -98,19 +98,23 @@ function edmfx_sgs_flux_tendency!(Yₜ, Y, p, t, colidx, turbconv_model::EDMFX) ) end - # diffusive flux + # momentum diffusive flux ᶠρaK_u = p.ᶠtemp_scalar @. ᶠρaK_u[colidx] = ᶠinterp(ᶜρa⁰[colidx]) * ᶠinterp(ᶜK_u[colidx]) + ᶠstrain_rate = p.ᶠtemp_UVWxUVW + compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu⁰[colidx]) + @. Yₜ.c.uₕ[colidx] -= C12( + ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx], + ) + # apply boundary condition for momentum flux ᶜdivᵥ_uₕ = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0)) ⊗ C12(FT(0), FT(0))), bottom = Operators.SetValue(sfc_conditions.ρ_flux_uₕ[colidx]), ) @. Yₜ.c.uₕ[colidx] -= - ᶜdivᵥ_uₕ(-(ᶠρaK_u[colidx] * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx] + ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx] end - # TODO: Add momentum mass flux - # TODO: Add tracer flux return nothing @@ -128,15 +132,15 @@ function edmfx_sgs_flux_tendency!( FT = Spaces.undertype(axes(Y.c)) n = n_mass_flux_subdomains(turbconv_model) (; edmfx_upwinding, sfc_conditions) = p - (; ᶠu³, ᶜh_tot, ᶜspecific) = p + (; ᶠu³, ᶜu, ᶜh_tot, ᶜspecific) = p (; ᶜρaʲs, ᶠu³ʲs, ᶜh_totʲs, ᶜq_totʲs) = p - (; ᶜh_tot⁰, ᶜK_u, ᶜK_h) = p + (; ᶜK_u, ᶜK_h) = p (; dt) = p.simulation ᶜJ = Fields.local_geometry_field(Y.c).J ᶠgradᵥ = Operators.GradientC2F() if p.atmos.edmfx_sgs_flux - # mass flux + # energy mass flux # TODO: check if there is contribution from the environment ᶠu³_diff_colidx = p.ᶠtemp_CT3[colidx] ᶜh_tot_diff_colidx = ᶜq_tot_diff_colidx = p.ᶜtemp_scalar[colidx] @@ -154,7 +158,7 @@ function edmfx_sgs_flux_tendency!( ) end - # diffusive flux + # energy diffusive flux ᶠρaK_h = p.ᶠtemp_scalar @. ᶠρaK_h[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_h[colidx]) @@ -163,10 +167,10 @@ function edmfx_sgs_flux_tendency!( bottom = Operators.SetValue(sfc_conditions.ρ_flux_h_tot[colidx]), ) @. Yₜ.c.ρe_tot[colidx] -= - ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot⁰[colidx]))) + ᶜdivᵥ_ρe_tot(-(ᶠρaK_h[colidx] * ᶠgradᵥ(ᶜh_tot[colidx]))) if !(p.atmos.moisture_model isa DryModel) - # mass flux + # specific humidity mass flux for j in 1:n @. ᶠu³_diff_colidx = ᶠu³ʲs.:($$j)[colidx] - ᶠu³[colidx] @. ᶜq_tot_diff_colidx = @@ -182,7 +186,7 @@ function edmfx_sgs_flux_tendency!( ) end - # diffusive flux + # specific humidity diffusive flux ᶜdivᵥ_ρq_tot = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0))), bottom = Operators.SetValue( @@ -194,19 +198,23 @@ function edmfx_sgs_flux_tendency!( ) end - # diffusive flux + # momentum diffusive flux ᶠρaK_u = p.ᶠtemp_scalar @. ᶠρaK_u[colidx] = ᶠinterp(Y.c.ρ[colidx]) * ᶠinterp(ᶜK_u[colidx]) + ᶠstrain_rate = p.ᶠtemp_UVWxUVW + compute_strain_rate_face!(ᶠstrain_rate[colidx], ᶜu[colidx]) + @. Yₜ.c.uₕ[colidx] -= C12( + ᶜdivᵥ(-(2 * ᶠρaK_u[colidx] * ᶠstrain_rate[colidx])) / Y.c.ρ[colidx], + ) + # apply boundary condition for momentum flux ᶜdivᵥ_uₕ = Operators.DivergenceF2C( top = Operators.SetValue(C3(FT(0)) ⊗ C12(FT(0), FT(0))), bottom = Operators.SetValue(sfc_conditions.ρ_flux_uₕ[colidx]), ) @. Yₜ.c.uₕ[colidx] -= - ᶜdivᵥ_uₕ(-(ᶠρaK_u[colidx] * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx] + ᶜdivᵥ_uₕ(-(FT(0) * ᶠgradᵥ(Y.c.uₕ[colidx]))) / Y.c.ρ[colidx] end - # TODO: Add momentum mass flux - # TODO: Add tracer flux return nothing diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl index 09a2e116049..befa7660e6a 100644 --- a/src/solver/type_getters.jl +++ b/src/solver/type_getters.jl @@ -533,10 +533,9 @@ function get_callbacks(parsed_args, simulation, atmos, params) ) end - return SciMLBase.CallbackSet(callbacks...) + return callbacks end - function get_cache( Y, parsed_args, @@ -607,6 +606,101 @@ function get_simulation(config::AtmosConfig, comms_ctx) return sim end +function get_diagnostics(parsed_args, atmos_model) + + # We either get the diagnostics section in the YAML file, or we return an empty list + # (which will result in an empty list being created by the map below) + yaml_diagnostics = get(parsed_args, "diagnostics", []) + + # ALLOWED_REDUCTIONS is the collection of reductions we support. The keys are the + # strings that have to be provided in the YAML file. The values are tuples with the + # function that has to be passed to reduction_time_func and the one that has to passed + # to pre_output_hook! + + # We make "nothing" a string so that we can accept also the word "nothing", in addition + # to the absence of the value + # + # NOTE: Everything has to be lowercase in ALLOWED_REDUCTIONS (so that we can match + # "max" and "Max") + ALLOWED_REDUCTIONS = Dict( + "nothing" => (nothing, nothing), # nothing is: just dump the variable + "max" => (max, nothing), + "min" => (min, nothing), + "average" => ((+), CAD.average_pre_output_hook!), + ) + + # The default writer is HDF5 + ALLOWED_WRITERS = Dict( + "nothing" => CAD.HDF5Writer(), + "h5" => CAD.HDF5Writer(), + "hdf5" => CAD.HDF5Writer(), + "nc" => CAD.NetCDFWriter(), + "netcdf" => CAD.NetCDFWriter(), + ) + + diagnostics = map(yaml_diagnostics) do yaml_diag + # Return "nothing" if "reduction_time" is not in the YAML block + # + # We also normalize everything to lowercase, so that can accept "max" but + # also "Max" + reduction_time_yaml = + lowercase(get(yaml_diag, "reduction_time", "nothing")) + + if !haskey(ALLOWED_REDUCTIONS, reduction_time_yaml) + error("reduction $reduction_time_yaml not implemented") + else + reduction_time_func, pre_output_hook! = + ALLOWED_REDUCTIONS[reduction_time_yaml] + end + + writer_ext = lowercase(get(yaml_diag, "writer", "nothing")) + + if !haskey(ALLOWED_WRITERS, writer_ext) + error("writer $writer_ext not implemented") + else + writer = ALLOWED_WRITERS[writer_ext] + end + + name = get(yaml_diag, "name", nothing) + + haskey(yaml_diag, "period") || + error("period keyword required for diagnostics") + + period_seconds = time_to_seconds(yaml_diag["period"]) + + if isnothing(name) + name = CAD.descriptive_short_name( + CAD.get_diagnostic_variable(yaml_diag["short_name"]), + period_seconds, + reduction_time_func, + pre_output_hook!, + ) + end + + if isnothing(reduction_time_func) + compute_every = period_seconds + else + compute_every = :timestep + end + + return CAD.ScheduledDiagnosticTime( + variable = CAD.get_diagnostic_variable(yaml_diag["short_name"]), + output_every = period_seconds, + compute_every = compute_every, + reduction_time_func = reduction_time_func, + pre_output_hook! = pre_output_hook!, + output_writer = writer, + output_short_name = name, + ) + end + + if parsed_args["output_default_diagnostics"] + return [CAD.default_diagnostics(atmos_model)..., diagnostics...] + else + return collect(diagnostics) + end +end + function args_integrator(parsed_args, Y, p, tspan, ode_algo, callback) (; atmos, simulation) = p (; dt) = simulation @@ -693,6 +787,10 @@ function get_integrator(config::AtmosConfig) atmos = get_atmos(config, params) numerics = get_numerics(config.parsed_args) simulation = get_simulation(config, config.comms_ctx) + if config.parsed_args["log_params"] + filepath = joinpath(simulation.output_dir, "$(job_id)_parameters.toml") + CP.log_parameter_information(config.toml_dict, filepath) + end initial_condition = get_initial_condition(config.parsed_args) surface_setup = get_surface_setup(config.parsed_args) @@ -742,8 +840,83 @@ function get_integrator(config::AtmosConfig) callback = get_callbacks(config.parsed_args, simulation, atmos, params) end @info "get_callbacks: $s" - @info "n_steps_per_cycle_per_cb: $(n_steps_per_cycle_per_cb(callback, simulation.dt))" - @info "n_steps_per_cycle: $(n_steps_per_cycle(callback, simulation.dt))" + + # Initialize diagnostics + s = @timed_str begin + diagnostics = get_diagnostics(config.parsed_args, atmos) + end + @info "initializing diagnostics: $s" + + # First, we convert all the ScheduledDiagnosticTime into ScheduledDiagnosticIteration, + # ensuring that there is consistency in the timestep and the periods and translating + # those periods that depended on the timestep + diagnostics_iterations = [ + CAD.ScheduledDiagnosticIterations(d, simulation.dt) for d in diagnostics + ] + + # For diagnostics that perform reductions, the storage is used for the values computed + # at each call. Reductions also save the accumulated value in diagnostic_accumulators. + diagnostic_storage = Dict() + diagnostic_accumulators = Dict() + diagnostic_counters = Dict() + + # NOTE: The diagnostics_callbacks are not called at the initial timestep + s = @timed_str begin + diagnostics_functions = CAD.get_callbacks_from_diagnostics( + diagnostics_iterations, + diagnostic_storage, + diagnostic_accumulators, + diagnostic_counters, + ) + end + @info "Prepared diagnostic callbacks: $s" + + # It would be nice to just pass the callbacks to the integrator. However, this leads to + # a significant increase in compile time for reasons that are not known. For this + # reason, we only add one callback to the integrator, and this function takes care of + # executing the other callbacks. This single function is orchestrate_diagnostics + + function orchestrate_diagnostics(integrator) + diagnostics_to_be_run = + filter(d -> integrator.step % d.cbf.n == 0, diagnostics_functions) + + for diag_func in diagnostics_to_be_run + diag_func.f!(integrator) + end + end + + diagnostic_callbacks = + call_every_n_steps(orchestrate_diagnostics, skip_first = true) + + # We need to ensure the precomputed quantities are indeed precomputed + + # TODO: Remove this when we can assume that the precomputed_quantities are in sync with + # the state + sync_precomputed = call_every_n_steps( + (int) -> set_precomputed_quantities!(int.u, int.p, int.t), + ) + + # The generic constructor for SciMLBase.CallbackSet has to split callbacks into discrete + # and continuous. This is not hard, but can introduce significant latency. However, all + # the callbacks in ClimaAtmos are discrete_callbacks, so we directly pass this + # information to the constructor + continuous_callbacks = tuple() + discrete_callbacks = (callback..., sync_precomputed, diagnostic_callbacks) + + s = @timed_str begin + all_callbacks = + SciMLBase.CallbackSet(continuous_callbacks, discrete_callbacks) + end + @info "Prepared SciMLBase.CallbackSet callbacks: $s" + steps_cycle_non_diag = + n_steps_per_cycle_per_cb(all_callbacks, simulation.dt) + steps_cycle_diag = + n_steps_per_cycle_per_cb_diagnostic(diagnostics_functions) + steps_cycle = lcm([steps_cycle_non_diag..., steps_cycle_diag...]) + @info "n_steps_per_cycle_per_cb (non diagnostics): $steps_cycle_non_diag" + @info "n_steps_per_cycle_per_cb_diagnostic: $steps_cycle_diag" + @info "n_steps_per_cycle (non diagnostics): $steps_cycle" + tspan = (t_start, simulation.t_end) s = @timed_str begin integrator_args, integrator_kwargs = args_integrator( @@ -752,7 +925,7 @@ function get_integrator(config::AtmosConfig) p, tspan, ode_algo, - callback, + all_callbacks, ) end @@ -760,5 +933,35 @@ function get_integrator(config::AtmosConfig) integrator = SciMLBase.init(integrator_args...; integrator_kwargs...) end @info "init integrator: $s" + + s = @timed_str begin + for diag in diagnostics_iterations + variable = diag.variable + try + # The first time we call compute! we use its return value. All + # the subsequent times (in the callbacks), we will write the + # result in place + diagnostic_storage[diag] = variable.compute!( + nothing, + integrator.u, + integrator.p, + integrator.t, + ) + diagnostic_counters[diag] = 1 + # If it is not a reduction, call the output writer as well + if isnothing(diag.reduction_time_func) + diag.output_writer(diagnostic_storage[diag], diag, integrator) + else + # Add to the accumulator + diagnostic_accumulators[diag] = + copy(diagnostic_storage[diag]) + end + catch e + error("Could not compute diagnostic $(variable.long_name): $e") + end + end + end + @info "Init diagnostics: $s" + return integrator end diff --git a/src/solver/types.jl b/src/solver/types.jl index c0137a7bf43..598cefa675b 100644 --- a/src/solver/types.jl +++ b/src/solver/types.jl @@ -466,3 +466,11 @@ function AtmosConfig(; return AtmosConfig{FT, TD, PA, C}(toml_dict, config, comms_ctx) end Base.eltype(::AtmosConfig{FT}) where {FT} = FT + +# In order to specify C2F operator boundary conditions with 0 instead of FT(0), +# we need to tell ClimaCore how to convert AxisTensor components from Int to FT. +# TODO: Move this monkey patch to ClimaCore in the next release. +using ClimaCore.Geometry: AxisTensor, components +AxisTensor{T′, N, A, S′}(a::AxisTensor{T, N, A, S}) where {T, N, A, S, T′, S′} = + AxisTensor(axes(a), S′(components(a))) +Base.convert(::Type{AT}, a::AxisTensor) where {AT <: AxisTensor} = AT(a) diff --git a/src/solver/yaml_helper.jl b/src/solver/yaml_helper.jl index 399532a30ce..800ac198583 100644 --- a/src/solver/yaml_helper.jl +++ b/src/solver/yaml_helper.jl @@ -74,6 +74,13 @@ function override_default_config(config_dict::AbstractDict;) v = config_dict[k] config[k] = isnothing(default_config[k]) ? v : default_type(v) end + + # The "diagnostics" entry is a more complex type that doesn't fit the schema described in + # the previous lines. So, we manually add it. + if haskey(config_dict, "diagnostics") + config["diagnostics"] = config_dict["diagnostics"] + end + return config end diff --git a/src/utils/abbreviations.jl b/src/utils/abbreviations.jl index a374cbd3960..9db0375ccea 100644 --- a/src/utils/abbreviations.jl +++ b/src/utils/abbreviations.jl @@ -75,11 +75,3 @@ const ᶜdivᵥ_stencil = Operators.Operator2Stencil(ᶜdivᵥ) const ᶜadvdivᵥ_stencil = Operators.Operator2Stencil(ᶜadvdivᵥ) const ᶠinterp_stencil = Operators.Operator2Stencil(ᶠinterp) const ᶠgradᵥ_stencil = Operators.Operator2Stencil(ᶠgradᵥ) - -# In order to specify C2F operator boundary conditions with 0 instead of FT(0), -# we need to tell ClimaCore how to convert AxisTensor components from Int to FT. -# TODO: Move this monkey patch to ClimaCore in the next release. -using ClimaCore.Geometry: AxisTensor, components -AxisTensor{T′, N, A, S′}(a::AxisTensor{T, N, A, S}) where {T, N, A, S, T′, S′} = - AxisTensor(axes(a), S′(components(a))) -Base.convert(::Type{AT}, a::AxisTensor) where {AT <: AxisTensor} = AT(a) diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl index a6eb38ed21b..4d0a5b6c7f2 100644 --- a/src/utils/utilities.jl +++ b/src/utils/utilities.jl @@ -79,12 +79,12 @@ compute_kinetic!(κ::Fields.Field, Y::Fields.FieldVector) = compute_kinetic!(κ, Y.c.uₕ, Y.f.u₃) """ - compute_strain_rate!(ϵ::Field, u::Field) + compute_strain_rate_center!(ϵ::Field, u::Field) Compute the strain_rate at cell centers, storing in `ϵ` from velocity at cell faces. """ -function compute_strain_rate!(ϵ::Fields.Field, u::Fields.Field) +function compute_strain_rate_center!(ϵ::Fields.Field, u::Fields.Field) @assert eltype(u) <: C123 axis_uvw = Geometry.UVWAxis() @. ϵ = @@ -94,6 +94,28 @@ function compute_strain_rate!(ϵ::Fields.Field, u::Fields.Field) ) / 2 end +""" + compute_strain_rate_face!(ϵ::Field, u::Field) + +Compute the strain_rate at cell faces, storing in `ϵ` from +velocity at cell centers. +""" +function compute_strain_rate_face!(ϵ::Fields.Field, u::Fields.Field) + @assert eltype(u) <: C123 + ∇ᵥuvw_boundary = + Geometry.outer(Geometry.WVector(0), Geometry.UVWVector(0, 0, 0)) + ᶠgradᵥ = Operators.GradientC2F( + bottom = Operators.SetGradient(∇ᵥuvw_boundary), + top = Operators.SetGradient(∇ᵥuvw_boundary), + ) + axis_uvw = Geometry.UVWAxis() + @. ϵ = + ( + Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u))) + + adjoint(Geometry.project((axis_uvw,), ᶠgradᵥ(UVW(u)))) + ) / 2 +end + """ g³³_field(field) diff --git a/toml/edmfx_box_advection.toml b/toml/edmfx_box_advection.toml index d39ecbfaa8f..fbffbcba315 100644 --- a/toml/edmfx_box_advection.toml +++ b/toml/edmfx_box_advection.toml @@ -5,6 +5,6 @@ type = "float" [EDMF_min_area] alias = "min_area" -value = 1.0e-2 +value = 1.0e-3 type = "float" description = "Minimum area fraction per updraft. Parameter not described in the literature." diff --git a/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml b/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml new file mode 100644 index 00000000000..91de57be7f7 --- /dev/null +++ b/toml/sphere_baroclinic_wave_rhoe_equilmoist_expvdiff.toml @@ -0,0 +1,10 @@ +[C_H] +alias = "C_H" +value = 0.0 +type = "float" + +[C_E] +alias = "C_E" +value = 1 +type = "float" +