Skip to content

Commit b470097

Browse files
tcclevengertcclevenger
committed
Use field util horiz_contraction to compute means
1 parent 485e692 commit b470097

File tree

2 files changed

+81
-102
lines changed

2 files changed

+81
-102
lines changed

‎components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp‎

Lines changed: 73 additions & 98 deletions
Original file line numberDiff line numberDiff line change
@@ -35,11 +35,21 @@ void IOPForcing::set_grids(const std::shared_ptr<const GridsManager> grids_manag
3535
"Error! IOPDataManager not setup by driver, but IOPForcing"
3636
"being used as an ATM process.\n");
3737

38-
// Compute number of buffer views
38+
// Create helper fields for finding horizontal means
39+
auto level_only_scalar_layout = scalar3d_mid.clone().strip_dim(0);
40+
auto level_only_vector_layout = vector3d_mid.clone().strip_dim(0);
3941
const auto iop_nudge_tq = m_iop_data_manager->get_params().get<bool>("iop_nudge_tq");
4042
const auto iop_nudge_uv = m_iop_data_manager->get_params().get<bool>("iop_nudge_uv");
41-
if (iop_nudge_tq) m_buffer.num_1d_scalar_nlev += 2;
42-
if (iop_nudge_uv) m_buffer.num_1d_scalar_nlev += 2;
43+
if (iop_nudge_tq or iop_nudge_uv) {
44+
create_helper_field("horiz_mean_weights", scalar2d, grid_name);
45+
}
46+
if (iop_nudge_tq) {
47+
create_helper_field("qv_mean", level_only_scalar_layout, grid_name);
48+
create_helper_field("t_mean", level_only_scalar_layout, grid_name);
49+
}
50+
if (iop_nudge_uv) {
51+
create_helper_field("horiz_winds_mean", level_only_vector_layout, grid_name);
52+
}
4353
}
4454
// =========================================================================================
4555
void IOPForcing::
@@ -61,46 +71,22 @@ set_computed_group_impl (const FieldGroup& group)
6171
// =========================================================================================
6272
size_t IOPForcing::requested_buffer_size_in_bytes() const
6373
{
64-
const int nlev_packs = ekat::npack<Pack>(m_num_levs);
65-
const int nlevi_packs = ekat::npack<Pack>(m_num_levs+1);
66-
67-
const size_t temp_view_bytes =
68-
m_buffer.num_1d_scalar_nlev*nlev_packs*sizeof(Pack);
69-
7074
// Number of bytes needed by the WorkspaceManager passed to shoc_main
75+
const int nlevi_packs = ekat::npack<Pack>(m_num_levs+1);
7176
const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs);
7277
const size_t wsm_bytes = WorkspaceMgr::get_total_bytes_needed(nlevi_packs, 7+m_num_tracers, policy);
7378

74-
return temp_view_bytes + wsm_bytes;
79+
return wsm_bytes;
7580
}
7681
// =========================================================================================
7782
void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager)
7883
{
7984
EKAT_REQUIRE_MSG(buffer_manager.allocated_bytes() >= requested_buffer_size_in_bytes(),
8085
"Error! Buffers size not sufficient.\n");
8186

82-
const int nlev_packs = ekat::npack<Pack>(m_num_levs);
8387
const int nlevi_packs = ekat::npack<Pack>(m_num_levs+1);
84-
8588
Pack* mem = reinterpret_cast<Pack*>(buffer_manager.get_memory());
8689

87-
// Temp view data
88-
using mean_view_t = decltype(m_buffer.qv_mean);
89-
const auto iop_nudge_tq = m_iop_data_manager->get_params().get<bool>("iop_nudge_tq");
90-
const auto iop_nudge_uv = m_iop_data_manager->get_params().get<bool>("iop_nudge_uv");
91-
if (iop_nudge_tq) {
92-
m_buffer.qv_mean = mean_view_t(mem, nlev_packs);
93-
mem += m_buffer.qv_mean.size();
94-
m_buffer.t_mean = mean_view_t(mem, nlev_packs);
95-
mem += m_buffer.t_mean.size();
96-
}
97-
if (iop_nudge_uv) {
98-
m_buffer.u_mean = mean_view_t(mem, nlev_packs);
99-
mem += m_buffer.u_mean.size();
100-
m_buffer.v_mean = mean_view_t(mem, nlev_packs);
101-
mem += m_buffer.v_mean.size();
102-
}
103-
10490
// WSM data
10591
m_buffer.wsm_data = mem;
10692

@@ -112,6 +98,22 @@ void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager)
11298
EKAT_REQUIRE_MSG(used_mem==requested_buffer_size_in_bytes(), "Error! Used memory != requested memory for IOPForcing.\n");
11399
}
114100
// =========================================================================================
101+
void IOPForcing::create_helper_field (const std::string& name,
102+
const FieldLayout& layout,
103+
const std::string& grid_name)
104+
{
105+
using namespace ekat::units;
106+
FieldIdentifier id(name,layout,Units::nondimensional(),grid_name);
107+
108+
// Create the field. Init with NaN's, so we spot instances of uninited memory usage
109+
Field f(id);
110+
f.get_header().get_alloc_properties().request_allocation();
111+
f.allocate_view();
112+
f.deep_copy(ekat::ScalarTraits<Real>::invalid());
113+
114+
m_helper_fields[name] = f;
115+
}
116+
// =========================================================================================
115117
void IOPForcing::initialize_impl (const RunType run_type)
116118
{
117119
// Set field property checks for the fields in this process
@@ -126,6 +128,12 @@ void IOPForcing::initialize_impl (const RunType run_type)
126128
const auto nlevi_packs = ekat::npack<Pack>(m_num_levs+1);
127129
const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs);
128130
m_workspace_mgr.setup(m_buffer.wsm_data, nlevi_packs, 7+m_num_tracers, policy);
131+
132+
// Compute field for horizontal contraction weights (1/num_global_dofs)
133+
const auto iop_nudge_tq = m_iop_data_manager->get_params().get<bool>("iop_nudge_tq");
134+
const auto iop_nudge_uv = m_iop_data_manager->get_params().get<bool>("iop_nudge_uv");
135+
const Real one_over_num_dofs = 1.0/m_grid->get_num_global_dofs();
136+
if (iop_nudge_tq or iop_nudge_uv) m_helper_fields.at("horiz_mean_weights").deep_copy(one_over_num_dofs);
129137
}
130138
// =========================================================================================
131139
KOKKOS_FUNCTION
@@ -437,55 +445,22 @@ void IOPForcing::run_impl (const double dt)
437445
// and observed quantities of T, Q, u, and v
438446
if (iop_nudge_tq or iop_nudge_uv) {
439447
// Compute domain mean of qv, T_mid, u, and v
440-
const auto qv_mean = m_buffer.qv_mean;
441-
const auto t_mean = m_buffer.t_mean;
442-
const auto u_mean = m_buffer.u_mean;
443-
const auto v_mean = m_buffer.v_mean;
444-
445-
const auto qv_mean_h = Kokkos::create_mirror_view(qv_mean);
446-
const auto t_mean_h = Kokkos::create_mirror_view(t_mean);
447-
const auto u_mean_h = Kokkos::create_mirror_view(u_mean);
448-
const auto v_mean_h = Kokkos::create_mirror_view(v_mean);
449-
450-
const auto num_global_cols = m_grid->get_num_global_dofs();
451-
for (int k=0; k<m_num_levs; ++k) {
452-
if (iop_nudge_tq){
453-
Real& qv_mean_k = qv_mean_h(k/Pack::n)[k%Pack::n];
454-
Real& t_mean_k = t_mean_h(k/Pack::n)[k%Pack::n];
455-
Kokkos::parallel_reduce("compute_domain_means_tq",
456-
m_num_cols,
457-
KOKKOS_LAMBDA (const int icol, Real& q_sum, Real& t_sum) {
458-
q_sum += qv(icol, k/Pack::n)[k%Pack::n];
459-
t_sum += T_mid(icol, k/Pack::n)[k%Pack::n];
460-
}, qv_mean_k, t_mean_k);
461-
462-
m_comm.all_reduce(&qv_mean_k, 1, MPI_SUM);
463-
m_comm.all_reduce(&t_mean_k, 1, MPI_SUM);
464-
465-
qv_mean_k /= num_global_cols;
466-
t_mean_k /= num_global_cols;
467-
}
468-
if (iop_nudge_uv){
469-
Real& u_mean_k = u_mean_h(k/Pack::n)[k%Pack::n];
470-
Real& v_mean_k = v_mean_h(k/Pack::n)[k%Pack::n];
471-
Kokkos::parallel_reduce("compute_domain_means_uv",
472-
m_num_cols,
473-
KOKKOS_LAMBDA (const int icol, Real& u_sum, Real& v_sum) {
474-
u_sum += horiz_winds(icol, 0, k/Pack::n)[k%Pack::n];
475-
v_sum += horiz_winds(icol, 1, k/Pack::n)[k%Pack::n];
476-
}, u_mean_k, v_mean_k);
477-
478-
m_comm.all_reduce(&u_mean_k, 1, MPI_SUM);
479-
m_comm.all_reduce(&v_mean_k, 1, MPI_SUM);
480-
481-
u_mean_k /= num_global_cols;
482-
v_mean_k /= num_global_cols;
483-
}
448+
view_1d<Real> qv_mean, t_mean;
449+
view_2d<Real> horiz_winds_mean;
450+
if (iop_nudge_tq){
451+
horiz_contraction(m_helper_fields.at("qv_mean"), get_field_out("qv"),
452+
m_helper_fields.at("horiz_mean_weights"), m_comm);
453+
qv_mean = m_helper_fields.at("qv_mean").get_view<Real*>();
454+
455+
horiz_contraction(m_helper_fields.at("t_mean"), get_field_out("T_mid"),
456+
m_helper_fields.at("horiz_mean_weights"), m_comm);
457+
t_mean = m_helper_fields.at("t_mean").get_view<Real*>();
458+
}
459+
if (iop_nudge_uv){
460+
horiz_contraction(m_helper_fields.at("horiz_winds_mean"), get_field_out("horiz_winds"),
461+
m_helper_fields.at("horiz_mean_weights"), m_comm);
462+
horiz_winds_mean = m_helper_fields.at("horiz_winds_mean").get_view<Real**>();
484463
}
485-
Kokkos::deep_copy(qv_mean, qv_mean_h);
486-
Kokkos::deep_copy(t_mean, t_mean_h);
487-
Kokkos::deep_copy(u_mean, u_mean_h);
488-
Kokkos::deep_copy(v_mean, v_mean_h);
489464

490465
// Apply relaxation
491466
const auto rtau = std::max(dt, iop_nudge_tscale);
@@ -512,28 +487,28 @@ void IOPForcing::run_impl (const double dt)
512487
});
513488
team.team_barrier();
514489

515-
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) {
516-
if (iop_nudge_tq) {
517-
// Restrict nudging of T and qv to certain levels if requested by user
518-
// IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high
519-
// is in units of [hPa], thus convert iop_nudge_tq_low/high
520-
Mask nudge_level(false);
521-
int max_size = hyam.size();
522-
for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) {
523-
const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop;
524-
nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100
525-
and
526-
pressure_from_iop >= iop_nudge_tq_high*100);
527-
}
528-
529-
qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0);
530-
T_mid_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0);
531-
}
532-
if (iop_nudge_uv) {
533-
u_i(k).update(u_mean(k) - u_iop(k), -dt/rtau, 1.0);
534-
v_i(k).update(v_mean(k) - v_iop(k), -dt/rtau, 1.0);
490+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) {
491+
if (iop_nudge_tq) {
492+
// Restrict nudging of T and qv to certain levels if requested by user
493+
// IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high
494+
// is in units of [hPa], thus convert iop_nudge_tq_low/high
495+
Mask nudge_level(false);
496+
int max_size = hyam.size();
497+
for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) {
498+
const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop;
499+
nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100
500+
and
501+
pressure_from_iop >= iop_nudge_tq_high*100);
535502
}
536-
});
503+
504+
qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0);
505+
T_mid_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0);
506+
}
507+
if (iop_nudge_uv) {
508+
u_i(k).update(horiz_winds_mean(0, k) - u_iop(k), -dt/rtau, 1.0);
509+
v_i(k).update(horiz_winds_mean(1, k) - v_iop(k), -dt/rtau, 1.0);
510+
}
511+
});
537512

538513
// Release WS views
539514
ws.release_many_contiguous<1>({&ref_p_mid});

‎components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,11 @@ class IOPForcing : public scream::AtmosphereProcess
116116

117117
void finalize_impl () {}
118118

119+
// Creates an helper field, not to be shared with the AD's FieldManager
120+
void create_helper_field (const std::string& name,
121+
const FieldLayout& layout,
122+
const std::string& grid_name);
123+
119124
void set_computed_group_impl (const FieldGroup& group);
120125

121126
// Computes total number of bytes needed for local variables
@@ -132,13 +137,12 @@ class IOPForcing : public scream::AtmosphereProcess
132137
Int m_num_tracers;
133138

134139
struct Buffer {
135-
int num_1d_scalar_nlev = 0;
136-
137-
uview_1d<Pack> qv_mean, t_mean, u_mean, v_mean;
138-
139140
Pack* wsm_data;
140141
};
141142

143+
// Some helper fields.
144+
std::map<std::string,Field> m_helper_fields;
145+
142146
// Struct which contains local variables
143147
Buffer m_buffer;
144148

0 commit comments

Comments
 (0)