Add warn_allocations_diagnostics

config/default_configs/default_config.yml

@@ -261,3 +261,6 @@ output_default_diagnostics:
 netcdf_interpolate_z_over_msl:
   help: "Interpolate diagnostics in the NetCDF files in such a way that `z` is the elevation from the sea level (as opposed to the elevation from the surface)"
   value: false
+warn_allocations_diagnostics:
+  help: "When true, a dry-run for all the diagnostics is perform to check whether the functions allocate additional memory (which reduces performances)"
+  value: false

config/default_configs/default_perf.yml

@@ -10,3 +10,4 @@ forcing: "held_suarez"
 precip_model: "0M"
 dt_save_to_sol: "Inf"
 rad: "allskywithclear"
+warn_allocations_diagnostics: true

src/solver/type_getters.jl

@@ -925,6 +925,21 @@ function get_simulation(config::AtmosConfig)
     end
     @info "Init diagnostics: $s"
 
+    if config.parsed_args["warn_allocations_diagnostics"]
+        for diag in diagnostics_iterations
+            # We write over the storage space we have already prepared (and filled) before
+            allocs = @allocated diag.variable.compute!(
+                diagnostic_storage[diag],
+                integrator.u,
+                integrator.p,
+                integrator.t,
+            )
+            if allocs > 10 * 1024
+                @warn "Diagnostics $(diag.output_short_name) allocates $allocs bytes"
+            end
+        end
+    end
+
     return AtmosSimulation(
         sim_info.job_id,
         sim_info.output_dir,