From 254bba50e53b8b4933709f898055f0c3f0548346 Mon Sep 17 00:00:00 2001 From: tcclevenger Date: Fri, 1 Nov 2024 00:31:22 -0600 Subject: [PATCH 1/5] Create IOPForcing ATM process Moves IOP forcing from within HOMME interface to it's own physics process, called directly after dynamics process. Also, changes name of IntensiveObservationPeriod class to IOPDataManager. --- .../cime_config/namelist_defaults_scream.xml | 4 + .../eamxx/src/control/atmosphere_driver.cpp | 46 +- .../eamxx/src/control/atmosphere_driver.hpp | 8 +- .../atmosphere_surface_coupling_importer.cpp | 21 +- .../eamxx/src/dynamics/homme/CMakeLists.txt | 1 - .../src/dynamics/homme/eamxx_homme_iop.cpp | 610 ------------------ .../homme/eamxx_homme_process_interface.cpp | 9 - .../homme/eamxx_homme_process_interface.hpp | 38 -- .../eamxx/src/mct_coupling/CMakeLists.txt | 1 + components/eamxx/src/physics/CMakeLists.txt | 10 +- .../src/physics/iop_forcing/CMakeLists.txt | 5 + .../eamxx_iop_forcing_process_interface.cpp | 546 ++++++++++++++++ .../eamxx_iop_forcing_process_interface.hpp | 153 +++++ .../eamxx/src/physics/register_physics.hpp | 6 + .../rrtmgp/eamxx_rrtmgp_process_interface.cpp | 6 +- .../shoc/eamxx_shoc_process_interface.cpp | 2 +- .../spa/eamxx_spa_process_interface.cpp | 8 +- .../eamxx/src/physics/spa/spa_functions.hpp | 16 +- .../src/physics/spa/spa_functions_impl.hpp | 4 +- components/eamxx/src/share/CMakeLists.txt | 2 +- .../IOPDataManager.cpp} | 32 +- .../IOPDataManager.hpp} | 24 +- .../share/atm_process/atmosphere_process.hpp | 10 +- .../atm_process/atmosphere_process_group.hpp | 4 +- .../CMakeLists.txt | 2 +- .../input.yaml | 2 +- 26 files changed, 813 insertions(+), 757 deletions(-) delete mode 100644 components/eamxx/src/dynamics/homme/eamxx_homme_iop.cpp create mode 100644 components/eamxx/src/physics/iop_forcing/CMakeLists.txt create mode 100644 components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp create mode 100644 components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp rename components/eamxx/src/share/{iop/intensive_observation_period.cpp => atm_process/IOPDataManager.cpp} (98%) rename components/eamxx/src/share/{iop/intensive_observation_period.hpp => atm_process/IOPDataManager.hpp} (92%) diff --git a/components/eamxx/cime_config/namelist_defaults_scream.xml b/components/eamxx/cime_config/namelist_defaults_scream.xml index 6d4c36b81bd4..5b50f7031092 100644 --- a/components/eamxx/cime_config/namelist_defaults_scream.xml +++ b/components/eamxx/cime_config/namelist_defaults_scream.xml @@ -427,6 +427,9 @@ be lost if SCREAM_HACK_XML is not enabled. false + + + 1,2 @@ -556,6 +559,7 @@ be lost if SCREAM_HACK_XML is not enabled. mac_aero_mic,rrtmgp + iop_forcing,mac_aero_mic,rrtmgp diff --git a/components/eamxx/src/control/atmosphere_driver.cpp b/components/eamxx/src/control/atmosphere_driver.cpp index 0f1cb1e31ab0..e935cff5cfcb 100644 --- a/components/eamxx/src/control/atmosphere_driver.cpp +++ b/components/eamxx/src/control/atmosphere_driver.cpp @@ -180,24 +180,24 @@ init_time_stamps (const util::TimeStamp& run_t0, const util::TimeStamp& case_t0, } void AtmosphereDriver:: -setup_iop () +setup_iop_data_manager () { // At this point, must have comm, params, initialized timestamps created. check_ad_status(s_comm_set | s_params_set | s_ts_inited); // Check to make sure iop is not already initialized - EKAT_REQUIRE_MSG(not m_iop, "Error! setup_iop() is called, but IOP already set up.\n"); + EKAT_REQUIRE_MSG(not m_iop_data_manager, "Error! setup_iop_data_manager() is called, but IOP already set up.\n"); // This function should only be called if we are enabling IOP const bool enable_iop = m_atm_params.sublist("driver_options").get("enable_iop", false); - EKAT_REQUIRE_MSG(enable_iop, "Error! setup_iop() is called, but enable_iop=false " + EKAT_REQUIRE_MSG(enable_iop, "Error! setup_iop_data_manager() is called, but enable_iop=false " "in driver_options parameters.\n"); // Params must include iop_options sublist. const auto iop_sublist_exists = m_atm_params.isSublist("iop_options"); EKAT_REQUIRE_MSG(iop_sublist_exists, - "Error! setup_iop() is called, but no iop_options " + "Error! setup_iop_data_manager() is called, but no iop_options " "defined in parameters.\n"); const auto iop_params = m_atm_params.sublist("iop_options"); @@ -206,15 +206,15 @@ setup_iop () const auto hyam = phys_grid->get_geometry_data("hyam"); const auto hybm = phys_grid->get_geometry_data("hybm"); - m_iop = std::make_shared(m_atm_comm, - iop_params, - m_run_t0, - nlevs, - hyam, - hybm); + m_iop_data_manager = std::make_shared(m_atm_comm, + iop_params, + m_run_t0, + nlevs, + hyam, + hybm); // Set IOP object in atm processes - m_atm_process_group->set_iop(m_iop); + m_atm_process_group->set_iop_data_manager(m_iop_data_manager); } void AtmosphereDriver::create_atm_processes() @@ -295,7 +295,7 @@ void AtmosphereDriver::create_grids() const bool enable_iop = m_atm_params.sublist("driver_options").get("enable_iop", false); if (enable_iop) { - setup_iop (); + setup_iop_data_manager (); } // Set the grids in the processes. Do this by passing the grids manager. @@ -1203,7 +1203,7 @@ void AtmosphereDriver::set_initial_conditions () } } - if (m_iop) { + if (m_iop_data_manager) { // For runs with IOP, call to setup io grids and lat // lon information needed for reading from file // We use a single topo file for both GLL and PG2 runs. All @@ -1213,13 +1213,13 @@ void AtmosphereDriver::set_initial_conditions () for (const auto& it : m_field_mgrs) { const auto& grid_name = it.first; if (ic_fields_names[grid_name].size() > 0 or - topography_eamxx_fields_names[grid_name].size() > 0) { + topography_eamxx_fields_names[grid_name].size() > 0) { const auto& file_name = grid_name == "Physics GLL" ? ic_pl.get("Filename") : ic_pl.get("topography_filename"); - m_iop->setup_io_info(file_name, it.second->get_grid()); + m_iop_data_manager->setup_io_info(file_name, it.second->get_grid()); } } } @@ -1231,12 +1231,12 @@ void AtmosphereDriver::set_initial_conditions () m_atm_logger->info(" [EAMxx] IC filename: " + file_name); for (const auto& it : m_field_mgrs) { const auto& grid_name = it.first; - if (not m_iop) { + if (not m_iop_data_manager) { read_fields_from_file (ic_fields_names[grid_name],it.second->get_grid(),file_name,m_current_ts); } else { // For IOP enabled, we load from file and copy data from the closest // lat/lon column to every other column - m_iop->read_fields_from_file_for_iop(file_name, + m_iop_data_manager->read_fields_from_file_for_iop(file_name, ic_fields_names[grid_name], m_current_ts, it.second); @@ -1306,7 +1306,7 @@ void AtmosphereDriver::set_initial_conditions () m_atm_logger->info(" filename: " + file_name); for (const auto& it : m_field_mgrs) { const auto& grid_name = it.first; - if (not m_iop) { + if (not m_iop_data_manager) { // Topography files always use "ncol_d" for the GLL grid value of ncol. // To ensure we read in the correct value, we must change the name for that dimension auto io_grid = it.second->get_grid(); @@ -1322,7 +1322,7 @@ void AtmosphereDriver::set_initial_conditions () } else { // For IOP enabled, we load from file and copy data from the closest // lat/lon column to every other column - m_iop->read_fields_from_file_for_iop(file_name, + m_iop_data_manager->read_fields_from_file_for_iop(file_name, topography_file_fields_names[grid_name], topography_eamxx_fields_names[grid_name], m_current_ts, @@ -1347,16 +1347,16 @@ void AtmosphereDriver::set_initial_conditions () m_atm_params.sublist("provenance").set("topography_file","NONE"); } - if (m_iop) { + if (m_iop_data_manager) { // Load IOP data file data for initial time stamp - m_iop->read_iop_file_data(m_current_ts); + m_iop_data_manager->read_iop_file_data(m_current_ts); // Now that ICs are processed, set appropriate fields using IOP file data. // Since ICs are loaded on GLL grid, we set those fields only and dynamics // will take care of the rest (for PG2 case). if (m_field_mgrs.count("Physics GLL") > 0) { const auto& fm = m_field_mgrs.at("Physics GLL"); - m_iop->set_fields_from_iop_data(fm); + m_iop_data_manager->set_fields_from_iop_data(fm); } } @@ -1754,7 +1754,7 @@ void AtmosphereDriver::finalize ( /* inputs? */ ) { } // Destroy iop - m_iop = nullptr; + m_iop_data_manager = nullptr; // Destroy the buffer manager m_memory_buffer = nullptr; diff --git a/components/eamxx/src/control/atmosphere_driver.hpp b/components/eamxx/src/control/atmosphere_driver.hpp index a3acfba5d945..9b191371b36d 100644 --- a/components/eamxx/src/control/atmosphere_driver.hpp +++ b/components/eamxx/src/control/atmosphere_driver.hpp @@ -2,7 +2,6 @@ #define SCREAM_ATMOSPHERE_DRIVER_HPP #include "control/surface_coupling_utils.hpp" -#include "share/iop/intensive_observation_period.hpp" #include "share/field/field_manager.hpp" #include "share/grid/grids_manager.hpp" #include "share/util/scream_time_stamp.hpp" @@ -11,6 +10,7 @@ #include "share/io/scorpio_input.hpp" #include "share/atm_process/ATMBufferManager.hpp" #include "share/atm_process/SCDataManager.hpp" +#include "share/atm_process/IOPDataManager.hpp" #include "ekat/logging/ekat_logger.hpp" #include "ekat/mpi/ekat_comm.hpp" @@ -72,8 +72,8 @@ class AtmosphereDriver // Set AD params void init_scorpio (const int atm_id = 0); - // Setup IntensiveObservationPeriod - void setup_iop (); + // Setup IOPDataManager + void setup_iop_data_manager (); // Create atm processes, without initializing them void create_atm_processes (); @@ -217,7 +217,7 @@ class AtmosphereDriver std::shared_ptr m_surface_coupling_import_data_manager; std::shared_ptr m_surface_coupling_export_data_manager; - std::shared_ptr m_iop; + std::shared_ptr m_iop_data_manager; // This is the time stamp at the beginning of the time step. util::TimeStamp m_current_ts; diff --git a/components/eamxx/src/control/atmosphere_surface_coupling_importer.cpp b/components/eamxx/src/control/atmosphere_surface_coupling_importer.cpp index 2c3360a3b4f7..385e57cae556 100644 --- a/components/eamxx/src/control/atmosphere_surface_coupling_importer.cpp +++ b/components/eamxx/src/control/atmosphere_surface_coupling_importer.cpp @@ -208,8 +208,8 @@ void SurfaceCouplingImporter::do_import(const bool called_during_initialization) }); #endif - if (m_iop) { - if (m_iop->get_params().get("iop_srf_prop")) { + if (m_iop_data_manager) { + if (m_iop_data_manager->get_params().get("iop_srf_prop")) { // Overwrite imports with data from IOP file overwrite_iop_imports(called_during_initialization); } @@ -221,9 +221,12 @@ void SurfaceCouplingImporter::overwrite_iop_imports (const bool called_during_in using policy_type = KokkosTypes::RangePolicy; using C = physics::Constants; - const auto has_lhflx = m_iop->has_iop_field("lhflx"); - const auto has_shflx = m_iop->has_iop_field("shflx"); - const auto has_Tg = m_iop->has_iop_field("Tg"); + const auto has_lhflx = m_iop_data_manager->has_iop_field("lhflx"); + const auto has_shflx = m_iop_data_manager->has_iop_field("shflx"); + const auto has_Tg = m_iop_data_manager->has_iop_field("Tg"); + + // Read IOP file for current time step, if necessary + m_iop_data_manager->read_iop_file_data(timestamp()); static constexpr Real latvap = C::LatVap; static constexpr Real stebol = C::stebol; @@ -243,19 +246,19 @@ void SurfaceCouplingImporter::overwrite_iop_imports (const bool called_during_in // Store IOP surf data into col_val Real col_val(std::nan("")); if (fname == "surf_evap" && has_lhflx) { - const auto f = m_iop->get_iop_field("lhflx"); + const auto f = m_iop_data_manager->get_iop_field("lhflx"); f.sync_to_host(); col_val = f.get_view()()/latvap; } else if (fname == "surf_sens_flux" && has_shflx) { - const auto f = m_iop->get_iop_field("shflx"); + const auto f = m_iop_data_manager->get_iop_field("shflx"); f.sync_to_host(); col_val = f.get_view()(); } else if (fname == "surf_radiative_T" && has_Tg) { - const auto f = m_iop->get_iop_field("Tg"); + const auto f = m_iop_data_manager->get_iop_field("Tg"); f.sync_to_host(); col_val = f.get_view()(); } else if (fname == "surf_lw_flux_up" && has_Tg) { - const auto f = m_iop->get_iop_field("Tg"); + const auto f = m_iop_data_manager->get_iop_field("Tg"); f.sync_to_host(); col_val = stebol*std::pow(f.get_view()(), 4); } else { diff --git a/components/eamxx/src/dynamics/homme/CMakeLists.txt b/components/eamxx/src/dynamics/homme/CMakeLists.txt index b6a69a3605f9..8e2b3577567a 100644 --- a/components/eamxx/src/dynamics/homme/CMakeLists.txt +++ b/components/eamxx/src/dynamics/homme/CMakeLists.txt @@ -146,7 +146,6 @@ macro (CreateDynamicsLib HOMME_TARGET NP PLEV QSIZE) ${SCREAM_DYNAMICS_SRC_DIR}/eamxx_homme_process_interface.cpp ${SCREAM_DYNAMICS_SRC_DIR}/eamxx_homme_fv_phys.cpp ${SCREAM_DYNAMICS_SRC_DIR}/eamxx_homme_rayleigh_friction.cpp - ${SCREAM_DYNAMICS_SRC_DIR}/eamxx_homme_iop.cpp ${SCREAM_DYNAMICS_SRC_DIR}/physics_dynamics_remapper.cpp ${SCREAM_DYNAMICS_SRC_DIR}/homme_grids_manager.cpp ${SCREAM_DYNAMICS_SRC_DIR}/interface/homme_context_mod.F90 diff --git a/components/eamxx/src/dynamics/homme/eamxx_homme_iop.cpp b/components/eamxx/src/dynamics/homme/eamxx_homme_iop.cpp deleted file mode 100644 index 9c04a3f1ba6c..000000000000 --- a/components/eamxx/src/dynamics/homme/eamxx_homme_iop.cpp +++ /dev/null @@ -1,610 +0,0 @@ -#include "eamxx_homme_process_interface.hpp" - -// EAMxx includes -#include "dynamics/homme/homme_dimensions.hpp" -#include "dynamics/homme/homme_dynamics_helpers.hpp" -#include "physics/share/physics_constants.hpp" -#include "share/iop/intensive_observation_period.hpp" -#include "share/util/scream_column_ops.hpp" - -// Homme includes -#include "Context.hpp" -#include "ColumnOps.hpp" -#include "HommexxEnums.hpp" -#include "HybridVCoord.hpp" -#include "SimulationParams.hpp" -#include "Types.hpp" - -// SCREAM includes -#include "share/util/scream_common_physics_functions.hpp" - -// EKAT includes -#include "ekat/ekat_workspace.hpp" -#include "ekat/kokkos/ekat_kokkos_types.hpp" - -namespace scream { - -// Compute effects of large scale subsidence on T, q, u, and v. -KOKKOS_FUNCTION -void HommeDynamics:: -advance_iop_subsidence(const KT::MemberType& team, - const int nlevs, - const Real dt, - const Real ps, - const view_1d& pmid, - const view_1d& pint, - const view_1d& pdel, - const view_1d& omega, - const Workspace& workspace, - const view_1d& u, - const view_1d& v, - const view_1d& T, - const view_2d& Q) -{ - using ColOps = ColumnOps; - using C = physics::Constants; - constexpr Real Rair = C::Rair; - constexpr Real Cpair = C::Cpair; - - const auto n_q_tracers = Q.extent_int(0); - const auto nlev_packs = ekat::npack(nlevs); - - // Get some temporary views from WS - uview_1d omega_int, delta_u, delta_v, delta_T, tmp; - workspace.take_many_contiguous_unsafe<4>({"omega_int", "delta_u", "delta_v", "delta_T"}, - {&omega_int, &delta_u, &delta_v, &delta_T}); - const auto delta_Q_slot = workspace.take_macro_block("delta_Q", n_q_tracers); - uview_2d delta_Q(delta_Q_slot.data(), n_q_tracers, nlev_packs); - - auto s_pmid = ekat::scalarize(pmid); - auto s_omega = ekat::scalarize(omega); - auto s_delta_u = ekat::scalarize(delta_u); - auto s_delta_v = ekat::scalarize(delta_v); - auto s_delta_T = ekat::scalarize(delta_T); - auto s_delta_Q = ekat::scalarize(delta_Q); - auto s_omega_int = ekat::scalarize(omega_int); - - // Compute omega on the interface grid by using a weighted average in pressure - const int pack_begin = 1/Pack::n, pack_end = (nlevs-1)/Pack::n; - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pack_begin, pack_end+1), [&] (const int k){ - auto range_pack = ekat::range(k*Pack::n); - range_pack.set(range_pack<1, 1); - Pack pmid_k, pmid_km1, omega_k, omega_km1; - ekat::index_and_shift<-1>(s_pmid, range_pack, pmid_k, pmid_km1); - ekat::index_and_shift<-1>(s_omega, range_pack, omega_k, omega_km1); - - const auto weight = (pint(k) - pmid_km1)/(pmid_k - pmid_km1); - omega_int(k).set(range_pack>=1 and range_pack<=nlevs-1, - weight*omega_k + (1-weight)*omega_km1); - }); - omega_int(0)[0] = 0; - omega_int(nlevs/Pack::n)[nlevs%Pack::n] = 0; - - // Compute delta views for u, v, T, and Q (e.g., u(k+1) - u(k), k=0,...,nlevs-2) - ColOps::compute_midpoint_delta(team, nlevs-1, u, delta_u); - ColOps::compute_midpoint_delta(team, nlevs-1, v, delta_v); - ColOps::compute_midpoint_delta(team, nlevs-1, T, delta_T); - for (int iq=0; iq(k*Pack::n); - const auto at_top = range_pack==0; - const auto not_at_top = not at_top; - const auto at_bot = range_pack==nlevs-1; - const auto not_at_bot = not at_bot; - const bool any_at_top = at_top.any(); - const bool any_at_bot = at_bot.any(); - - // Get delta(k-1) packs. The range pack should not - // contain index 0 (so that we don't attempt to access - // k=-1 index) or index > nlevs-2 (since delta_* views - // are size nlevs-1). - auto range_pack_for_m1_shift = range_pack; - range_pack_for_m1_shift.set(range_pack<1, 1); - range_pack_for_m1_shift.set(range_pack>nlevs-2, nlevs-2); - Pack delta_u_k, delta_u_km1, - delta_v_k, delta_v_km1, - delta_T_k, delta_T_km1; - ekat::index_and_shift<-1>(s_delta_u, range_pack_for_m1_shift, delta_u_k, delta_u_km1); - ekat::index_and_shift<-1>(s_delta_v, range_pack_for_m1_shift, delta_v_k, delta_v_km1); - ekat::index_and_shift<-1>(s_delta_T, range_pack_for_m1_shift, delta_T_k, delta_T_km1); - - // At the top and bottom of the model, set the end points for - // delta_*_k and delta_*_km1 to be the first and last entries - // of delta_*, respectively. - if (any_at_top) { - delta_u_k.set(at_top, s_delta_u(0)); - delta_v_k.set(at_top, s_delta_v(0)); - delta_T_k.set(at_top, s_delta_T(0)); - } - if (any_at_bot) { - delta_u_km1.set(at_bot, s_delta_u(nlevs-2)); - delta_v_km1.set(at_bot, s_delta_v(nlevs-2)); - delta_T_km1.set(at_bot, s_delta_T(nlevs-2)); - } - - // Get omega_int(k+1) pack. The range pack should not - // contain index > nlevs-1 (since omega_int is size nlevs+1). - auto range_pack_for_p1_shift = range_pack; - range_pack_for_p1_shift.set(range_pack>nlevs-1, nlevs-1); - Pack omega_int_k, omega_int_kp1; - ekat::index_and_shift<1>(s_omega_int, range_pack, omega_int_k, omega_int_kp1); - - const auto fac = (dt/2)/pdel(k); - - // Update u - u(k).update(not_at_bot, fac*omega_int_kp1*delta_u_k, -1, 1); - u(k).update(not_at_top, fac*omega_int_k*delta_u_km1, -1, 1); - - // Update v - v(k).update(not_at_bot, fac*omega_int_kp1*delta_v_k, -1, 1); - v(k).update(not_at_top, fac*omega_int_k*delta_v_km1, -1, 1); - - // Before updating T, first scale using thermal - // expansion term due to LS vertical advection - T(k) *= 1 + (dt*Rair/Cpair)*omega(k)/pmid(k); - - // Update T - T(k).update(not_at_bot, fac*omega_int_kp1*delta_T_k, -1, 1); - T(k).update(not_at_top, fac*omega_int_k*delta_T_km1, -1, 1); - - // Update Q - Pack delta_tracer_k, delta_tracer_km1; - for (int iq=0; iq(s_delta_tracer, range_pack_for_m1_shift, delta_tracer_k, delta_tracer_km1); - if (any_at_top) delta_tracer_k.set(at_top, s_delta_tracer(0)); - if (any_at_bot) delta_tracer_km1.set(at_bot, s_delta_tracer(nlevs-2)); - - Q(iq, k).update(not_at_bot, fac*omega_int_kp1*delta_tracer_k, -1, 1); - Q(iq, k).update(not_at_top, fac*omega_int_k*delta_tracer_km1, -1, 1); - } - }); - - // Release WS views - workspace.release_macro_block(delta_Q_slot, n_q_tracers); - workspace.release_many_contiguous<4>({&omega_int, &delta_u, &delta_v, &delta_T}); -} - -// Apply large scale forcing for temperature and water vapor as provided by the IOP file -KOKKOS_FUNCTION -void HommeDynamics:: -advance_iop_forcing(const KT::MemberType& team, - const int nlevs, - const Real dt, - const view_1d& divT, - const view_1d& divq, - const view_1d& T, - const view_1d& qv) -{ - const auto nlev_packs = ekat::npack(nlevs); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&] (const int k) { - T(k).update(divT(k), dt, 1.0); - qv(k).update(divq(k), dt, 1.0); - }); -} - -// Provide coriolis forcing to u and v winds, using large scale winds specified in IOP forcing file. -KOKKOS_FUNCTION -void HommeDynamics:: -iop_apply_coriolis(const KT::MemberType& team, - const int nlevs, - const Real dt, - const Real lat, - const view_1d& u_ls, - const view_1d& v_ls, - const view_1d& u, - const view_1d& v) -{ - using C = physics::Constants; - constexpr Real pi = C::Pi; - constexpr Real earth_rotation = C::omega; - - // Compute coriolis force - const auto fcor = 2*earth_rotation*std::sin(lat*pi/180); - - const auto nlev_packs = ekat::npack(nlevs); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&] (const int k) { - const auto u_cor = v(k) - v_ls(k); - const auto v_cor = u(k) - u_ls(k); - u(k).update(u_cor, dt*fcor, 1.0); - v(k).update(v_cor, -dt*fcor, 1.0); - }); -} - -void HommeDynamics:: -apply_iop_forcing(const Real dt) -{ - using ESU = ekat::ExeSpaceUtils; - using PF = PhysicsFunctions; - using ColOps = ColumnOps; - - // Homme objects - const auto& c = Homme::Context::singleton(); - const auto& hvcoord = c.get(); - const auto& params = c.get(); - - // Dimensions - constexpr int NGP = HOMMEXX_NP; - constexpr int NLEV = HOMMEXX_NUM_LEV; - constexpr int NLEVI = HOMMEXX_NUM_LEV_P; - const auto nelem = m_dyn_grid->get_num_local_dofs()/(NGP*NGP); - const auto total_levels = m_dyn_grid->get_num_vertical_levels(); - const auto qsize = params.qsize; - - // Sanity checks since we will be switching between ekat::Pack - // and Homme::Scalar view types - EKAT_ASSERT_MSG(NLEV == ekat::npack(total_levels), - "Error! Dimension for vectorized Homme levels does not match level dimension " - "of the packed views used here. Check that Pack typedef is using a pack size " - "consistent with Homme's vector size.\n"); - EKAT_ASSERT_MSG(NLEVI == ekat::npack(total_levels+1), - "Error! Dimension for vectorized Homme levels does not match level dimension " - "of the packed views used here. Check that Pack typedef is using a pack size " - "consistent with Homme's vector size.\n"); - - // Hybrid coord values - const auto ps0 = hvcoord.ps0; - const auto hyam = m_dyn_grid->get_geometry_data("hyam").get_view(); - const auto hybm = m_dyn_grid->get_geometry_data("hybm").get_view(); - const auto hyai = m_dyn_grid->get_geometry_data("hyai").get_view(); - const auto hybi = m_dyn_grid->get_geometry_data("hybi").get_view(); - - // Homme element states - auto ps_dyn = get_internal_field("ps_dyn").get_view(); - auto dp3d_dyn = get_internal_field("dp3d_dyn").get_view(); - auto vtheta_dp_dyn = get_internal_field("vtheta_dp_dyn").get_view(); - auto phi_int_dyn = get_internal_field("phi_int_dyn").get_view(); - auto v_dyn = get_internal_field("v_dyn").get_view(); - auto Q_dyn = m_helper_fields.at("Q_dyn").get_view(); - auto Qdp_dyn = get_internal_field("Qdp_dyn").get_view(); - - // Load data from IOP files, if necessary - m_iop->read_iop_file_data(timestamp()); - - // Define local IOP param values - const auto iop_dosubsidence = m_iop->get_params().get("iop_dosubsidence"); - const auto iop_coriolis = m_iop->get_params().get("iop_coriolis"); - const auto iop_nudge_tq = m_iop->get_params().get("iop_nudge_tq"); - const auto iop_nudge_uv = m_iop->get_params().get("iop_nudge_uv"); - const auto use_large_scale_wind = m_iop->get_params().get("use_large_scale_wind"); - const auto use_3d_forcing = m_iop->get_params().get("use_3d_forcing"); - const auto lat = m_iop->get_params().get("target_latitude"); - const auto iop_nudge_tscale = m_iop->get_params().get("iop_nudge_tscale"); - const auto iop_nudge_tq_low = m_iop->get_params().get("iop_nudge_tq_low"); - const auto iop_nudge_tq_high = m_iop->get_params().get("iop_nudge_tq_high"); - - // Define local IOP field views - const Real ps_iop = m_iop->get_iop_field("Ps").get_view()(); - view_1d omega, divT, divq, u_ls, v_ls, qv_iop, t_iop, u_iop, v_iop; - divT = use_3d_forcing ? m_iop->get_iop_field("divT3d").get_view() - : m_iop->get_iop_field("divT").get_view(); - divq = use_3d_forcing ? m_iop->get_iop_field("divq3d").get_view() - : m_iop->get_iop_field("divq").get_view(); - if (iop_dosubsidence) { - omega = m_iop->get_iop_field("omega").get_view(); - } - if (iop_coriolis) { - u_ls = m_iop->get_iop_field("u_ls").get_view(); - v_ls = m_iop->get_iop_field("v_ls").get_view(); - } - if (iop_nudge_tq) { - qv_iop = m_iop->get_iop_field("q").get_view(); - t_iop = m_iop->get_iop_field("T").get_view(); - } - if (iop_nudge_uv) { - u_iop = use_large_scale_wind ? m_iop->get_iop_field("u_ls").get_view() - : m_iop->get_iop_field("u").get_view(); - v_iop = use_large_scale_wind ? m_iop->get_iop_field("v_ls").get_view() - : m_iop->get_iop_field("v").get_view(); - } - - // Team policy and workspace manager for eamxx - const auto policy_iop = ESU::get_default_team_policy(nelem*NGP*NGP, NLEV); - - // TODO: Create a memory buffer for this class - // and add the below WSM and views - WorkspaceMgr iop_wsm(NLEVI, 7+qsize, policy_iop); - view_Nd - temperature("temperature", nelem, NGP, NGP, NLEV); - - // Lambda for computing temperature - auto compute_temperature = [&] () { - Kokkos::parallel_for("compute_temperature_for_iop", policy_iop, KOKKOS_LAMBDA (const KT::MemberType& team) { - const int ie = team.league_rank()/(NGP*NGP); - const int igp = (team.league_rank()/NGP)%NGP; - const int jgp = team.league_rank()%NGP; - - // Get temp views from workspace - auto ws = iop_wsm.get_workspace(team); - uview_1d pmid; - ws.take_many_contiguous_unsafe<1>({"pmid"},{&pmid}); - - auto ps_i = ps_dyn(ie, igp, jgp); - auto dp3d_i = ekat::subview(dp3d_dyn, ie, igp, jgp); - auto vtheta_dp_i = ekat::subview(vtheta_dp_dyn, ie, igp, jgp); - auto qv_i = ekat::subview(Q_dyn, ie, 0, igp, jgp); - auto temperature_i = ekat::subview(temperature, ie, igp, jgp); - - // Compute reference pressures and layer thickness. - // TODO: Allow geometry data to allocate packsize - auto s_pmid = ekat::scalarize(pmid); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, total_levels), [&](const int& k) { - s_pmid(k) = hyam(k)*ps0 + hybm(k)*ps_i; - }); - team.team_barrier(); - - // Compute temperature from virtual potential temperature - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, NLEV), [&] (const int k) { - auto T_val = vtheta_dp_i(k); - T_val /= dp3d_i(k); - T_val = PF::calculate_temperature_from_virtual_temperature(T_val,qv_i(k)); - temperature_i(k) = PF::calculate_T_from_theta(T_val,pmid(k)); - }); - - // Release WS views - ws.release_many_contiguous<1>({&pmid}); - }); - }; - - // Preprocess some homme states to get temperature - compute_temperature(); - Kokkos::fence(); - - // Apply IOP forcing - Kokkos::parallel_for("apply_iop_forcing", policy_iop, KOKKOS_LAMBDA (const KT::MemberType& team) { - const int ie = team.league_rank()/(NGP*NGP); - const int igp = (team.league_rank()/NGP)%NGP; - const int jgp = team.league_rank()%NGP; - - // Get temp views from workspace - auto ws = iop_wsm.get_workspace(team); - uview_1d pmid, pint, pdel; - ws.take_many_contiguous_unsafe<3>({"pmid", "pint", "pdel"}, - {&pmid, &pint, &pdel}); - - auto ps_i = ps_dyn(ie, igp, jgp); - auto u_i = ekat::subview(v_dyn, ie, 0, igp, jgp); - auto v_i = ekat::subview(v_dyn, ie, 1, igp, jgp); - auto temperature_i = ekat::subview(temperature, ie, igp, jgp); - auto qv_i = ekat::subview(Q_dyn, ie, 0, igp, jgp); - auto Q_i = Kokkos::subview(Q_dyn, ie, Kokkos::ALL(), igp, jgp, Kokkos::ALL()); - - // Compute reference pressures and layer thickness. - // TODO: Allow geometry data to allocate packsize - auto s_pmid = ekat::scalarize(pmid); - auto s_pint = ekat::scalarize(pint); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, total_levels+1), [&](const int& k) { - s_pint(k) = hyai(k)*ps0 + hybi(k)*ps_i; - if (k < total_levels) { - s_pmid(k) = hyam(k)*ps0 + hybm(k)*ps_i; - } - }); - team.team_barrier(); - ColOps::compute_midpoint_delta(team, total_levels, pint, pdel); - team.team_barrier(); - - if (iop_dosubsidence) { - // Compute subsidence due to large-scale forcing - advance_iop_subsidence(team, total_levels, dt, ps_i, pmid, pint, pdel, omega, ws, u_i, v_i, temperature_i, Q_i); - } - - // Update T and qv according to large scale forcing as specified in IOP file. - advance_iop_forcing(team, total_levels, dt, divT, divq, temperature_i, qv_i); - - if (iop_coriolis) { - // Apply coriolis forcing to u and v winds - iop_apply_coriolis(team, total_levels, dt, lat, u_ls, v_ls, u_i, v_i); - } - - // Release WS views - ws.release_many_contiguous<3>({&pmid, &pint, &pdel}); - }); - Kokkos::fence(); - - // Postprocess homme states Qdp and vtheta_dp - Kokkos::parallel_for("compute_qdp_and_vtheta_dp", policy_iop, KOKKOS_LAMBDA (const KT::MemberType& team) { - const int ie = team.league_rank()/(NGP*NGP); - const int igp = (team.league_rank()/NGP)%NGP; - const int jgp = team.league_rank()%NGP; - - // Get temp views from workspace - auto ws = iop_wsm.get_workspace(team); - uview_1d pmid, pint, pdel; - ws.take_many_contiguous_unsafe<3>({"pmid", "pint", "pdel"}, - {&pmid, &pint, &pdel}); - - auto ps_i = ps_dyn(ie, igp, jgp); - auto dp3d_i = ekat::subview(dp3d_dyn, ie, igp, jgp); - auto vtheta_dp_i = ekat::subview(vtheta_dp_dyn, ie, igp, jgp); - auto qv_i = ekat::subview(Q_dyn, ie, 0, igp, jgp); - auto Q_i = Kokkos::subview(Q_dyn, ie, Kokkos::ALL(), igp, jgp, Kokkos::ALL()); - auto Qdp_i = Kokkos::subview(Qdp_dyn, ie, Kokkos::ALL(), igp, jgp, Kokkos::ALL()); - auto temperature_i = ekat::subview(temperature, ie, igp, jgp); - - // Compute reference pressures and layer thickness. - // TODO: Allow geometry data to allocate packsize - auto s_pmid = ekat::scalarize(pmid); - auto s_pint = ekat::scalarize(pint); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, total_levels+1), [&](const int& k) { - s_pint(k) = hyai(k)*ps0 + hybi(k)*ps_i; - if (k < total_levels) { - s_pmid(k) = hyam(k)*ps0 + hybm(k)*ps_i; - } - }); - - team.team_barrier(); - - // Compute Qdp from updated Q - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, NLEV*qsize), [&] (const int k) { - const int ilev = k/qsize; - const int q = k%qsize; - - Qdp_i(q, ilev) = Q_i(q, ilev)*dp3d_i(ilev); - // For BFB on restarts, Q needs to be updated after we compute Qdp - Q_i(q, ilev) = Qdp_i(q, ilev)/dp3d_i(ilev); - }); - team.team_barrier(); - - // Convert updated temperature back to psuedo density virtual potential temperature - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, NLEV), [&] (const int k) { - const auto th = PF::calculate_theta_from_T(temperature_i(k),pmid(k)); - vtheta_dp_i(k) = PF::calculate_virtual_temperature(th,qv_i(k))*dp3d_i(k); - }); - - // Release WS views - ws.release_many_contiguous<3>({&pmid, &pint, &pdel}); - }); - - if (iop_nudge_tq or iop_nudge_uv) { - // Nudge the domain based on the domain mean - // and observed quantities of T, Q, u, and - - if (iop_nudge_tq) { - // Compute temperature - compute_temperature(); - Kokkos::fence(); - } - - // Compute domain mean of qv, temperature, u, and v - - // TODO: add to local mem buffer - view_1d qv_mean, t_mean, u_mean, v_mean; - if (iop_nudge_tq) { - qv_mean = view_1d("u_mean", NLEV), - t_mean = view_1d("v_mean", NLEV); - } - if (iop_nudge_uv){ - u_mean = view_1d("u_mean", NLEV), - v_mean = view_1d("v_mean", NLEV); - } - - const auto qv_mean_h = Kokkos::create_mirror_view(qv_mean); - const auto t_mean_h = Kokkos::create_mirror_view(t_mean); - const auto u_mean_h = Kokkos::create_mirror_view(u_mean); - const auto v_mean_h = Kokkos::create_mirror_view(v_mean); - - for (int k=0; kget_num_global_dofs(); - t_mean_k /= m_dyn_grid->get_num_global_dofs(); - } - if (iop_nudge_uv){ - Real& u_mean_k = u_mean_h(k/Pack::n)[k%Pack::n]; - Real& v_mean_k = v_mean_h(k/Pack::n)[k%Pack::n]; - Kokkos::parallel_reduce("compute_domain_means_uv", - nelem*NGP*NGP, - KOKKOS_LAMBDA (const int idx, Real& u_sum, Real& v_sum) { - const int ie = idx/(NGP*NGP); - const int igp = (idx/NGP)%NGP; - const int jgp = idx%NGP; - - u_sum += v_dyn(ie, 0, igp, jgp, k/Pack::n)[k%Pack::n]; - v_sum += v_dyn(ie, 1, igp, jgp, k/Pack::n)[k%Pack::n]; - }, - u_mean_k, - v_mean_k); - - m_comm.all_reduce(&u_mean_k, 1, MPI_SUM); - m_comm.all_reduce(&v_mean_k, 1, MPI_SUM); - - u_mean_k /= m_dyn_grid->get_num_global_dofs(); - v_mean_k /= m_dyn_grid->get_num_global_dofs(); - } - } - Kokkos::deep_copy(qv_mean, qv_mean_h); - Kokkos::deep_copy(t_mean, t_mean_h); - Kokkos::deep_copy(u_mean, u_mean_h); - Kokkos::deep_copy(v_mean, v_mean_h); - - // Apply relaxation - const auto rtau = std::max(dt, iop_nudge_tscale); - Kokkos::parallel_for("apply_domain_relaxation", - policy_iop, - KOKKOS_LAMBDA (const KT::MemberType& team) { - - const int ie = team.league_rank()/(NGP*NGP); - const int igp = (team.league_rank()/NGP)%NGP; - const int jgp = team.league_rank()%NGP; - - // Get temp views from workspace - auto ws = iop_wsm.get_workspace(team); - uview_1d pmid; - ws.take_many_contiguous_unsafe<1>({"pmid"},{&pmid}); - - auto ps_i = ps_dyn(ie, igp, jgp); - auto dp3d_i = ekat::subview(dp3d_dyn, ie, igp, jgp); - auto vtheta_dp_i = ekat::subview(vtheta_dp_dyn, ie, igp, jgp); - auto qv_i = ekat::subview(Q_dyn, ie, 0, igp, jgp); - auto temperature_i = ekat::subview(temperature, ie, igp, jgp); - auto u_i = ekat::subview(v_dyn, ie, 0, igp, jgp); - auto v_i = ekat::subview(v_dyn, ie, 1, igp, jgp); - - // Compute reference pressures and layer thickness. - // TODO: Allow geometry data to allocate packsize - auto s_pmid = ekat::scalarize(pmid); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, total_levels), [&](const int& k) { - s_pmid(k) = hyam(k)*ps0 + hybm(k)*ps_i; - }); - team.team_barrier(); - - if (iop_nudge_tq or iop_nudge_uv) { - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, NLEV), [&](const int& k) { - if (iop_nudge_tq) { - // Restrict nudging of T and qv to certain levels if requested by user - // IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high - // is in units of [hPa], thus convert iop_nudge_tq_low/high - Mask nudge_level(false); - int max_size = hyam.size(); - for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) { - const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop; - nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100 - and - pressure_from_iop >= iop_nudge_tq_high*100); - } - - qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0); - temperature_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0); - - // Convert updated temperature back to virtual potential temperature - const auto th = PF::calculate_theta_from_T(temperature_i(k),pmid(k)); - vtheta_dp_i(k) = PF::calculate_virtual_temperature(th,qv_i(k))*dp3d_i(k); - } - if (iop_nudge_uv) { - u_i(k).update(u_mean(k) - u_iop(k), -dt/rtau, 1.0); - v_i(k).update(v_mean(k) - v_iop(k), -dt/rtau, 1.0); - } - }); - } - - // Release WS views - ws.release_many_contiguous<1>({&pmid}); - }); - } -} - -} // namespace scream diff --git a/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.cpp b/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.cpp index 8b7495ffd73f..0ccbbbc4d02a 100644 --- a/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.cpp +++ b/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.cpp @@ -426,11 +426,6 @@ void HommeDynamics::initialize_impl (const RunType run_type) if (run_type==RunType::Initial) { initialize_homme_state (); } else { - if (m_iop) { - // We need to reload IOP data after restarting - m_iop->read_iop_file_data(timestamp()); - } - restart_homme_state (); } @@ -669,10 +664,6 @@ void HommeDynamics::homme_post_process (const double dt) { get_internal_field("w_int_dyn").get_header().get_alloc_properties().reset_subview_idx(tl.n0); } - if (m_iop) { - apply_iop_forcing(dt); - } - if (fv_phys_active()) { fv_phys_post_process(); // Apply Rayleigh friction to update temperature and horiz_winds diff --git a/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.hpp b/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.hpp index 93dff0cd72e3..b3e01f8aa302 100644 --- a/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.hpp +++ b/components/eamxx/src/dynamics/homme/eamxx_homme_process_interface.hpp @@ -109,44 +109,6 @@ class HommeDynamics : public AtmosphereProcess void rayleigh_friction_init (); void rayleigh_friction_apply (const Real dt) const; - // IOP functions - void apply_iop_forcing(const Real dt); - - KOKKOS_FUNCTION - static void advance_iop_subsidence(const KT::MemberType& team, - const int nlevs, - const Real dt, - const Real ps, - const view_1d& pmid, - const view_1d& pint, - const view_1d& pdel, - const view_1d& omega, - const Workspace& workspace, - const view_1d& u, - const view_1d& v, - const view_1d& T, - const view_2d& Q); - - KOKKOS_FUNCTION - static void advance_iop_forcing(const KT::MemberType& team, - const int nlevs, - const Real dt, - const view_1d& divT, - const view_1d& divq, - const view_1d& T, - const view_1d& qv); - - - KOKKOS_FUNCTION - static void iop_apply_coriolis(const KT::MemberType& team, - const int nlevs, - const Real dt, - const Real lat, - const view_1d& u_ls, - const view_1d& v_ls, - const view_1d& u, - const view_1d& v); - public: // Fast boolean function returning whether Physics PGN is being used. bool fv_phys_active() const; diff --git a/components/eamxx/src/mct_coupling/CMakeLists.txt b/components/eamxx/src/mct_coupling/CMakeLists.txt index 308cd1776234..39f864e728aa 100644 --- a/components/eamxx/src/mct_coupling/CMakeLists.txt +++ b/components/eamxx/src/mct_coupling/CMakeLists.txt @@ -38,6 +38,7 @@ set (SCREAM_LIBS eamxx_cosp cld_fraction spa + iop_forcing nudging diagnostics tms diff --git a/components/eamxx/src/physics/CMakeLists.txt b/components/eamxx/src/physics/CMakeLists.txt index e0e89e60f80d..f9beda35a20c 100644 --- a/components/eamxx/src/physics/CMakeLists.txt +++ b/components/eamxx/src/physics/CMakeLists.txt @@ -8,8 +8,10 @@ add_subdirectory(p3) if (SCREAM_DOUBLE_PRECISION) add_subdirectory(rrtmgp) add_subdirectory(cosp) + add_subdirectory(tms) + add_subdirectory(iop_forcing) else() - message(STATUS "WARNING: RRTMGP and COSP only supported for double precision builds; skipping") + message(STATUS "WARNING: RRTMGP, COSP, TMS, and IOPForcing only supported for double precision builds; skipping") endif() add_subdirectory(shoc) add_subdirectory(cld_fraction) @@ -21,8 +23,4 @@ add_subdirectory(nudging) if (SCREAM_ENABLE_MAM) add_subdirectory(mam) endif() -if (SCREAM_DOUBLE_PRECISION) - add_subdirectory(tms) -else() - message(STATUS "WARNING: TMS only supported for double precision builds; skipping") -endif() + diff --git a/components/eamxx/src/physics/iop_forcing/CMakeLists.txt b/components/eamxx/src/physics/iop_forcing/CMakeLists.txt new file mode 100644 index 000000000000..093ceac73c91 --- /dev/null +++ b/components/eamxx/src/physics/iop_forcing/CMakeLists.txt @@ -0,0 +1,5 @@ +add_library(iop_forcing eamxx_iop_forcing_process_interface.cpp) +target_compile_definitions(iop_forcing PUBLIC EAMXX_HAS_IOP_FORCING) +target_link_libraries(iop_forcing physics_share scream_share) + +target_link_libraries(eamxx_physics INTERFACE iop_forcing) diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp new file mode 100644 index 000000000000..9aac3fdf03ac --- /dev/null +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp @@ -0,0 +1,546 @@ +#include "physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp" + +#include "share/property_checks/field_within_interval_check.hpp" + +namespace scream +{ +// ========================================================================================= +void IOPForcing::set_grids(const std::shared_ptr grids_manager) +{ + using namespace ekat::units; + + m_grid = grids_manager->get_grid("Physics"); + const auto& grid_name = m_grid->name(); + + m_num_cols = m_grid->get_num_local_dofs(); // Number of columns on this rank + m_num_levs = m_grid->get_num_vertical_levels(); // Number of levels per column + + // Define the different field layouts that will be used for this process + FieldLayout scalar2d = m_grid->get_2d_scalar_layout(); + FieldLayout scalar3d_mid = m_grid->get_3d_scalar_layout(true); + FieldLayout vector3d_mid = m_grid->get_3d_vector_layout(true,2); + + constexpr int pack_size = Pack::n; + + add_field("ps", scalar2d, Pa, grid_name); + + add_field("horiz_winds", vector3d_mid, m/s, grid_name, pack_size); + add_field("T_mid", scalar3d_mid, K, grid_name, pack_size); + + add_tracer("qv", m_grid, kg/kg, pack_size); + add_group("tracers", grid_name, pack_size, Bundling::Required); + + // Sanity check that iop data manager is setup by driver + EKAT_REQUIRE_MSG(m_iop_data_manager, + "Error! IOPDataManager not setup by driver, but IOPForcing" + "being used as an ATM process.\n"); + + // Compute number of buffer views + const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); + const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); + if (iop_nudge_tq) m_buffer.num_1d_scalar_nlev += 2; + if (iop_nudge_uv) m_buffer.num_1d_scalar_nlev += 2; +} +// ========================================================================================= +void IOPForcing:: +set_computed_group_impl (const FieldGroup& group) +{ + EKAT_REQUIRE_MSG(group.m_info->size() >= 1, + "Error! IOPForcing requires at least qv as tracer input.\n"); + + const auto& name = group.m_info->m_group_name; + + EKAT_REQUIRE_MSG(name=="tracers", + "Error! IOPForcing was not expecting a field group called '" << name << "\n"); + + EKAT_REQUIRE_MSG(group.m_info->m_bundled, + "Error! IOPForcing expects bundled fields for tracers.\n"); + + m_num_tracers = group.m_info->size(); +} +// ========================================================================================= +size_t IOPForcing::requested_buffer_size_in_bytes() const +{ + const int nlev_packs = ekat::npack(m_num_levs); + const int nlevi_packs = ekat::npack(m_num_levs+1); + + const size_t temp_view_bytes = + m_buffer.num_1d_scalar_nlev*nlev_packs*sizeof(Pack); + + // Number of bytes needed by the WorkspaceManager passed to shoc_main + const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs); + const size_t wsm_bytes = WorkspaceMgr::get_total_bytes_needed(nlevi_packs, 7+m_num_tracers, policy); + + return temp_view_bytes + wsm_bytes; +} +// ========================================================================================= +void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager) +{ + EKAT_REQUIRE_MSG(buffer_manager.allocated_bytes() >= requested_buffer_size_in_bytes(), + "Error! Buffers size not sufficient.\n"); + + const int nlev_packs = ekat::npack(m_num_levs); + const int nlevi_packs = ekat::npack(m_num_levs+1); + + Pack* mem = reinterpret_cast(buffer_manager.get_memory()); + + // Temp view data + using mean_view_t = decltype(m_buffer.qv_mean); + const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); + const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); + if (iop_nudge_tq) { + m_buffer.qv_mean = mean_view_t(mem, nlev_packs); + mem += m_buffer.qv_mean.size(); + m_buffer.t_mean = mean_view_t(mem, nlev_packs); + mem += m_buffer.t_mean.size(); + } + if (iop_nudge_uv) { + m_buffer.u_mean = mean_view_t(mem, nlev_packs); + mem += m_buffer.u_mean.size(); + m_buffer.v_mean = mean_view_t(mem, nlev_packs); + mem += m_buffer.v_mean.size(); + } + + // WSM data + m_buffer.wsm_data = mem; + + const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs); + const size_t wsm_npacks = WorkspaceMgr::get_total_bytes_needed(nlevi_packs, 7+m_num_tracers, policy)/sizeof(Pack); + mem += wsm_npacks; + + size_t used_mem = (reinterpret_cast(mem) - buffer_manager.get_memory())*sizeof(Real); + EKAT_REQUIRE_MSG(used_mem==requested_buffer_size_in_bytes(), "Error! Used memory != requested memory for IOPForcing.\n"); +} +// ========================================================================================= +void IOPForcing::initialize_impl (const RunType run_type) +{ + // Set field property checks for the fields in this process + using Interval = FieldWithinIntervalCheck; + add_postcondition_check(get_field_out("T_mid"),m_grid,100.0,500.0,false); + add_postcondition_check(get_field_out("horiz_winds"),m_grid,-400.0,400.0,false); + // For qv, ensure it doesn't get negative, by allowing repair of any neg value. + // TODO: use a repairable lb that clips only "small" negative values + add_postcondition_check(get_field_out("qv"),m_grid,0,0.2,true); + + // Setup WSM for internal local variables + const auto nlevi_packs = ekat::npack(m_num_levs+1); + const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs); + m_workspace_mgr.setup(m_buffer.wsm_data, nlevi_packs, 7+m_num_tracers, policy); +} +// ========================================================================================= +KOKKOS_FUNCTION +void IOPForcing:: +advance_iop_subsidence(const MemberType& team, + const int nlevs, + const Real dt, + const Real ps, + const view_1d& ref_p_mid, + const view_1d& ref_p_int, + const view_1d& ref_p_del, + const view_1d& omega, + const Workspace& workspace, + const view_1d& u, + const view_1d& v, + const view_1d& T, + const view_2d& Q) +{ + constexpr Real Rair = C::Rair; + constexpr Real Cpair = C::Cpair; + + const auto n_q_tracers = Q.extent_int(0); + const auto nlev_packs = ekat::npack(nlevs); + + // Get some temporary views from WS + uview_1d omega_int, delta_u, delta_v, delta_T, tmp; + workspace.take_many_contiguous_unsafe<4>({"omega_int", "delta_u", "delta_v", "delta_T"}, + {&omega_int, &delta_u, &delta_v, &delta_T}); + const auto delta_Q_slot = workspace.take_macro_block("delta_Q", n_q_tracers); + uview_2d delta_Q(delta_Q_slot.data(), n_q_tracers, nlev_packs); + + auto s_ref_p_mid = ekat::scalarize(ref_p_mid); + auto s_omega = ekat::scalarize(omega); + auto s_delta_u = ekat::scalarize(delta_u); + auto s_delta_v = ekat::scalarize(delta_v); + auto s_delta_T = ekat::scalarize(delta_T); + auto s_delta_Q = ekat::scalarize(delta_Q); + auto s_omega_int = ekat::scalarize(omega_int); + + // Compute omega on the interface grid by using a weighted average in pressure + const int pack_begin = 1/Pack::n, pack_end = (nlevs-1)/Pack::n; + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pack_begin, pack_end+1), [&] (const int k){ + auto range_pack = ekat::range(k*Pack::n); + range_pack.set(range_pack<1, 1); + Pack ref_p_mid_k, ref_p_mid_km1, omega_k, omega_km1; + ekat::index_and_shift<-1>(s_ref_p_mid, range_pack, ref_p_mid_k, ref_p_mid_km1); + ekat::index_and_shift<-1>(s_omega, range_pack, omega_k, omega_km1); + + const auto weight = (ref_p_int(k) - ref_p_mid_km1)/(ref_p_mid_k - ref_p_mid_km1); + omega_int(k).set(range_pack>=1 and range_pack<=nlevs-1, + weight*omega_k + (1-weight)*omega_km1); + }); + omega_int(0)[0] = 0; + omega_int(nlevs/Pack::n)[nlevs%Pack::n] = 0; + + // Compute delta views for u, v, T, and Q (e.g., u(k+1) - u(k), k=0,...,nlevs-2) + ColOps::compute_midpoint_delta(team, nlevs-1, u, delta_u); + ColOps::compute_midpoint_delta(team, nlevs-1, v, delta_v); + ColOps::compute_midpoint_delta(team, nlevs-1, T, delta_T); + for (int iq=0; iq(k*Pack::n); + const auto at_top = range_pack==0; + const auto not_at_top = not at_top; + const auto at_bot = range_pack==nlevs-1; + const auto not_at_bot = not at_bot; + const bool any_at_top = at_top.any(); + const bool any_at_bot = at_bot.any(); + + // Get delta(k-1) packs. The range pack should not + // contain index 0 (so that we don't attempt to access + // k=-1 index) or index > nlevs-2 (since delta_* views + // are size nlevs-1). + auto range_pack_for_m1_shift = range_pack; + range_pack_for_m1_shift.set(range_pack<1, 1); + range_pack_for_m1_shift.set(range_pack>nlevs-2, nlevs-2); + Pack delta_u_k, delta_u_km1, + delta_v_k, delta_v_km1, + delta_T_k, delta_T_km1; + ekat::index_and_shift<-1>(s_delta_u, range_pack_for_m1_shift, delta_u_k, delta_u_km1); + ekat::index_and_shift<-1>(s_delta_v, range_pack_for_m1_shift, delta_v_k, delta_v_km1); + ekat::index_and_shift<-1>(s_delta_T, range_pack_for_m1_shift, delta_T_k, delta_T_km1); + + // At the top and bottom of the model, set the end points for + // delta_*_k and delta_*_km1 to be the first and last entries + // of delta_*, respectively. + if (any_at_top) { + delta_u_k.set(at_top, s_delta_u(0)); + delta_v_k.set(at_top, s_delta_v(0)); + delta_T_k.set(at_top, s_delta_T(0)); + } + if (any_at_bot) { + delta_u_km1.set(at_bot, s_delta_u(nlevs-2)); + delta_v_km1.set(at_bot, s_delta_v(nlevs-2)); + delta_T_km1.set(at_bot, s_delta_T(nlevs-2)); + } + + // Get omega_int(k+1) pack. The range pack should not + // contain index > nlevs-1 (since omega_int is size nlevs+1). + auto range_pack_for_p1_shift = range_pack; + range_pack_for_p1_shift.set(range_pack>nlevs-1, nlevs-1); + Pack omega_int_k, omega_int_kp1; + ekat::index_and_shift<1>(s_omega_int, range_pack, omega_int_k, omega_int_kp1); + + const auto fac = (dt/2)/ref_p_del(k); + + // Update u + u(k).update(not_at_bot, fac*omega_int_kp1*delta_u_k, -1, 1); + u(k).update(not_at_top, fac*omega_int_k*delta_u_km1, -1, 1); + + // Update v + v(k).update(not_at_bot, fac*omega_int_kp1*delta_v_k, -1, 1); + v(k).update(not_at_top, fac*omega_int_k*delta_v_km1, -1, 1); + + // Before updating T, first scale using thermal + // expansion term due to LS vertical advection + T(k) *= 1 + (dt*Rair/Cpair)*omega(k)/ref_p_mid(k); + + // Update T + T(k).update(not_at_bot, fac*omega_int_kp1*delta_T_k, -1, 1); + T(k).update(not_at_top, fac*omega_int_k*delta_T_km1, -1, 1); + + // Update Q + Pack delta_tracer_k, delta_tracer_km1; + for (int iq=0; iq(s_delta_tracer, range_pack_for_m1_shift, delta_tracer_k, delta_tracer_km1); + if (any_at_top) delta_tracer_k.set(at_top, s_delta_tracer(0)); + if (any_at_bot) delta_tracer_km1.set(at_bot, s_delta_tracer(nlevs-2)); + + Q(iq, k).update(not_at_bot, fac*omega_int_kp1*delta_tracer_k, -1, 1); + Q(iq, k).update(not_at_top, fac*omega_int_k*delta_tracer_km1, -1, 1); + } + }); + + // Release WS views + workspace.release_macro_block(delta_Q_slot, n_q_tracers); + workspace.release_many_contiguous<4>({&omega_int, &delta_u, &delta_v, &delta_T}); +} +// ========================================================================================= +KOKKOS_FUNCTION +void IOPForcing:: +advance_iop_forcing(const MemberType& team, + const int nlevs, + const Real dt, + const view_1d& divT, + const view_1d& divq, + const view_1d& T, + const view_1d& qv) +{ + const auto nlev_packs = ekat::npack(nlevs); + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&] (const int k) { + T(k).update(divT(k), dt, 1.0); + qv(k).update(divq(k), dt, 1.0); + }); +} +// ========================================================================================= +KOKKOS_FUNCTION +void IOPForcing:: +iop_apply_coriolis(const MemberType& team, + const int nlevs, + const Real dt, + const Real lat, + const view_1d& u_ls, + const view_1d& v_ls, + const view_1d& u, + const view_1d& v) +{ + constexpr Real pi = C::Pi; + constexpr Real earth_rotation = C::omega; + + // Compute coriolis force + const auto fcor = 2*earth_rotation*std::sin(lat*pi/180); + + const auto nlev_packs = ekat::npack(nlevs); + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&] (const int k) { + const auto u_cor = v(k) - v_ls(k); + const auto v_cor = u(k) - u_ls(k); + u(k).update(u_cor, dt*fcor, 1.0); + v(k).update(v_cor, -dt*fcor, 1.0); + }); +} +// ========================================================================================= +void IOPForcing::run_impl (const double dt) +{ + // Pack dimensions + const auto nlev_packs = ekat::npack(m_num_levs); + + // Hybrid coord values + const auto ps0 = C::P0; + const auto hyam = m_grid->get_geometry_data("hyam").get_view(); + const auto hybm = m_grid->get_geometry_data("hybm").get_view(); + const auto hyai = m_grid->get_geometry_data("hyai").get_view(); + const auto hybi = m_grid->get_geometry_data("hybi").get_view(); + + // Get FM fields + const auto ps = get_field_in("ps").get_view(); + const auto horiz_winds = get_field_out("horiz_winds").get_view(); + const auto T_mid = get_field_out("T_mid").get_view(); + const auto qv = get_field_out("qv").get_view(); + const auto Q = get_group_out("tracers").m_bundle->get_view(); + + // Load data from IOP files, if necessary + m_iop_data_manager->read_iop_file_data(timestamp()); + + // Define local IOP param values + const auto iop_dosubsidence = m_iop_data_manager->get_params().get("iop_dosubsidence"); + const auto iop_coriolis = m_iop_data_manager->get_params().get("iop_coriolis"); + const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); + const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); + const auto use_large_scale_wind = m_iop_data_manager->get_params().get("use_large_scale_wind"); + const auto use_3d_forcing = m_iop_data_manager->get_params().get("use_3d_forcing"); + const auto target_lat = m_iop_data_manager->get_params().get("target_latitude"); + const auto iop_nudge_tscale = m_iop_data_manager->get_params().get("iop_nudge_tscale"); + const auto iop_nudge_tq_low = m_iop_data_manager->get_params().get("iop_nudge_tq_low"); + const auto iop_nudge_tq_high = m_iop_data_manager->get_params().get("iop_nudge_tq_high"); + + // Define local IOP field views + const Real ps_iop = m_iop_data_manager->get_iop_field("Ps").get_view()(); + view_1d omega, divT, divq, u_ls, v_ls, qv_iop, t_iop, u_iop, v_iop; + divT = use_3d_forcing ? m_iop_data_manager->get_iop_field("divT3d").get_view() + : m_iop_data_manager->get_iop_field("divT").get_view(); + divq = use_3d_forcing ? m_iop_data_manager->get_iop_field("divq3d").get_view() + : m_iop_data_manager->get_iop_field("divq").get_view(); + if (iop_dosubsidence) { + omega = m_iop_data_manager->get_iop_field("omega").get_view(); + } + if (iop_coriolis) { + u_ls = m_iop_data_manager->get_iop_field("u_ls").get_view(); + v_ls = m_iop_data_manager->get_iop_field("v_ls").get_view(); + } + if (iop_nudge_tq) { + qv_iop = m_iop_data_manager->get_iop_field("q").get_view(); + t_iop = m_iop_data_manager->get_iop_field("T").get_view(); + } + if (iop_nudge_uv) { + u_iop = use_large_scale_wind ? m_iop_data_manager->get_iop_field("u_ls").get_view() + : m_iop_data_manager->get_iop_field("u").get_view(); + v_iop = use_large_scale_wind ? m_iop_data_manager->get_iop_field("v_ls").get_view() + : m_iop_data_manager->get_iop_field("v").get_view(); + } + + // Team policy and workspace manager for eamxx + const auto policy_iop = ESU::get_default_team_policy(m_num_cols, nlev_packs); + + // Reset internal WSM variables. + m_workspace_mgr.reset_internals(); + + // Avoid implicit capture of this + auto wsm = m_workspace_mgr; + auto num_levs = m_num_levs; + + // Apply IOP forcing + Kokkos::parallel_for("apply_iop_forcing", policy_iop, KOKKOS_LAMBDA (const MemberType& team) { + const int icol = team.league_rank(); + + auto ps_i = ps(icol); + auto u_i = Kokkos::subview(horiz_winds, icol, 0, Kokkos::ALL()); + auto v_i = Kokkos::subview(horiz_winds, icol, 1, Kokkos::ALL()); + auto T_mid_i = ekat::subview(T_mid, icol); + auto qv_i = ekat::subview(qv, icol); + auto Q_i = Kokkos::subview(Q, icol, Kokkos::ALL(), Kokkos::ALL()); + + auto ws = wsm.get_workspace(team); + uview_1d ref_p_mid, ref_p_int, ref_p_del; + ws.take_many_contiguous_unsafe<3>({"ref_p_mid", "ref_p_int", "ref_p_del"}, + {&ref_p_mid, &ref_p_int, &ref_p_del}); + + // Compute reference pressures and layer thickness. + // TODO: Allow geometry data to allocate packsize + auto s_ref_p_mid = ekat::scalarize(ref_p_mid); + auto s_ref_p_int = ekat::scalarize(ref_p_int); + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, num_levs+1), [&](const int& k) { + s_ref_p_int(k) = hyai(k)*ps0 + hybi(k)*ps_i; + if (k < num_levs) { + s_ref_p_mid(k) = hyam(k)*ps0 + hybm(k)*ps_i; + } + }); + team.team_barrier(); + ColOps::compute_midpoint_delta(team, num_levs, ref_p_int, ref_p_del); + team.team_barrier(); + + if (iop_dosubsidence) { + // Compute subsidence due to large-scale forcing + advance_iop_subsidence(team, num_levs, dt, ps_i, ref_p_mid, ref_p_int, ref_p_del, omega, ws, u_i, v_i, T_mid_i, Q_i); + } + + // Update T and qv according to large scale forcing as specified in IOP file. + advance_iop_forcing(team, num_levs, dt, divT, divq, T_mid_i, qv_i); + + if (iop_coriolis) { + // Apply coriolis forcing to u and v winds + iop_apply_coriolis(team, num_levs, dt, target_lat, u_ls, v_ls, u_i, v_i); + } + + // Release WS views + ws.release_many_contiguous<3>({&ref_p_mid, &ref_p_int, &ref_p_del}); + }); + Kokkos::fence(); + + // Nudge the domain based on the domain mean + // and observed quantities of T, Q, u, and v + if (iop_nudge_tq or iop_nudge_uv) { + // Compute domain mean of qv, T_mid, u, and v + const auto qv_mean = m_buffer.qv_mean; + const auto t_mean = m_buffer.t_mean; + const auto u_mean = m_buffer.u_mean; + const auto v_mean = m_buffer.v_mean; + + const auto qv_mean_h = Kokkos::create_mirror_view(qv_mean); + const auto t_mean_h = Kokkos::create_mirror_view(t_mean); + const auto u_mean_h = Kokkos::create_mirror_view(u_mean); + const auto v_mean_h = Kokkos::create_mirror_view(v_mean); + + const auto num_global_cols = m_grid->get_num_global_dofs(); + for (int k=0; k ref_p_mid; + ws.take_many_contiguous_unsafe<1>({"ref_p_mid"},{&ref_p_mid}); + + // Compute reference pressures and layer thickness. + // TODO: Allow geometry data to allocate packsize + auto s_ref_p_mid = ekat::scalarize(ref_p_mid); + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, num_levs), [&](const int& k) { + s_ref_p_mid(k) = hyam(k)*ps0 + hybm(k)*ps_i; + }); + team.team_barrier(); + + if (iop_nudge_tq or iop_nudge_uv) { + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) { + if (iop_nudge_tq) { + // Restrict nudging of T and qv to certain levels if requested by user + // IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high + // is in units of [hPa], thus convert iop_nudge_tq_low/high + Mask nudge_level(false); + int max_size = hyam.size(); + for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) { + const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop; + nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100 + and + pressure_from_iop >= iop_nudge_tq_high*100); + } + + qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0); + T_mid_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0); + } + if (iop_nudge_uv) { + u_i(k).update(u_mean(k) - u_iop(k), -dt/rtau, 1.0); + v_i(k).update(v_mean(k) - v_iop(k), -dt/rtau, 1.0); + } + }); + } + + // Release WS views + ws.release_many_contiguous<1>({&ref_p_mid}); + }); + } +} +// ========================================================================================= +} // namespace scream diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp new file mode 100644 index 000000000000..535043f401e5 --- /dev/null +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp @@ -0,0 +1,153 @@ +#ifndef SCREAM_IOP_FORCING_HPP +#define SCREAM_IOP_FORCING_HPP + +#include "ekat/ekat_parameter_list.hpp" +#include "ekat/ekat_workspace.hpp" + +#include "share/atm_process/atmosphere_process.hpp" +#include "share/atm_process/ATMBufferManager.hpp" +#include "share/util/scream_column_ops.hpp" + +#include "physics/share/physics_constants.hpp" + +#include + +namespace scream +{ +/* + * The class responsible for running EAMxx with an intensive + * observation period (IOP). + * + * The AD should store exactly ONE instance of this class stored + * in its list of subcomponents (the AD should make sure of this). + * + * Currently the only use case is the doubly + * periodic model (DP-SCREAM). + */ + +class IOPForcing : public scream::AtmosphereProcess +{ + // Typedefs for process + using KT = ekat::KokkosTypes; + using ESU = ekat::ExeSpaceUtils; + using Pack = ekat::Pack; + using IntPack = ekat::Pack; + using Mask = ekat::Mask; + using WorkspaceMgr = ekat::WorkspaceManager; + using Workspace = WorkspaceMgr::Workspace; + + using MemberType = KT::MemberType; + template + using view_1d = KT::view_1d; + template + using view_2d = KT::view_2d; + template + using uview_1d = ekat::Unmanaged>; + template + using uview_2d = ekat::Unmanaged>; + + using ColOps = ColumnOps; + using C = physics::Constants; + + + +public: + + // Constructors + IOPForcing (const ekat::Comm& comm, const ekat::ParameterList& params) + : AtmosphereProcess(comm, params) {} + + // The type of subcomponent + AtmosphereProcessType type () const { return AtmosphereProcessType::Physics; } + + // The name of the subcomponent + std::string name () const { return "iop"; } + + // Set the grid + void set_grids (const std::shared_ptr grids_manager); + +#ifndef KOKKOS_ENABLE_CUDA + // Cuda requires methods enclosing __device__ lambda's to be public +protected: +#endif + + void initialize_impl (const RunType run_type); + + // Compute effects of large scale subsidence on T, q, u, and v. + KOKKOS_FUNCTION + static void advance_iop_subsidence(const KT::MemberType& team, + const int nlevs, + const Real dt, + const Real ps, + const view_1d& pmid, + const view_1d& pint, + const view_1d& pdel, + const view_1d& omega, + const Workspace& workspace, + const view_1d& u, + const view_1d& v, + const view_1d& T, + const view_2d& Q); + + // Apply large scale forcing for temperature and water vapor as provided by the IOP file + KOKKOS_FUNCTION + static void advance_iop_forcing(const KT::MemberType& team, + const int nlevs, + const Real dt, + const view_1d& divT, + const view_1d& divq, + const view_1d& T, + const view_1d& qv); + + // Provide coriolis forcing to u and v winds, using large scale winds specified in IOP forcing file. + KOKKOS_FUNCTION + static void iop_apply_coriolis(const KT::MemberType& team, + const int nlevs, + const Real dt, + const Real lat, + const view_1d& u_ls, + const view_1d& v_ls, + const view_1d& u, + const view_1d& v); + + void run_impl (const double dt); + +protected: + + void finalize_impl () {} + + void set_computed_group_impl (const FieldGroup& group); + + // Computes total number of bytes needed for local variables + size_t requested_buffer_size_in_bytes() const; + + // Set local variables using memory provided by + // the ATMBufferManager + void init_buffers(const ATMBufferManager &buffer_manager); + + // Keep track of field dimensions and other scalar values + // needed in IOP + Int m_num_cols; + Int m_num_levs; + Int m_num_tracers; + + struct Buffer { + int num_1d_scalar_nlev = 0; + + uview_1d qv_mean, t_mean, u_mean, v_mean; + + Pack* wsm_data; + }; + + // Struct which contains local variables + Buffer m_buffer; + + // WSM for internal local variables + WorkspaceMgr m_workspace_mgr; + + std::shared_ptr m_grid; +}; // class IOPForcing + +} // namespace scream + +#endif // SCREAM_IOP_FORCING_HPP diff --git a/components/eamxx/src/physics/register_physics.hpp b/components/eamxx/src/physics/register_physics.hpp index 99956bb75f5e..dc1ce5745d38 100644 --- a/components/eamxx/src/physics/register_physics.hpp +++ b/components/eamxx/src/physics/register_physics.hpp @@ -41,6 +41,9 @@ #ifdef EAMXX_HAS_ML_CORRECTION #include "physics/ml_correction/eamxx_ml_correction_process_interface.hpp" #endif +#ifdef EAMXX_HAS_IOP_FORCING +#include "physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp" +#endif namespace scream { @@ -82,6 +85,9 @@ inline void register_physics () { #ifdef EAMXX_HAS_ML_CORRECTION proc_factory.register_product("MLCorrection",&create_atmosphere_process); #endif +#ifdef EAMXX_HAS_IOP_FORCING + proc_factory.register_product("iop_forcing",&create_atmosphere_process); +#endif // If no physics was enabled, silence compile warning about unused var (void) proc_factory; diff --git a/components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.cpp b/components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.cpp index 0709bb1a37ff..0121741963d9 100644 --- a/components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.cpp +++ b/components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.cpp @@ -99,16 +99,16 @@ void RRTMGPRadiation::set_grids(const std::shared_ptr grids_ m_ncol = m_grid->get_num_local_dofs(); m_nlay = m_grid->get_num_vertical_levels(); - if (m_iop) { + if (m_iop_data_manager) { // For IOP runs, we need to use the lat/lon from the // IOP files instead of the geometry data. Deep copy // to device and sync to host since both will be used. m_lat = m_grid->get_geometry_data("lat").clone(); - m_lat.deep_copy(m_iop->get_params().get("target_latitude")); + m_lat.deep_copy(m_iop_data_manager->get_params().get("target_latitude")); m_lat.sync_to_host(); m_lon = m_grid->get_geometry_data("lon").clone(); - m_lon.deep_copy(m_iop->get_params().get("target_longitude")); + m_lon.deep_copy(m_iop_data_manager->get_params().get("target_longitude")); m_lon.sync_to_host(); } else { m_lat = m_grid->get_geometry_data("lat"); diff --git a/components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.cpp b/components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.cpp index ac620e19add5..1bbb17f84806 100644 --- a/components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.cpp +++ b/components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.cpp @@ -449,7 +449,7 @@ void SHOCMacrophysics::initialize_impl (const RunType run_type) const auto ncols = m_num_cols; view_1d cell_length("cell_length", ncols); if (m_grid->has_geometry_data("dx_short")) { - // We must be running with IntensiveObservationPeriod on, with a planar geometry + // In this case IOP is running with a planar geometry auto dx = m_grid->get_geometry_data("dx_short").get_view()(); Kokkos::deep_copy(cell_length, dx*1000); // convert km -> m } else { diff --git a/components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp b/components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp index 61c61b878c2f..c27af18f85b2 100644 --- a/components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp +++ b/components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp @@ -64,11 +64,11 @@ void SPA::set_grids(const std::shared_ptr grids_manager) // where a single column of data corresponding to the closest lat/lon pair to // the IOP lat/lon parameters is read from file, and that column data is mapped // to all columns of the IdentityRemapper source fields. - EKAT_REQUIRE_MSG(spa_map_file == "" or spa_map_file == "None" or not m_iop, + EKAT_REQUIRE_MSG(spa_map_file == "" or spa_map_file == "None" or not m_iop_data_manager, "Error! Cannot define spa_remap_file for cases with an Intensive Observation Period defined. " "The IOP class defines it's own remap from file data -> model data.\n"); - SPAHorizInterp = SPAFunc::create_horiz_remapper (m_grid,spa_data_file,spa_map_file, m_iop!=nullptr); + SPAHorizInterp = SPAFunc::create_horiz_remapper (m_grid,spa_data_file,spa_map_file, m_iop_data_manager!=nullptr); // Grab a sw and lw field from the horiz interp, and check sw/lw dim against what we hardcoded in this class auto nswbands_data = SPAHorizInterp->get_src_field(4).get_header().get_identifier().get_layout().dim("swband"); @@ -128,8 +128,8 @@ void SPA::set_grids(const std::shared_ptr grids_manager) // AtmosphereInput object (for reading into standard // grids) or a SpaFunctions::IOPReader (for reading into // an IOP grid). - if (m_iop) { - SPAIOPDataReader = SPAFunc::create_spa_data_reader(m_iop,SPAHorizInterp,spa_data_file); + if (m_iop_data_manager) { + SPAIOPDataReader = SPAFunc::create_spa_data_reader(m_iop_data_manager,SPAHorizInterp,spa_data_file); } else { SPADataReader = SPAFunc::create_spa_data_reader(SPAHorizInterp,spa_data_file); } diff --git a/components/eamxx/src/physics/spa/spa_functions.hpp b/components/eamxx/src/physics/spa/spa_functions.hpp index 6bdba702e70f..444b2d9992e8 100644 --- a/components/eamxx/src/physics/spa/spa_functions.hpp +++ b/components/eamxx/src/physics/spa/spa_functions.hpp @@ -4,7 +4,7 @@ #include "share/grid/abstract_grid.hpp" #include "share/grid/remap/abstract_remapper.hpp" #include "share/io/scorpio_input.hpp" -#include "share/iop/intensive_observation_period.hpp" +#include "share/atm_process/IOPDataManager.hpp" #include "share/util/scream_time_stamp.hpp" #include "share/scream_types.hpp" @@ -30,7 +30,7 @@ struct SPAFunctions using gid_type = AbstractGrid::gid_type; - using iop_ptr_type = std::shared_ptr; + using iop_data_ptr_type = std::shared_ptr; template using view_1d = typename KT::template view_1d; @@ -128,11 +128,11 @@ struct SPAFunctions }; // SPAInput struct IOPReader { - IOPReader (iop_ptr_type& iop_, + IOPReader (iop_data_ptr_type& iop_, const std::string file_name_, const std::vector& io_fields_, const std::shared_ptr& io_grid_) - : iop(iop_), file_name(file_name_) + : iop_data_manager(iop_), file_name(file_name_) { field_mgr = std::make_shared(io_grid_); for (auto& f : io_fields_) { @@ -141,14 +141,14 @@ struct SPAFunctions } // Set IO info for this grid and file in IOP object - iop->setup_io_info(file_name, io_grid_); + iop_data_manager->setup_io_info(file_name, io_grid_); } void read_variables(const int time_index, const util::TimeStamp& ts) { - iop->read_fields_from_file_for_iop(file_name, field_names, ts, field_mgr, time_index); + iop_data_manager->read_fields_from_file_for_iop(file_name, field_names, ts, field_mgr, time_index); } - iop_ptr_type iop; + iop_data_ptr_type iop_data_manager; std::string file_name; std::vector field_names; std::shared_ptr field_mgr; @@ -175,7 +175,7 @@ struct SPAFunctions static std::shared_ptr create_spa_data_reader ( - iop_ptr_type& iop, + iop_data_ptr_type& iop_data_manager, const std::shared_ptr& horiz_remapper, const std::string& spa_data_file); diff --git a/components/eamxx/src/physics/spa/spa_functions_impl.hpp b/components/eamxx/src/physics/spa/spa_functions_impl.hpp index 287f69a98896..e566d0c985c7 100644 --- a/components/eamxx/src/physics/spa/spa_functions_impl.hpp +++ b/components/eamxx/src/physics/spa/spa_functions_impl.hpp @@ -162,7 +162,7 @@ template std::shared_ptr::IOPReader> SPAFunctions:: create_spa_data_reader ( - iop_ptr_type& iop, + iop_data_ptr_type& iop_data_manager, const std::shared_ptr& horiz_remapper, const std::string& spa_data_file) { @@ -171,7 +171,7 @@ create_spa_data_reader ( io_fields.push_back(horiz_remapper->get_src_field(i)); } const auto io_grid = horiz_remapper->get_src_grid(); - return std::make_shared(iop, spa_data_file, io_fields, io_grid); + return std::make_shared(iop_data_manager, spa_data_file, io_fields, io_grid); } /*-----------------------------------------------------------------*/ diff --git a/components/eamxx/src/share/CMakeLists.txt b/components/eamxx/src/share/CMakeLists.txt index 105b39a98086..dd290dfb364a 100644 --- a/components/eamxx/src/share/CMakeLists.txt +++ b/components/eamxx/src/share/CMakeLists.txt @@ -8,6 +8,7 @@ set(SHARE_SRC atm_process/atmosphere_process_group.cpp atm_process/atmosphere_process_dag.cpp atm_process/atmosphere_diagnostic.cpp + atm_process/IOPDataManager.cpp field/field_alloc_prop.cpp field/field_identifier.cpp field/field_header.cpp @@ -27,7 +28,6 @@ set(SHARE_SRC grid/remap/horiz_interp_remapper_data.cpp grid/remap/refining_remapper_p2p.cpp grid/remap/vertical_remapper.cpp - iop/intensive_observation_period.cpp property_checks/property_check.cpp property_checks/field_nan_check.cpp property_checks/field_within_interval_check.cpp diff --git a/components/eamxx/src/share/iop/intensive_observation_period.cpp b/components/eamxx/src/share/atm_process/IOPDataManager.cpp similarity index 98% rename from components/eamxx/src/share/iop/intensive_observation_period.cpp rename to components/eamxx/src/share/atm_process/IOPDataManager.cpp index 1d5f7c7ba336..2e2c3552d8e4 100644 --- a/components/eamxx/src/share/iop/intensive_observation_period.cpp +++ b/components/eamxx/src/share/atm_process/IOPDataManager.cpp @@ -1,7 +1,7 @@ #include "share/grid/point_grid.hpp" #include "share/io/scorpio_input.hpp" #include "share/io/scream_scorpio_interface.hpp" -#include "share/iop/intensive_observation_period.hpp" +#include "share/atm_process/IOPDataManager.hpp" #include "ekat/ekat_assert.hpp" #include "ekat/util/ekat_lin_interp.hpp" @@ -27,13 +27,13 @@ namespace ekat { namespace scream { namespace control { -IntensiveObservationPeriod:: -IntensiveObservationPeriod(const ekat::Comm& comm, - const ekat::ParameterList& params, - const util::TimeStamp& run_t0, - const int model_nlevs, - const Field& hyam, - const Field& hybm) +IOPDataManager:: +IOPDataManager(const ekat::Comm& comm, + const ekat::ParameterList& params, + const util::TimeStamp& run_t0, + const int model_nlevs, + const Field& hyam, + const Field& hybm) { m_comm = comm; m_params = params; @@ -72,14 +72,14 @@ IntensiveObservationPeriod(const ekat::Comm& comm, initialize_iop_file(run_t0, model_nlevs); } -IntensiveObservationPeriod:: -~IntensiveObservationPeriod () +IOPDataManager:: +~IOPDataManager () { const auto iop_file = m_params.get("iop_file"); scorpio::release_file(iop_file); } -void IntensiveObservationPeriod:: +void IOPDataManager:: initialize_iop_file(const util::TimeStamp& run_t0, int model_nlevs) { @@ -310,7 +310,7 @@ initialize_iop_file(const util::TimeStamp& run_t0, m_helper_fields.insert({"model_pressure", model_pressure}); } -void IntensiveObservationPeriod:: +void IOPDataManager:: setup_io_info(const std::string& file_name, const grid_ptr& grid) { @@ -397,7 +397,7 @@ setup_io_info(const std::string& file_name, } } -void IntensiveObservationPeriod:: +void IOPDataManager:: read_fields_from_file_for_iop (const std::string& file_name, const vos& field_names_nc, const vos& field_names_eamxx, @@ -501,7 +501,7 @@ read_fields_from_file_for_iop (const std::string& file_name, } } -void IntensiveObservationPeriod:: +void IOPDataManager:: read_iop_file_data (const util::TimeStamp& current_ts) { // Query to see if we need to load data from IOP file. @@ -749,7 +749,7 @@ read_iop_file_data (const util::TimeStamp& current_ts) m_time_info.time_idx_of_current_data = iop_file_time_idx; } -void IntensiveObservationPeriod:: +void IOPDataManager:: set_fields_from_iop_data(const field_mgr_ptr field_mgr) { if (m_params.get("zero_non_iop_tracers") && field_mgr->has_group("tracers")) { @@ -858,7 +858,7 @@ set_fields_from_iop_data(const field_mgr_ptr field_mgr) }); } -void IntensiveObservationPeriod:: +void IOPDataManager:: correct_temperature_and_water_vapor(const field_mgr_ptr field_mgr) { // Find the first valid level index for t_iop, i.e., first non-zero entry diff --git a/components/eamxx/src/share/iop/intensive_observation_period.hpp b/components/eamxx/src/share/atm_process/IOPDataManager.hpp similarity index 92% rename from components/eamxx/src/share/iop/intensive_observation_period.hpp rename to components/eamxx/src/share/atm_process/IOPDataManager.hpp index 860eb082c427..0ca5e2ef6d10 100644 --- a/components/eamxx/src/share/iop/intensive_observation_period.hpp +++ b/components/eamxx/src/share/atm_process/IOPDataManager.hpp @@ -14,11 +14,9 @@ namespace scream { namespace control { /* - * Class which provides functionality for running EAMxx with an intensive - * observation period (IOP). Currently the only use case is the doubly - * periodic model (DP-SCREAM). + * Class which data for an intensive observation period (IOP). */ -class IntensiveObservationPeriod +class IOPDataManager { using vos = std::vector; using field_mgr_ptr = std::shared_ptr; @@ -47,15 +45,15 @@ class IntensiveObservationPeriod // - run_t0: Initial timestamp for the simulation // - model_nlevs: Number of vertical levels in the simulation. Needed since // the iop file contains a (potentially) different number of levels - IntensiveObservationPeriod(const ekat::Comm& comm, - const ekat::ParameterList& params, - const util::TimeStamp& run_t0, - const int model_nlevs, - const Field& hyam, - const Field& hybm); + IOPDataManager(const ekat::Comm& comm, + const ekat::ParameterList& params, + const util::TimeStamp& run_t0, + const int model_nlevs, + const Field& hyam, + const Field& hybm); - // Default destructor - ~IntensiveObservationPeriod(); + // Destructor + ~IOPDataManager(); // Read data from IOP file and store internally. void read_iop_file_data(const util::TimeStamp& current_ts); @@ -197,7 +195,7 @@ class IntensiveObservationPeriod std::map m_iop_file_varnames; std::map m_iop_field_surface_varnames; std::map m_iop_field_type; -}; // class IntensiveObservationPeriod +}; // class IOPDataManager } // namespace control } // namespace scream diff --git a/components/eamxx/src/share/atm_process/atmosphere_process.hpp b/components/eamxx/src/share/atm_process/atmosphere_process.hpp index 43fc35a33185..016321ea506e 100644 --- a/components/eamxx/src/share/atm_process/atmosphere_process.hpp +++ b/components/eamxx/src/share/atm_process/atmosphere_process.hpp @@ -1,10 +1,10 @@ #ifndef SCREAM_ATMOSPHERE_PROCESS_HPP #define SCREAM_ATMOSPHERE_PROCESS_HPP -#include "share/iop/intensive_observation_period.hpp" #include "share/atm_process/atmosphere_process_utils.hpp" #include "share/atm_process/ATMBufferManager.hpp" #include "share/atm_process/SCDataManager.hpp" +#include "share/atm_process/IOPDataManager.hpp" #include "share/field/field_identifier.hpp" #include "share/field/field_manager.hpp" #include "share/property_checks/property_check.hpp" @@ -83,7 +83,7 @@ class AtmosphereProcess : public ekat::enable_shared_from_this; - using iop_ptr = std::shared_ptr; + using iop_data_ptr = std::shared_ptr; // Base constructor to set MPI communicator and params AtmosphereProcess (const ekat::Comm& comm, const ekat::ParameterList& params); @@ -280,8 +280,8 @@ class AtmosphereProcess : public ekat::enable_shared_from_this get_logger () const { @@ -597,7 +597,7 @@ class AtmosphereProcess : public ekat::enable_shared_from_thisset_iop(iop); + atm_proc->set_iop_data_manager(iop_data_manager); } } diff --git a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/CMakeLists.txt b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/CMakeLists.txt index f56d8beefa60..db5c8585943b 100644 --- a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/CMakeLists.txt +++ b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/CMakeLists.txt @@ -11,7 +11,7 @@ CreateDynamicsLib("theta-l_kokkos" 4 72 10) set (TEST_LABELS "dynamics;driver;tms;shoc;cld;spa;p3;rrtmgp;physics;dp") CreateUnitTest(homme_shoc_cld_spa_p3_rrtmgp_pg2_dp "homme_shoc_cld_spa_p3_rrtmgp_pg2_dp.cpp" LABELS ${TEST_LABELS} - LIBS cld_fraction tms shoc spa p3 scream_rrtmgp ${dynLibName} scream_control diagnostics + LIBS cld_fraction tms shoc spa iop_forcing p3 scream_rrtmgp ${dynLibName} scream_control diagnostics MPI_RANKS ${TEST_RANK_START} ${TEST_RANK_END} FIXTURES_SETUP_INDIVIDUAL ${FIXTURES_BASE_NAME} ) diff --git a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/input.yaml b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/input.yaml index 0950f6bfdbc2..47d0ed110869 100644 --- a/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/input.yaml +++ b/components/eamxx/tests/multi-process/dynamics_physics/homme_shoc_cld_spa_p3_rrtmgp_pg2_dp/input.yaml @@ -40,7 +40,7 @@ atmosphere_processes: homme: Moisture: moist physics: - atm_procs_list: [mac_aero_mic,rrtmgp] + atm_procs_list: [iop_forcing,mac_aero_mic,rrtmgp] schedule_type: Sequential Type: Group mac_aero_mic: From 10892147ab4810f1ad4216e2f7745ee9d56f83e6 Mon Sep 17 00:00:00 2001 From: tcclevenger Date: Thu, 5 Dec 2024 12:33:21 -0500 Subject: [PATCH 2/5] Remove unnecessary if-block if (iop_nudge_tq or iop_nudge_uv) { ... if (iop_nudge_tq or iop_nudge_uv) { ... } } --- .../physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp index 9aac3fdf03ac..ee662cd81b86 100644 --- a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp @@ -512,7 +512,6 @@ void IOPForcing::run_impl (const double dt) }); team.team_barrier(); - if (iop_nudge_tq or iop_nudge_uv) { Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) { if (iop_nudge_tq) { // Restrict nudging of T and qv to certain levels if requested by user @@ -535,7 +534,6 @@ void IOPForcing::run_impl (const double dt) v_i(k).update(v_mean(k) - v_iop(k), -dt/rtau, 1.0); } }); - } // Release WS views ws.release_many_contiguous<1>({&ref_p_mid}); From 9348923cfe54e21dadde1025a7fba49e26588257 Mon Sep 17 00:00:00 2001 From: tcclevenger Date: Thu, 5 Dec 2024 12:34:04 -0500 Subject: [PATCH 3/5] Use field util horiz_contraction to compute means --- .../eamxx_iop_forcing_process_interface.cpp | 172 ++++++++---------- .../eamxx_iop_forcing_process_interface.hpp | 12 +- 2 files changed, 82 insertions(+), 102 deletions(-) diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp index ee662cd81b86..1543f5de63e9 100644 --- a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp @@ -1,5 +1,6 @@ #include "physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp" +#include "share/field/field_utils.hpp" #include "share/property_checks/field_within_interval_check.hpp" namespace scream @@ -35,11 +36,21 @@ void IOPForcing::set_grids(const std::shared_ptr grids_manag "Error! IOPDataManager not setup by driver, but IOPForcing" "being used as an ATM process.\n"); - // Compute number of buffer views + // Create helper fields for finding horizontal means + auto level_only_scalar_layout = scalar3d_mid.clone().strip_dim(0); + auto level_only_vector_layout = vector3d_mid.clone().strip_dim(0); const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); - if (iop_nudge_tq) m_buffer.num_1d_scalar_nlev += 2; - if (iop_nudge_uv) m_buffer.num_1d_scalar_nlev += 2; + if (iop_nudge_tq or iop_nudge_uv) { + create_helper_field("horiz_mean_weights", scalar2d, grid_name); + } + if (iop_nudge_tq) { + create_helper_field("qv_mean", level_only_scalar_layout, grid_name); + create_helper_field("t_mean", level_only_scalar_layout, grid_name); + } + if (iop_nudge_uv) { + create_helper_field("horiz_winds_mean", level_only_vector_layout, grid_name); + } } // ========================================================================================= void IOPForcing:: @@ -61,17 +72,12 @@ set_computed_group_impl (const FieldGroup& group) // ========================================================================================= size_t IOPForcing::requested_buffer_size_in_bytes() const { - const int nlev_packs = ekat::npack(m_num_levs); - const int nlevi_packs = ekat::npack(m_num_levs+1); - - const size_t temp_view_bytes = - m_buffer.num_1d_scalar_nlev*nlev_packs*sizeof(Pack); - // Number of bytes needed by the WorkspaceManager passed to shoc_main + const int nlevi_packs = ekat::npack(m_num_levs+1); const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs); const size_t wsm_bytes = WorkspaceMgr::get_total_bytes_needed(nlevi_packs, 7+m_num_tracers, policy); - return temp_view_bytes + wsm_bytes; + return wsm_bytes; } // ========================================================================================= void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager) @@ -79,28 +85,9 @@ void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager) EKAT_REQUIRE_MSG(buffer_manager.allocated_bytes() >= requested_buffer_size_in_bytes(), "Error! Buffers size not sufficient.\n"); - const int nlev_packs = ekat::npack(m_num_levs); const int nlevi_packs = ekat::npack(m_num_levs+1); - Pack* mem = reinterpret_cast(buffer_manager.get_memory()); - // Temp view data - using mean_view_t = decltype(m_buffer.qv_mean); - const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); - const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); - if (iop_nudge_tq) { - m_buffer.qv_mean = mean_view_t(mem, nlev_packs); - mem += m_buffer.qv_mean.size(); - m_buffer.t_mean = mean_view_t(mem, nlev_packs); - mem += m_buffer.t_mean.size(); - } - if (iop_nudge_uv) { - m_buffer.u_mean = mean_view_t(mem, nlev_packs); - mem += m_buffer.u_mean.size(); - m_buffer.v_mean = mean_view_t(mem, nlev_packs); - mem += m_buffer.v_mean.size(); - } - // WSM data m_buffer.wsm_data = mem; @@ -112,6 +99,22 @@ void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager) EKAT_REQUIRE_MSG(used_mem==requested_buffer_size_in_bytes(), "Error! Used memory != requested memory for IOPForcing.\n"); } // ========================================================================================= +void IOPForcing::create_helper_field (const std::string& name, + const FieldLayout& layout, + const std::string& grid_name) +{ + using namespace ekat::units; + FieldIdentifier id(name,layout,Units::nondimensional(),grid_name); + + // Create the field. Init with NaN's, so we spot instances of uninited memory usage + Field f(id); + f.get_header().get_alloc_properties().request_allocation(); + f.allocate_view(); + f.deep_copy(ekat::ScalarTraits::invalid()); + + m_helper_fields[name] = f; +} +// ========================================================================================= void IOPForcing::initialize_impl (const RunType run_type) { // Set field property checks for the fields in this process @@ -126,6 +129,12 @@ void IOPForcing::initialize_impl (const RunType run_type) const auto nlevi_packs = ekat::npack(m_num_levs+1); const auto policy = ESU::get_default_team_policy(m_num_cols, nlevi_packs); m_workspace_mgr.setup(m_buffer.wsm_data, nlevi_packs, 7+m_num_tracers, policy); + + // Compute field for horizontal contraction weights (1/num_global_dofs) + const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); + const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); + const Real one_over_num_dofs = 1.0/m_grid->get_num_global_dofs(); + if (iop_nudge_tq or iop_nudge_uv) m_helper_fields.at("horiz_mean_weights").deep_copy(one_over_num_dofs); } // ========================================================================================= KOKKOS_FUNCTION @@ -437,55 +446,22 @@ void IOPForcing::run_impl (const double dt) // and observed quantities of T, Q, u, and v if (iop_nudge_tq or iop_nudge_uv) { // Compute domain mean of qv, T_mid, u, and v - const auto qv_mean = m_buffer.qv_mean; - const auto t_mean = m_buffer.t_mean; - const auto u_mean = m_buffer.u_mean; - const auto v_mean = m_buffer.v_mean; - - const auto qv_mean_h = Kokkos::create_mirror_view(qv_mean); - const auto t_mean_h = Kokkos::create_mirror_view(t_mean); - const auto u_mean_h = Kokkos::create_mirror_view(u_mean); - const auto v_mean_h = Kokkos::create_mirror_view(v_mean); - - const auto num_global_cols = m_grid->get_num_global_dofs(); - for (int k=0; k qv_mean, t_mean; + view_2d horiz_winds_mean; + if (iop_nudge_tq){ + horiz_contraction(m_helper_fields.at("qv_mean"), get_field_out("qv"), + m_helper_fields.at("horiz_mean_weights"), &m_comm); + qv_mean = m_helper_fields.at("qv_mean").get_view(); + + horiz_contraction(m_helper_fields.at("t_mean"), get_field_out("T_mid"), + m_helper_fields.at("horiz_mean_weights"), &m_comm); + t_mean = m_helper_fields.at("t_mean").get_view(); + } + if (iop_nudge_uv){ + horiz_contraction(m_helper_fields.at("horiz_winds_mean"), get_field_out("horiz_winds"), + m_helper_fields.at("horiz_mean_weights"), &m_comm); + horiz_winds_mean = m_helper_fields.at("horiz_winds_mean").get_view(); } - Kokkos::deep_copy(qv_mean, qv_mean_h); - Kokkos::deep_copy(t_mean, t_mean_h); - Kokkos::deep_copy(u_mean, u_mean_h); - Kokkos::deep_copy(v_mean, v_mean_h); // Apply relaxation const auto rtau = std::max(dt, iop_nudge_tscale); @@ -512,28 +488,28 @@ void IOPForcing::run_impl (const double dt) }); team.team_barrier(); - Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) { - if (iop_nudge_tq) { - // Restrict nudging of T and qv to certain levels if requested by user - // IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high - // is in units of [hPa], thus convert iop_nudge_tq_low/high - Mask nudge_level(false); - int max_size = hyam.size(); - for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) { - const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop; - nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100 - and - pressure_from_iop >= iop_nudge_tq_high*100); - } - - qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0); - T_mid_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0); - } - if (iop_nudge_uv) { - u_i(k).update(u_mean(k) - u_iop(k), -dt/rtau, 1.0); - v_i(k).update(v_mean(k) - v_iop(k), -dt/rtau, 1.0); + Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&](const int& k) { + if (iop_nudge_tq) { + // Restrict nudging of T and qv to certain levels if requested by user + // IOP pressure variable is in unitis of [Pa], while iop_nudge_tq_low/high + // is in units of [hPa], thus convert iop_nudge_tq_low/high + Mask nudge_level(false); + int max_size = hyam.size(); + for (int lev=k*Pack::n, p = 0; p < Pack::n && lev < max_size; ++lev, ++p) { + const auto pressure_from_iop = hyam(lev)*ps0 + hybm(lev)*ps_iop; + nudge_level.set(p, pressure_from_iop <= iop_nudge_tq_low*100 + and + pressure_from_iop >= iop_nudge_tq_high*100); } - }); + + qv_i(k).update(nudge_level, qv_mean(k) - qv_iop(k), -dt/rtau, 1.0); + T_mid_i(k).update(nudge_level, t_mean(k) - t_iop(k), -dt/rtau, 1.0); + } + if (iop_nudge_uv) { + u_i(k).update(horiz_winds_mean(0, k) - u_iop(k), -dt/rtau, 1.0); + v_i(k).update(horiz_winds_mean(1, k) - v_iop(k), -dt/rtau, 1.0); + } + }); // Release WS views ws.release_many_contiguous<1>({&ref_p_mid}); diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp index 535043f401e5..0af010ec8ec9 100644 --- a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp @@ -116,6 +116,11 @@ class IOPForcing : public scream::AtmosphereProcess void finalize_impl () {} + // Creates an helper field, not to be shared with the AD's FieldManager + void create_helper_field (const std::string& name, + const FieldLayout& layout, + const std::string& grid_name); + void set_computed_group_impl (const FieldGroup& group); // Computes total number of bytes needed for local variables @@ -132,13 +137,12 @@ class IOPForcing : public scream::AtmosphereProcess Int m_num_tracers; struct Buffer { - int num_1d_scalar_nlev = 0; - - uview_1d qv_mean, t_mean, u_mean, v_mean; - Pack* wsm_data; }; + // Some helper fields. + std::map m_helper_fields; + // Struct which contains local variables Buffer m_buffer; From 8b76d0f1d469b9fb97f740c837258eccce2798aa Mon Sep 17 00:00:00 2001 From: tcclevenger Date: Mon, 9 Dec 2024 08:42:16 -0500 Subject: [PATCH 4/5] EAMxx: Unrelated: remove left over debug print from scorpio_output.cpp --- components/eamxx/src/share/io/scorpio_output.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/components/eamxx/src/share/io/scorpio_output.cpp b/components/eamxx/src/share/io/scorpio_output.cpp index 4e9d5add8004..d51df07fde62 100644 --- a/components/eamxx/src/share/io/scorpio_output.cpp +++ b/components/eamxx/src/share/io/scorpio_output.cpp @@ -1367,7 +1367,6 @@ AtmosphereOutput::create_diagnostic (const std::string& diag_field_name) for (const auto& freq : diag->get_required_field_requests()) { const auto& fname = freq.fid.name(); if (!sim_field_mgr->has_field(fname)) { - std::cout << diag_field_name << " depends on the diag " << fname << "\n"; // This diag depends on another diag. Create and init the dependency if (m_diagnostics.count(fname)==0) { m_diagnostics[fname] = create_diagnostic(fname); From 88153c6f7004106a76ae4e24e387e7c5777b8918 Mon Sep 17 00:00:00 2001 From: tcclevenger Date: Tue, 10 Dec 2024 20:28:24 -0500 Subject: [PATCH 5/5] Use packs for *_mean views Fixes a bug where real views were being used as if they were packed. --- .../eamxx_iop_forcing_process_interface.cpp | 23 ++++++++++--------- .../eamxx_iop_forcing_process_interface.hpp | 3 ++- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp index 1543f5de63e9..930dd9636368 100644 --- a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp @@ -42,14 +42,14 @@ void IOPForcing::set_grids(const std::shared_ptr grids_manag const auto iop_nudge_tq = m_iop_data_manager->get_params().get("iop_nudge_tq"); const auto iop_nudge_uv = m_iop_data_manager->get_params().get("iop_nudge_uv"); if (iop_nudge_tq or iop_nudge_uv) { - create_helper_field("horiz_mean_weights", scalar2d, grid_name); + create_helper_field("horiz_mean_weights", scalar2d, grid_name, pack_size); } if (iop_nudge_tq) { - create_helper_field("qv_mean", level_only_scalar_layout, grid_name); - create_helper_field("t_mean", level_only_scalar_layout, grid_name); + create_helper_field("qv_mean", level_only_scalar_layout, grid_name, pack_size); + create_helper_field("t_mean", level_only_scalar_layout, grid_name, pack_size); } if (iop_nudge_uv) { - create_helper_field("horiz_winds_mean", level_only_vector_layout, grid_name); + create_helper_field("horiz_winds_mean", level_only_vector_layout, grid_name, pack_size); } } // ========================================================================================= @@ -101,14 +101,15 @@ void IOPForcing::init_buffers(const ATMBufferManager &buffer_manager) // ========================================================================================= void IOPForcing::create_helper_field (const std::string& name, const FieldLayout& layout, - const std::string& grid_name) + const std::string& grid_name, + const int ps) { using namespace ekat::units; FieldIdentifier id(name,layout,Units::nondimensional(),grid_name); // Create the field. Init with NaN's, so we spot instances of uninited memory usage Field f(id); - f.get_header().get_alloc_properties().request_allocation(); + f.get_header().get_alloc_properties().request_allocation(ps); f.allocate_view(); f.deep_copy(ekat::ScalarTraits::invalid()); @@ -446,21 +447,21 @@ void IOPForcing::run_impl (const double dt) // and observed quantities of T, Q, u, and v if (iop_nudge_tq or iop_nudge_uv) { // Compute domain mean of qv, T_mid, u, and v - view_1d qv_mean, t_mean; - view_2d horiz_winds_mean; + view_1d qv_mean, t_mean; + view_2d horiz_winds_mean; if (iop_nudge_tq){ horiz_contraction(m_helper_fields.at("qv_mean"), get_field_out("qv"), m_helper_fields.at("horiz_mean_weights"), &m_comm); - qv_mean = m_helper_fields.at("qv_mean").get_view(); + qv_mean = m_helper_fields.at("qv_mean").get_view(); horiz_contraction(m_helper_fields.at("t_mean"), get_field_out("T_mid"), m_helper_fields.at("horiz_mean_weights"), &m_comm); - t_mean = m_helper_fields.at("t_mean").get_view(); + t_mean = m_helper_fields.at("t_mean").get_view(); } if (iop_nudge_uv){ horiz_contraction(m_helper_fields.at("horiz_winds_mean"), get_field_out("horiz_winds"), m_helper_fields.at("horiz_mean_weights"), &m_comm); - horiz_winds_mean = m_helper_fields.at("horiz_winds_mean").get_view(); + horiz_winds_mean = m_helper_fields.at("horiz_winds_mean").get_view(); } // Apply relaxation diff --git a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp index 0af010ec8ec9..7cec311a231c 100644 --- a/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp +++ b/components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.hpp @@ -119,7 +119,8 @@ class IOPForcing : public scream::AtmosphereProcess // Creates an helper field, not to be shared with the AD's FieldManager void create_helper_field (const std::string& name, const FieldLayout& layout, - const std::string& grid_name); + const std::string& grid_name, + const int ps = 1); void set_computed_group_impl (const FieldGroup& group);