E3SM-Project
EAMxx: requesting a diagnostic field inside a diagnostic
The idea of our diagnostics is that we can query them inside each other, and we can do composable and fancy stuff online. It took me a while to realize this wasn't the case, and the diagnostic field dz was simply always zero. It's a bit bizarre, because clearly it is there as add_field<Required>("dz"...) didn't error, and get_field_in("dz") didn't error.
There might be something I am missing or something I am assuming... but there's a larger chance that we have bug :)
It should be fine. The IO layer should detect that diag A depends on diag B, and create diag B as well. Then, at runtime, it should evaluate B before A
Are you absolutely sure you did not accidentally init one of the deps of dz to 0?
Yeah, I was testing with release/tests/multi-process/physics_only/shoc_cld_spa_p3_rrtmgp with the collapsed code
example python code
import numpy as np
import xarray as xr
from matplotlib import pyplot as plt
ds = xr.open_dataset("/path/to/tests/release/tests/multi-process/physics_only/shoc_cld_spa_p3_rrtmgp/shoc_cld_spa_p3_rrtmgp_output.INSTANT.nsteps_x5.np1.2021-10-12-54000.nc")
def test_array_equal(a, b):
return np.abs(a[0] - b[0]).max()
print("T_mid_horiz_avg: ", test_array_equal(ds.T_mid_horiz_avg, ds.T_mid.weighted(ds.area).mean(("ncol"))).values)
print("T_mid_vert_avg: ", test_array_equal(ds.T_mid_vert_avg, ds.T_mid.mean(("lev"))).values)
print("T_mid_vert_sum: ", test_array_equal(ds.T_mid_vert_sum, ds.T_mid.sum(("lev"))).values)
print("T_mid_vert_sum_dz_weighted: ", test_array_equal(ds.T_mid_vert_sum_dz_weighted, ds.T_mid.weighted(ds.dz).sum(("lev"))).values)
print("T_mid_vert_sum_dp_weighted: ", test_array_equal(ds.T_mid_vert_sum_dp_weighted, ds.T_mid.weighted(ds.pseudo_density/9.8).sum(("lev"))).values)
print("T_mid_vert_avg_dz_weighted: ", test_array_equal(ds.T_mid_vert_avg_dz_weighted, ds.T_mid.weighted(ds.dz).mean(("lev"))).values)
print("T_mid_vert_avg_dp_weighted: ", test_array_equal(ds.T_mid_vert_avg_dp_weighted, ds.T_mid.weighted(ds.pseudo_density/9.8).mean(("lev"))).values)
ds.T_mid_vert_avg_dz_weighted.plot(label="T_mid_vert_avg_dz_weighted")
ds.T_mid.weighted(ds.dz).mean(("lev")).plot(label="ds.T_mid.weighted(ds.dz).mean(("lev"))")
plt.legend()
# etc.
@mahf708 I don't know if things changed, but using dz as a field (rather than computing it) inside VertContraction seems to work. I get
// T_mid_vert_avg_dz_weighted(1, 0-217)
241.3825, 241.8377, 242.1158, 241.7968, 240.6747, 240.9166, 241.32,
241.075, 239.7008, 239.9491, 240.2971, 240.134, 239.2592, 239.4428,
239.6947, 239.6984, 241.2755, 240.4428, 240.1321, 240.7285, 240.071,
239.9094, 239.9198, 239.868, 239.895, 239.6843, 239.7995, 239.8736,
238.5275, 238.674, 238.761, 238.8838, 236.2203, 236.2287, 235.674,
235.5095, 233.6617, 232.3857, 230.3823, 229.8725, 239.0149, 239.2543,
239.2695, 235.4251, 236.1464, 236.2182, 230.2704, 232.4637, 233.67,
...
which seems reasonable.. Maybe we can switch back to use dz and see if things work?
Thanks for double checking! It's time to switch back yes! Let's keep this issue open until I switch back for the few places I am doing the longer calc
Tangentially related: I just tested requesting z_mid inside of nudging via add_field and got all zeros. Patch below on top of very recent codebase for the lengthy z_mid calc needed otherwise...
patch
From 8ab9dd9a1de572b2113f1326878febd7bb962013 Mon Sep 17 00:00:00 2001
From: mahf708 <[email protected]>
Date: Wed, 12 Nov 2025 06:10:51 -0800
Subject: [PATCH] give users option to not nudge in the shoc-defined PBL
---
.../cime_config/namelist_defaults_eamxx.xml | 6 ++
.../eamxx_nudging_process_interface.cpp | 78 ++++++++++++++++++-
.../eamxx_nudging_process_interface.hpp | 2 +
3 files changed, 85 insertions(+), 1 deletion(-)
diff --git a/components/eamxx/cime_config/namelist_defaults_eamxx.xml b/components/eamxx/cime_config/namelist_defaults_eamxx.xml
index 1d432bfe4306..3226066f62ad 100644
--- a/components/eamxx/cime_config/namelist_defaults_eamxx.xml
+++ b/components/eamxx/cime_config/namelist_defaults_eamxx.xml
@@ -449,6 +449,12 @@ be lost if SCREAM_HACK_XML is not enabled.
>
0.0
</nudging_refine_remap_vert_cutoff>
+ <nudging_pbl_height_cutoff
+ type="logical"
+ doc="Whether to use PBL height from SHOC as cutoff for nudging. If true, nudging will be turned off in the PBL (Z .le. PBLH)."
+ >
+ false
+ </nudging_pbl_height_cutoff>
</nudging>
<!-- ML correction -->
diff --git a/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.cpp b/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.cpp
index 1f0bfe0b349e..8272df2cb1ca 100644
--- a/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.cpp
+++ b/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.cpp
@@ -1,6 +1,8 @@
#include "eamxx_nudging_process_interface.hpp"
#include "share/util/eamxx_universal_constants.hpp"
+#include "share/physics/physics_constants.hpp"
+#include "share/physics/eamxx_common_physics_functions.hpp"
#include "share/remap/refining_remapper_p2p.hpp"
#include "share/util/eamxx_utils.hpp"
#include "share/scorpio_interface/eamxx_scorpio_interface.hpp"
@@ -27,6 +29,7 @@ Nudging::Nudging (const ekat::Comm& comm, const ekat::ParameterList& params)
"nudging_refine_remap_mapfile", "no-file-given");
m_refine_remap_vert_cutoff = m_params.get<Real>(
"nudging_refine_remap_vert_cutoff", 0.0);
+ m_pbl_height_cutoff = m_params.get<bool>("nudging_pbl_height_cutoff", false);
auto src_pres_type = m_params.get<std::string>("source_pressure_type","TIME_DEPENDENT_3D_PROFILE");
if (src_pres_type=="TIME_DEPENDENT_3D_PROFILE") {
m_src_pres_type = TIME_DEPENDENT_3D_PROFILE;
@@ -88,9 +91,22 @@ void Nudging::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
if (ekat::contains(m_fields_nudge,"U") or ekat::contains(m_fields_nudge,"V")) {
add_field<Updated>("horiz_winds", horiz_wind_layout, m/s, grid_name, ps);
}
+ if (ekat::contains(m_fields_nudge,"o3_volume_mix_ratio")) {
+ add_field<Required>("o3_volume_mix_ratio", scalar3d_layout_mid, mol/mol, grid_name, ps);
+ }
/* ----------------------- WARNING --------------------------------*/
+ // get pbl_height from shoc if requested
+ if (m_pbl_height_cutoff) {
+ add_field<Required>("pbl_height", m_grid->get_2d_scalar_layout(), m, grid_name);
+ // We need T, qv, p and pseudo_density to compute z_mid (AGL)
+ add_field<Required>("T_mid", scalar3d_layout_mid, K, grid_name, ps);
+ add_field<Required>("qv", scalar3d_layout_mid, kg/kg, grid_name, ps);
+ add_field<Required>("p_mid", scalar3d_layout_mid, Pa, grid_name, ps);
+ add_field<Required>("pseudo_density", scalar3d_layout_mid, Pa, grid_name, ps);
+ }
+
//Now need to read in the file
if (m_src_pres_type == TIME_DEPENDENT_3D_PROFILE) {
m_num_src_levs = scorpio::get_dimlen(m_datafiles[0],"lev");
@@ -168,11 +184,13 @@ void Nudging::apply_tendency(Field& state, const Field& nudge, const Real dt) co
// Calculate the weight to apply the tendency
const Real dtend = dt / m_timescale;
+ using cview_1d = decltype(state.get_view<const Real*>());
using cview_2d = decltype(state.get_view<const Real**>());
auto state_view = state.get_view<Real**>();
auto nudge_view = nudge.get_view<Real**>();
- cview_2d w_view, pmid_view;
+ cview_1d pbl_height_view;
+ cview_2d w_view, pmid_view, z_mid_view;
if (m_use_weights) {
auto weights = get_helper_field("nudging_weights");
@@ -184,11 +202,63 @@ void Nudging::apply_tendency(Field& state, const Field& nudge, const Real dt) co
auto use_weights = m_use_weights;
auto cutoff = m_refine_remap_vert_cutoff;
+ auto pblh_cutoff = m_pbl_height_cutoff;
auto policy = Kokkos::MDRangePolicy<Kokkos::Rank<2>>({0, 0}, {m_num_cols, m_num_levs});
+
+ if (pblh_cutoff) {
+ // Compute z_mid
+ const auto T_mid_d = get_field_in("T_mid").get_view<const Real**>();
+ const auto qv_d = get_field_in("qv").get_view<const Real**>();
+ const auto p_mid_d = get_field_in("p_mid").get_view<const Real**>();
+ const auto pseudo_density_d = get_field_in("pseudo_density").get_view<const Real**>();
+ const auto z_mid_d = m_z_mid.get_view<Real**>();
+ const auto z_int_d = m_z_int.get_view<Real**>();
+ const auto ncol = m_num_cols;
+ const auto nlev = m_num_levs;
+
+ using KT = KokkosTypes<DefaultDevice>;
+ using ExeSpace = typename KT::ExeSpace;
+ using TPF = ekat::TeamPolicyFactory<ExeSpace>;
+ using PF = scream::PhysicsFunctions<DefaultDevice>;
+
+ const auto scan_policy = TPF::get_thread_range_parallel_scan_team_policy(ncol, nlev);
+ Kokkos::parallel_for(scan_policy, KOKKOS_LAMBDA (const KT::MemberType& team) {
+ const int i = team.league_rank();
+ const auto p_mid_s = ekat::subview(p_mid_d, i);
+ const auto T_mid_s = ekat::subview(T_mid_d, i);
+ const auto qv_s = ekat::subview(qv_d, i);
+ const auto z_int_s = ekat::subview(z_int_d, i);
+ const auto z_mid_s = ekat::subview(z_mid_d, i);
+ // We want z_mid to be AGL, so we set z_surf to 0.
+ const Real z_surf = 0.0;
+ const auto pseudo_density_s = ekat::subview(pseudo_density_d, i);
+
+ // 1. Compute dz (recycle z_mid_s as a temporary)
+ const auto dz_s = z_mid_s; //
+ PF::calculate_dz(team, pseudo_density_s, p_mid_s, T_mid_s, qv_s, dz_s);
+ team.team_barrier();
+
+ // 2. Compute z_int (vertical scan)
+ PF::calculate_z_int(team,nlev,dz_s,z_surf,z_int_s);
+ team.team_barrier();
+
+ // 3. Compute z_mid (int->mid interpolation)
+ PF::calculate_z_mid(team,nlev,z_int_s,z_mid_s);
+ team.team_barrier();
+ });
+ Kokkos::fence();
+ }
+ if (m_pbl_height_cutoff) {
+ pbl_height_view = get_field_in("pbl_height").get_view<const Real*>();
+ z_mid_view = m_z_mid.get_view<const Real**>();
+ }
auto update = KOKKOS_LAMBDA(const int& i, const int& j) {
if (cutoff>0 and pmid_view(i,j)>=cutoff) {
return;
}
+ if (pblh_cutoff and z_mid_view(i,j) < pbl_height_view(i)) {
+ return;
+ }
auto tend = nudge_view(i,j) - state_view(i,j);
if (use_weights) {
@@ -314,6 +384,12 @@ void Nudging::initialize_impl (const RunType /* run_type */)
AtmosphereInput src_weights_input(m_weights_file, m_grid, {nudging_weights},true);
src_weights_input.read_variables();
}
+
+ if (m_pbl_height_cutoff) {
+ // create z_mid and z_int fields
+ m_z_mid = create_helper_field("z_mid", layout_atm, m_grid->name());
+ m_z_int = create_helper_field("z_int", layout_atm, m_grid->name());
+ }
}
// =========================================================================================
diff --git a/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.hpp b/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.hpp
index 054ed6c1ee3a..621752e569fe 100644
--- a/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.hpp
+++ b/components/eamxx/src/physics/nudging/eamxx_nudging_process_interface.hpp
@@ -74,6 +74,8 @@ class Nudging : public AtmosphereProcess
int m_timescale;
bool m_use_weights;
bool m_skip_vert_interpolation;
+ bool m_pbl_height_cutoff;
+ Field m_z_mid, m_z_int;
std::vector<std::string> m_datafiles;
std::string m_static_vertical_pressure_file;
// add nudging weights for regional nudging update
Tangentially related: I just tested requesting z_mid inside of nudging via add_field and got all zeros.
add_field only requests that the FM allocates a field called "z_mid". Nobody computes it for you. Worse: since the FM now does have z_mid, IO will not create a diag for it, as it finds it in the FM. And even if you create it via create_diag, you would have to ensure all diags are registered in the factory before control reaches that point (which if it's Nudging::set_grids, then it's quite early).
Edit: we should register diags early on in the AD sequence, so that processes can just use them in case they need (and don't care about allocating an extra field).