Merge #829

829: fix allocations for snow and rain output r=trontrytel a=trontrytel

# PULL REQUEST

This PR fixes some allocations related to outputting rain and snow. See the discussion here: #816

`@charleskawczynski` - Is that what you had in mind?

## Purpose and Content
(Clearly and concisely state the purpose of this PR)

## Benefits and Risks
(State concisely the benefits to be derived from and the risks associated with this PR)

## Linked Issues
(Provide references to any link issues. Use closes #issuenum to automatically close an open issue)
- Fixes #
- Closes #

## PR Checklist
- [ ] This PR has a corresponding issue OR is linked to an SDI.
- [ ] I have followed CliMA's codebase [contribution](https://clima.github.io/ClimateMachine.jl/latest/Contributing/) and [style](https://clima.github.io/ClimateMachine.jl/latest/DevDocs/CodeStyle/) guidelines OR N/A.
- [ ] I have followed CliMA's [documentation policy](https://github.com/CliMA/policies/wiki/Documentation-Policy).
- [ ] I have checked all issues and PRs and I certify that this PR does not duplicate an open PR.
- [ ] I linted my code on my local machine prior to submission OR N/A.
- [ ] Unit tests are included OR N/A.
- [ ] Code used in an integration test OR N/A.
- [ ] All tests ran successfully on my local machine OR N/A.
- [ ] All classes, modules, and function contain docstrings OR N/A.
- [ ] Documentation has been added/updated OR N/A.


Co-authored-by: Anna Jaruga <[email protected]>

.buildkite/pipeline.yml

@@ -333,7 +333,7 @@ steps:
           slurm_mem: 20GB
 
       - label: ":balloon: Single column EDMF TRMM_LBA"
-        command: "julia --color=yes --project=examples examples/hybrid/driver.jl --config column --FLOAT_TYPE Float64 --hyperdiff false --moist equil --turbconv edmf --turbconv_case TRMM_LBA --split_ode false --ode_algo ODE.Euler --dt_save_to_sol 5mins --z_elem 82 --z_stretch false --z_max 16400 --job_id edmf_trmm --dt 1secs --t_end 6hours --regression_test true --anelastic_dycore true --apply_limiter false"
+        command: "julia --color=yes --project=examples examples/hybrid/driver.jl --config column --FLOAT_TYPE Float64 --hyperdiff false --moist equil --turbconv edmf --turbconv_case TRMM_LBA --split_ode false --ode_algo ODE.Euler --dt_save_to_sol 5mins --dt_save_to_disk 3hours --z_elem 82 --z_stretch false --z_max 16400 --job_id edmf_trmm --dt 1secs --t_end 6hours --regression_test true --anelastic_dycore true --apply_limiter false"
         artifact_paths: "edmf_trmm/*"
         agents:
           slurm_mem: 20GB

examples/hybrid/callbacks.jl

@@ -155,8 +155,8 @@ function save_to_disk_func(integrator)
     (; output_dir) = p.simulation
     Y = u
 
-    if :ρq_tot in propertynames(Y.c)
-        (; ᶜts, ᶜp, ᶜS_ρq_tot, params, ᶜK) = p
+    if :ᶜS_ρq_tot in propertynames(Y.c)
+        (; ᶜts, ᶜp, ᶜS_ρq_tot, ᶜ3d_rain, ᶜ3d_snow, params, ᶜK) = p
     else
         (; ᶜts, ᶜp, params, ᶜK) = p
     end
@@ -187,38 +187,45 @@ function save_to_disk_func(integrator)
         vorticity = ᶜvort,
     )
 
-    # cloudwater (liquid and ice), watervapor, precipitation, and RH for moist simulation
+    # cloudwater (liquid and ice), watervapor and RH for moist simulation
     if :ρq_tot in propertynames(Y.c)
         ᶜq = @. TD.PhasePartition(thermo_params, ᶜts)
         ᶜcloud_liquid = @. ᶜq.liq
         ᶜcloud_ice = @. ᶜq.ice
         ᶜwatervapor = @. TD.vapor_specific_humidity(ᶜq)
         ᶜRH = @. TD.relative_humidity(thermo_params, ᶜts)
 
-        # precipitation
-        @. ᶜS_ρq_tot =
-            Y.c.ρ * CM.Microphysics0M.remove_precipitation(
-                cm_params,
-                TD.PhasePartition(thermo_params, ᶜts),
-            )
-
-        # rain vs snow
-        ᶜ3d_rain = @. ifelse(ᶜT >= FT(273.15), ᶜS_ρq_tot, FT(0))
-        ᶜ3d_snow = @. ifelse(ᶜT < FT(273.15), ᶜS_ρq_tot, FT(0))
-        col_integrated_rain =
-            vertical∫_col(ᶜ3d_rain) ./ FT(CAP.ρ_cloud_liq(params))
-        col_integrated_snow =
-            vertical∫_col(ᶜ3d_snow) ./ FT(CAP.ρ_cloud_liq(params))
-
         moist_diagnostic = (;
             cloud_liquid = ᶜcloud_liquid,
             cloud_ice = ᶜcloud_ice,
             water_vapor = ᶜwatervapor,
-            precipitation_removal = ᶜS_ρq_tot,
-            column_integrated_rain = col_integrated_rain,
-            column_integrated_snow = col_integrated_snow,
             relative_humidity = ᶜRH,
         )
+        # precipitation
+        if :ᶜS_ρq_tot in propertynames(Y.c)
+
+            @. ᶜS_ρq_tot =
+                Y.c.ρ * CM.Microphysics0M.remove_precipitation(
+                    cm_params,
+                    TD.PhasePartition(thermo_params, ᶜts),
+                )
+
+            # rain vs snow
+            @. ᶜ3d_rain = ifelse(ᶜT >= FT(273.15), ᶜS_ρq_tot, FT(0))
+            @. ᶜ3d_snow = ifelse(ᶜT < FT(273.15), ᶜS_ρq_tot, FT(0))
+
+            col_integrated_rain =
+                vertical∫_col(ᶜ3d_rain) ./ FT(CAP.ρ_cloud_liq(params))
+            col_integrated_snow =
+                vertical∫_col(ᶜ3d_snow) ./ FT(CAP.ρ_cloud_liq(params))
+
+            moist_diagnostics = (
+                moist_diagnostics...,
+                precipitation_removal = ᶜS_ρq_tot,
+                column_integrated_rain = col_integrated_rain,
+                column_integrated_snow = col_integrated_snow,
+            )
+        end
     else
         moist_diagnostic = NamedTuple()
     end

examples/hybrid/sphere/baroclinic_wave_utilities.jl

@@ -151,6 +151,8 @@ microphysics_cache(Y, ::Nothing) = NamedTuple()
 microphysics_cache(Y, ::Microphysics0Moment) = (
     ᶜS_ρq_tot = similar(Y.c, FT),
     ᶜλ = similar(Y.c, FT),
+    ᶜ3d_rain = similar(Y.c, FT),
+    ᶜ3d_snow = similar(Y.c, FT),
     col_integrated_rain = similar(ClimaCore.Fields.level(Y.c.ρ, 1), FT),
     col_integrated_snow = similar(ClimaCore.Fields.level(Y.c.ρ, 1), FT),
 )

examples/hybrid/staggered_nonhydrostatic_model.jl

@@ -132,8 +132,6 @@ function default_cache(Y, params, spaces, numerics, simulation)
         ᶜts = similar(Y.c, ts_type),
         ᶜp = similar(Y.c, FT),
         ᶜT = similar(Y.c, FT),
-        ᶜ3d_rain = similar(Y.c, FT),
-        ᶜ3d_snow = similar(Y.c, FT),
         ᶜω³ = similar(Y.c, Geometry.Contravariant3Vector{FT}),
         ᶠω¹² = similar(Y.f, Geometry.Contravariant12Vector{FT}),
         ᶠu¹² = similar(Y.f, Geometry.Contravariant12Vector{FT}),

perf/flame.jl

@@ -57,7 +57,7 @@ end
 using Test
 allocs_limit = Dict()
 allocs_limit["flame_perf_target_rhoe"] = 10357712
-allocs_limit["flame_perf_target_rhoe_threaded"] = 90908208
+allocs_limit["flame_perf_target_rhoe_threaded"] = 90909168
 
 # Threaded allocations are not deterministic, so let's add a buffer
 # TODO: remove buffer, and threaded tests, when