Merge #399

399: Enhance AMIP and NCEP paperplots r=LenkaNovak a=LenkaNovak



Co-authored-by: LenkaNovak <[email protected]>

Author: bors[bot]

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -353,7 +353,7 @@ dates = (; date = [date], date0 = [date0], date1 = [Dates.firstdayofmonth(date0)
 User can write custom diagnostics in the `user_diagnostics.jl`.
 =#
 monthly_3d_diags = init_diagnostics(
-    (:T, :u, :q_tot),
+    (:T, :u, :q_tot, :q_liq_ice),
     atmos_sim.domain.center_space;
     save = Monthly(),
     operations = (; accumulate = TimeMean([Int(0)])),
@@ -362,7 +362,7 @@ monthly_3d_diags = init_diagnostics(
 )
 
 monthly_2d_diags = init_diagnostics(
-    (:precipitation, :toa, :T_sfc),
+    (:precipitation_rate, :toa_fluxes, :T_sfc, :tubulent_energy_fluxes),
     boundary_space;
     save = Monthly(),
     operations = (; accumulate = TimeMean([Int(0)])),
@@ -610,20 +610,24 @@ if ClimaComms.iamroot(comms_ctx)
             T = (:regrid, :zonal_mean),
             u = (:regrid, :zonal_mean),
             q_tot = (:regrid, :zonal_mean),
-            toa = (:regrid, :horizontal_slice),
-            precipitation = (:regrid, :horizontal_slice),
+            toa_fluxes = (:regrid, :horizontal_slice),
+            precipitation_rate = (:regrid, :horizontal_slice),
             T_sfc = (:regrid, :horizontal_slice),
+            tubulent_energy_fluxes = (:regrid, :horizontal_slice),
+            q_liq_ice = (:regrid, :zonal_mean),
         )
 
         plot_spec = (;
             T = (; clims = (190, 320), units = "K"),
             u = (; clims = (-50, 50), units = "m/s"),
-            q_tot = (; clims = (0, 50), units = "g/kg"),
-            toa = (; clims = (-250, 210), units = "W/m^2"),
-            precipitation = (clims = (0, 1e-6), units = "kg/m^2/s"),
+            q_tot = (; clims = (0, 30), units = "g/kg"),
+            toa_fluxes = (; clims = (-250, 250), units = "W/m^2"),
+            precipitation_rate = (clims = (0, 1e-4), units = "kg/m^2/s"),
             T_sfc = (clims = (225, 310), units = "K"),
+            tubulent_energy_fluxes = (; clims = (-250, 250), units = "W/m^2"),
+            q_liq_ice = (; clims = (0, 10), units = "g/kg"),
         )
-        amip_paperplots(
+        amip_data = amip_paperplots(
             post_spec,
             plot_spec,
             COUPLER_OUTPUT_DIR,
@@ -638,12 +642,13 @@ if ClimaComms.iamroot(comms_ctx)
             T = (:zonal_mean,),
             u = (:zonal_mean,),
             q_tot = (:zonal_mean,),
-            toa = (:horizontal_slice,),
-            precipitation = (:horizontal_slice,),
+            toa_fluxes = (:horizontal_slice,),
+            precipitation_rate = (:horizontal_slice,),
             T_sfc = (:horizontal_slice,),
+            tubulent_energy_fluxes = (:horizontal_slice,),
         )
         ncep_plot_spec = plot_spec
-        ncep_paperplots(
+        ncep_data = ncep_paperplots(
             ncep_post_spec,
             ncep_plot_spec,
             COUPLER_OUTPUT_DIR,

experiments/AMIP/modular/user_io/amip_visualizer.jl

@@ -19,19 +19,22 @@ function amip_paperplots(
     output_dir = ".",
     files_root = ".hdf5",
     fig_name = "amip_paperplots",
+    nlat = 180,
+    nlon = 360,
 )
 
     diags_names = propertynames(post_spec)
 
     all_plots = []
+    all_data = (;)
     for name in diags_names
         @info name
 
         # extract data
         diag_data = read_latest_model_data(name, files_dir, files_root)
 
         # postprocess
-        post_data = postprocess(name, diag_data, getproperty(post_spec, name))
+        post_data = postprocess(name, diag_data, getproperty(post_spec, name), nlat = nlat, nlon = nlon)
         post_data.data[1] = sum(post_data.data) == 0 ? post_data.data[1] + eps() : post_data.data[1] # avoids InexactError
 
         # create individual plots
@@ -42,21 +45,25 @@ function amip_paperplots(
         )
 
         push!(all_plots, p)
+
+        # create a named tuple with data
+        data = post_data.data
+        all_data = merge(all_data, [name => data])
     end
 
     # combine plots and save figure
     save_fig = Plots.plot(
         all_plots...,
-        size = (1500, 800),
-        right_margin = 12Plots.mm,
-        left_margin = 12Plots.mm,
-        bottom_margin = 12Plots.mm,
-        top_margin = 12Plots.mm,
+        size = (1500, 1200),
+        right_margin = 3Plots.mm,
+        left_margin = 3Plots.mm,
+        bottom_margin = 3Plots.mm,
+        top_margin = 3Plots.mm,
     )
 
     Plots.png(save_fig, joinpath(output_dir, fig_name * ".png"))
 
-    return all_plots
+    return all_data
 end
 
 """

experiments/AMIP/modular/user_io/ncep_visualizer.jl

@@ -13,9 +13,9 @@ include("plot_helper.jl")
         fig_name = "ncep_paperplots",
     )
 
-Coordinates the postprocessing and plotting of sample fields (specified in `post_spec`) 
+Coordinates the postprocessing and plotting of sample fields (specified in `post_spec`)
 of a particular monthly mean dataset (specified by `month_date`). Any plot NCEP- specific
-customization should be done here. 
+customization should be done here.
 """
 function ncep_paperplots(
     post_spec::NamedTuple,
@@ -35,6 +35,7 @@ function ncep_paperplots(
     diags_vnames = propertynames(post_spec)
 
     all_plots = []
+    all_data = (;)
     for vname in diags_vnames
         @info vname
 
@@ -50,21 +51,25 @@ function ncep_paperplots(
         p = plot(post_data, zmd_params = zonal_mean_params, hsd_params = (; getproperty(plot_spec, vname)...))
 
         push!(all_plots, p)
+
+        # create a named tuple with data
+        data = post_data.data
+        all_data = merge(all_data, [vname => data])
     end
 
     # combine plots and save figure
     save_fig = Plots.plot(
         all_plots...,
-        size = (1500, 800),
-        right_margin = 12Plots.mm,
-        left_margin = 12Plots.mm,
-        bottom_margin = 12Plots.mm,
-        top_margin = 12Plots.mm,
+        size = (1500, 1200),
+        right_margin = 3Plots.mm,
+        left_margin = 3Plots.mm,
+        bottom_margin = 3Plots.mm,
+        top_margin = 3Plots.mm,
     )
 
     Plots.png(save_fig, joinpath(output_dir, fig_name * ".png"))
 
-    return all_plots
+    return all_data
 end
 
 abstract type DataSource end
@@ -116,7 +121,7 @@ function get_var(data_source::NCEPMonthlyDataSource, ::Val{:q_tot})
     download_read_nc(data_source, https, ncep_vname)
 end
 
-function get_var(data_source::NCEPMonthlyDataSource, ::Val{:toa})
+function get_var(data_source::NCEPMonthlyDataSource, ::Val{:toa_fluxes})
     https_root = "https://downloads.psl.noaa.gov/Datasets/ncep.reanalysis2/Monthlies/gaussian_grid/"
     https_suffix = ".ntat.mon.mean.nc"
 
@@ -132,7 +137,7 @@ function get_var(data_source::NCEPMonthlyDataSource, ::Val{:toa})
     return (data_uswrf .- data_dswrf .+ data_ulwrf, coords)
 end
 
-function get_var(data_source::NCEPMonthlyDataSource, ::Val{:precipitation})
+function get_var(data_source::NCEPMonthlyDataSource, ::Val{:precipitation_rate})
     https = "https://downloads.psl.noaa.gov/Datasets/ncep.reanalysis2/Monthlies/gaussian_grid/prate.sfc.mon.mean.nc"
     ncep_vname = "prate"
     download_read_nc(data_source, https, ncep_vname)
@@ -144,3 +149,13 @@ function get_var(data_source::NCEPMonthlyDataSource, ::Val{:T_sfc})
     t_celsius, coords = download_read_nc(data_source, https, ncep_vname)
     return (t_celsius .+ 273.15, coords)
 end
+
+function get_var(data_source::NCEPMonthlyDataSource, ::Val{:tubulent_energy_fluxes})
+    https = "https://downloads.psl.noaa.gov/Datasets/ncep.reanalysis2/Monthlies/gaussian_grid/lhtfl.sfc.mon.mean.nc"
+    ncep_vname = "lhtfl"
+    lhtfl, coords = download_read_nc(data_source, https, ncep_vname)
+    https = "https://downloads.psl.noaa.gov/Datasets/ncep.reanalysis2/Monthlies/gaussian_grid/shtfl.sfc.mon.mean.nc"
+    ncep_vname = "shtfl"
+    shtfl, coords = download_read_nc(data_source, https, ncep_vname)
+    return (shtfl .+ lhtfl, coords)
+end

experiments/AMIP/modular/user_io/plot_helper.jl

@@ -10,9 +10,9 @@ Coordinates plotting based on parsed data types.
 """
 function plot(post_data::DataPackage; zmd_params = (;), hsd_params = (;))
 
-    if post_data isa ZLatData
+    if post_data.tag isa ZLatData
         plot_params = zmd_params
-    elseif post_data isa LatLonData
+    elseif post_data.tag isa LatLonData
         plot_params = hsd_params
     else
         plot_params = (;)
@@ -45,7 +45,7 @@ function contourf(
     units = " ",
 )
     clims = clims == nothing ? extrema(p.data) : clims
-    plot_p = Plots.contourf(
+    Plots.contourf(
         p.coords.lat,
         p.coords.lev,
         p.data',

experiments/AMIP/modular/user_io/user_diagnostics.jl

@@ -27,17 +27,30 @@ get_var(cs::CoupledSimulation, ::Val{:u}) =
 """
     get_var(cs::CoupledSimulation, ::Val{:q_tot})
 
-Total specific humidity (kg kg⁻¹).
+Total specific humidity (g kg⁻¹).
 """
 get_var(cs::CoupledSimulation, ::Val{:q_tot}) =
     cs.model_sims.atmos_sim.integrator.u.c.ρq_tot ./ cs.model_sims.atmos_sim.integrator.u.c.ρ .* float_type(cs)(1000)
 
+
+"""
+    get_var(cs::CoupledSimulation, ::Val{:q_liq_ice})
+
+Cloud specific humidity (g kg⁻¹).
 """
-    get_var(cs::CoupledSimulation, ::Val{:toa})
+function get_var(cs::CoupledSimulation, ::Val{:q_liq_ice})
+    p = cs.model_sims.atmos_sim.integrator.p
+    (; ᶜts, params) = p
+    thermo_params = CAP.thermodynamics_params(params)
+    TD.liquid_specific_humidity.(thermo_params, ᶜts) .* float_type(cs)(1000)
+end
+
+"""
+    get_var(cs::CoupledSimulation, ::Val{:toa_fluxes})
 
 Top of the atmosphere radiation fluxes (W m⁻²).
 """
-function get_var(cs::CoupledSimulation, ::Val{:toa})
+function get_var(cs::CoupledSimulation, ::Val{:toa_fluxes})
     atmos_sim = cs.model_sims.atmos_sim
     face_space = axes(atmos_sim.integrator.u.f)
     z = parent(Fields.coordinate_field(face_space).z)
@@ -66,11 +79,11 @@ function get_var(cs::CoupledSimulation, ::Val{:toa})
 end
 
 """
-    get_var(cs::CoupledSimulation, ::Val{:precipitation})
+    get_var(cs::CoupledSimulation, ::Val{:precipitation_rate})
 
-Precipitation (m m⁻²).
+Precipitation rate (Kg m⁻² s⁻¹).
 """
-get_var(cs::CoupledSimulation, ::Val{:precipitation}) =
+get_var(cs::CoupledSimulation, ::Val{:precipitation_rate}) =
     .-swap_space!(
         zeros(cs.boundary_space),
         cs.model_sims.atmos_sim.integrator.p.col_integrated_rain .+
@@ -85,4 +98,12 @@ Combined surface temperature (K).
 """
 get_var(cs::CoupledSimulation, ::Val{:T_sfc}) = swap_space!(zeros(cs.boundary_space), cs.fields.T_S)
 
+"""
+    get_var(cs::CoupledSimulation, ::Val{:tubulent_energy_fluxes})
+
+Combined aerodynamic turbulent energy surface fluxes (W m⁻²).
+"""
+get_var(cs::CoupledSimulation, ::Val{:tubulent_energy_fluxes}) =
+    swap_space!(zeros(cs.boundary_space), cs.fields.F_turb_energy)
+
 # land diagnotics

perf/flame.jl

@@ -18,7 +18,7 @@ run_name_list =
 run_name = run_name_list[parse(Int, ARGS[2])]
 allocs_limit = Dict()
 allocs_limit["perf_default_modular_unthreaded"] = 8638304
-allocs_limit["perf_coarse_single_modular"] = 12671552
+allocs_limit["perf_coarse_single_modular"] = 18280800
 allocs_limit["perf_target_amip_n32_shortrun"] = 172134848
 
 # number of time steps used for profiling