Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

EAMxx: add IOPRemapper and use it in SPA and IOPDataManager #6914

Open
wants to merge 10 commits into
base: master
Choose a base branch
from
Open

EAMxx: add IOPRemapper and use it in SPA and IOPDataManager #6914

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
841d35e
EAMxx: added a few utility methods to Field/FieldLayout
bartgol Jan 17, 2025
f5c3f80
EAMxx: add an IOP (horiz) remapper
bartgol Jan 17, 2025
e5ff284
EAMxx: refactor SPA to use IOPRemapper and prune IOPReader
bartgol Jan 17, 2025
da7eb3a
EAMxx: fix setup of iop remapper test
bartgol Jan 20, 2025
d055cb5
EAMxx: fix issue with catch rng seed in unit tests
bartgol Jan 20, 2025
661dcce
EAMxx: fix compilation error on CUDA
bartgol Jan 20, 2025
7f0287e
EAMxx: use IOPRemapper in IOPDataManager
bartgol Jan 20, 2025
5ac5995
EAMxx: fix linking error in SCMInput sources
bartgol Jan 21, 2025
4570448
EAMxx: fix warning in VerticalRemapper
bartgol Jan 21, 2025
5ac2a9f
EAMxx: fix bad usage of macro for IOP remapper
bartgol Jan 21, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
124 changes: 34 additions & 90 deletions components/eamxx/src/control/atmosphere_driver.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1001,11 +1001,14 @@ void AtmosphereDriver::restart_model ()
continue;
}
const auto& restart_group = it.second->get_groups_info().at("RESTART");
std::vector<std::string> fnames;
std::vector<Field> fields;
for (const auto& fn : restart_group->m_fields_names) {
fnames.push_back(fn);
fields.push_back(it.second->get_field(fn));
}
read_fields_from_file (fields,it.second->get_grid(),filename);
for (auto& f : fields) {
f.get_header().get_tracking().update_time_stamp(m_current_ts);
}
read_fields_from_file (fnames,it.second->get_grid(),filename,m_current_ts);
}

// Restart the num steps counter in the atm time stamp
Expand Down Expand Up @@ -1267,17 +1270,20 @@ void AtmosphereDriver::set_initial_conditions ()
// Now loop over all grids, and load from file the needed fields on each grid (if any).
const auto& file_name = ic_pl.get<std::string>("Filename");
m_atm_logger->info(" [EAMxx] IC filename: " + file_name);
for (const auto& it : m_field_mgrs) {
const auto& grid_name = it.first;
for (const auto& [grid_name,fm] : m_field_mgrs) {
std::vector<Field> ic_fields;
for (const auto& fn : ic_fields_names[grid_name]) {
ic_fields.push_back(fm->get_field(fn));
}
if (not m_iop_data_manager) {
read_fields_from_file (ic_fields_names[grid_name],it.second->get_grid(),file_name,m_current_ts);
read_fields_from_file (ic_fields,fm->get_grid(),file_name);
} else {
// For IOP enabled, we load from file and copy data from the closest
// lat/lon column to every other column
m_iop_data_manager->read_fields_from_file_for_iop(file_name,
ic_fields_names[grid_name],
m_current_ts,
it.second);
m_iop_data_manager->read_fields_from_file_for_iop(file_name,ic_fields,fm->get_grid());
}
for (auto& f : ic_fields) {
f.get_header().get_tracking().update_time_stamp(m_current_ts);
}
}
}
Expand Down Expand Up @@ -1342,29 +1348,33 @@ void AtmosphereDriver::set_initial_conditions ()
m_atm_logger->info(" [EAMxx] Reading topography from file ...");
const auto& file_name = ic_pl.get<std::string>("topography_filename");
m_atm_logger->info(" filename: " + file_name);
for (const auto& it : m_field_mgrs) {
const auto& grid_name = it.first;
for (const auto& [grid_name,fm] : m_field_mgrs) {
std::vector<Field> topo_fields;
int nfields = topography_eamxx_fields_names[grid_name].size();
for (int i=0; i<nfields; ++i) {
auto eamxx_fname = topography_eamxx_fields_names[grid_name][i];
auto file_fname = topography_file_fields_names[grid_name][i];
topo_fields.push_back(fm->get_field(eamxx_fname).alias(file_fname));
}

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();
auto io_grid = fm->get_grid();
if (grid_name=="Physics GLL") {
using namespace ShortFieldTagsNames;
auto grid = io_grid->clone(io_grid->name(),true);
grid->reset_field_tag_name(COL,"ncol_d");
io_grid = grid;
}
read_fields_from_file (topography_file_fields_names[grid_name],
topography_eamxx_fields_names[grid_name],
io_grid,file_name,m_current_ts);
read_fields_from_file (topo_fields,io_grid,file_name);
} else {
// For IOP enabled, we load from file and copy data from the closest
// lat/lon column to every other column
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,
it.second);
m_iop_data_manager->read_fields_from_file_for_iop(file_name,topo_fields,fm->get_grid());
}
for (auto& f : topo_fields) {
f.get_header().get_tracking().update_time_stamp(m_current_ts);
}
}
// Store in provenance list, for later usage in output file metadata
Expand Down Expand Up @@ -1469,83 +1479,17 @@ void AtmosphereDriver::set_initial_conditions ()
}

void AtmosphereDriver::
read_fields_from_file (const std::vector<std::string>& field_names_nc,
const std::vector<std::string>& field_names_eamxx,
const std::shared_ptr<const AbstractGrid>& grid,
const std::string& file_name,
const util::TimeStamp& t0)
{
EKAT_REQUIRE_MSG(field_names_nc.size()==field_names_eamxx.size(),
"Error! Field name arrays must have same size.\n");

if (field_names_nc.size()==0) {
return;
}

// NOTE: we cannot pass the field_mgr and m_grids_mgr, since the input
// grid may not be in the grids_manager and may not be the grid
// of the field mgr. This sounds weird, but there is a precise
// use case: when grid is a shallow clone of the fm grid, where
// we changed the name of some field tags (e.g., we set the name
// of COL to ncol_d). This is used when reading the topography,
// since the topo file *always* uses ncol_d for GLL points data,
// while a non-PG2 run would have the tag name be "ncol".
const auto& field_mgr = m_field_mgrs.at(grid->name());
std::vector<Field> fields;
for (size_t i=0; i<field_names_nc.size(); ++i) {
const auto& eamxx_name = field_names_eamxx[i];
const auto& nc_name = field_names_nc[i];
fields.push_back(field_mgr->get_field(eamxx_name).alias(nc_name));
}

AtmosphereInput ic_reader(file_name,grid,fields);
ic_reader.set_logger(m_atm_logger);
ic_reader.read_variables();
ic_reader.finalize();

for (auto& f : fields) {
// Set the initial time stamp
// NOTE: f is an alias of the field from field_mgr, so it shares all
// pointers to the metadata (except for the FieldIdentifier),
// so changing its timestamp will also change the timestamp
// of the field in field_mgr
f.get_header().get_tracking().update_time_stamp(t0);
}
}

void AtmosphereDriver::
read_fields_from_file (const std::vector<std::string>& field_names,
read_fields_from_file (const std::vector<Field>& fields,
const std::shared_ptr<const AbstractGrid>& grid,
const std::string& file_name,
const util::TimeStamp& t0)
const std::string& file_name)
{
if (field_names.size()==0) {
if (fields.size()==0) {
return;
}

// NOTE: we cannot pass the field_mgr and m_grids_mgr, since the input
// grid may not be in the grids_manager and may not be the grid
// of the field mgr. This sounds weird, but there is a precise
// use case: when grid is a shallow clone of the fm grid, where
// we changed the name of some field tags (e.g., we set the name
// of COL to ncol_d). This is used when reading the topography,
// since the topo file *always* uses ncol_d for GLL points data,
// while a non-PG2 run would have the tag name be "ncol".
const auto& field_mgr = m_field_mgrs.at(grid->name());
std::vector<Field> fields;
for (const auto& fn : field_names) {
fields.push_back(field_mgr->get_field(fn));
}

AtmosphereInput ic_reader(file_name,grid,fields);
ic_reader.set_logger(m_atm_logger);
ic_reader.read_variables();
ic_reader.finalize();

for (auto& f : fields) {
// Set the initial time stamp
f.get_header().get_tracking().update_time_stamp(t0);
}
}

void AtmosphereDriver::
Expand Down
17 changes: 3 additions & 14 deletions components/eamxx/src/control/atmosphere_driver.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,21 +185,10 @@ class AtmosphereDriver
void set_initial_conditions ();
void restart_model ();

// Read fields from a file when the names of the fields in
// EAMxx do not match exactly with the .nc file. Example is
// for topography data files, where GLL and PG2 grid have
// different naming conventions for phis.
void read_fields_from_file (const std::vector<std::string>& field_names_nc,
const std::vector<std::string>& field_names_eamxx,
// Read fields from a file
void read_fields_from_file (const std::vector<Field>& fields,
const std::shared_ptr<const AbstractGrid>& grid,
const std::string& file_name,
const util::TimeStamp& t0);
// Read fields from a file when the names of the fields in
// EAMxx match with the .nc file.
void read_fields_from_file (const std::vector<std::string>& field_names,
const std::shared_ptr<const AbstractGrid>& grid,
const std::string& file_name,
const util::TimeStamp& t0);
const std::string& file_name);
void register_groups ();

std::map<std::string,field_mgr_ptr> m_field_mgrs;
Expand Down
22 changes: 6 additions & 16 deletions components/eamxx/src/physics/spa/eamxx_spa_process_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,15 +60,12 @@ void SPA::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
auto spa_data_file = m_params.get<std::string>("spa_data_file");
auto spa_map_file = m_params.get<std::string>("spa_remap_file","");

// IOP cases cannot have a remap file. IOP file reader is itself a remaper,
// 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.
// IOP cases cannot have a remap file. We will create a IOPRemapper as the horiz remapper
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_data_manager!=nullptr);
SPAHorizInterp = SPAFunc::create_horiz_remapper (m_grid,spa_data_file,spa_map_file, m_iop_data_manager);

// 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");
Expand Down Expand Up @@ -124,15 +121,8 @@ void SPA::set_grids(const std::shared_ptr<const GridsManager> grids_manager)
Kokkos::deep_copy(spa_hybm,spa_hybm_h);
}

// 4. Create reader for spa data. The reader is either an
// AtmosphereInput object (for reading into standard
// grids) or a SpaFunctions::IOPReader (for reading into
// an IOP grid).
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);
}
// 4. Create reader for spa data.
SPADataReader = SPAFunc::create_spa_data_reader(SPAHorizInterp,spa_data_file);
}

// =========================================================================================
Expand Down Expand Up @@ -224,7 +214,7 @@ void SPA::initialize_impl (const RunType /* run_type */)
// Note: At the first time step, the data will be moved into spa_beg,
// and spa_end will be reloaded from file with the new month.
const int curr_month = timestamp().get_month()-1; // 0-based
SPAFunc::update_spa_data_from_file(SPADataReader,SPAIOPDataReader,timestamp(),curr_month,*SPAHorizInterp,SPAData_end);
SPAFunc::update_spa_data_from_file(*SPADataReader,curr_month,*SPAHorizInterp,SPAData_end);

// 6. Set property checks for fields in this process
using Interval = FieldWithinIntervalCheck;
Expand All @@ -246,7 +236,7 @@ void SPA::run_impl (const double dt)
/* Update the SPATimeState to reflect the current time, note the addition of dt */
SPATimeState.t_now = ts.frac_of_year_in_days();
/* Update time state and if the month has changed, update the data.*/
SPAFunc::update_spa_timestate(SPADataReader,SPAIOPDataReader,ts,*SPAHorizInterp,SPATimeState,SPAData_start,SPAData_end);
SPAFunc::update_spa_timestate(*SPADataReader,ts,*SPAHorizInterp,SPATimeState,SPAData_start,SPAData_end);

// Call the main SPA routine to get interpolated aerosol forcings.
const auto& pmid_tgt = get_field_in("p_mid").get_view<const Spack**>();
Expand Down
2 changes: 0 additions & 2 deletions components/eamxx/src/physics/spa/eamxx_spa_process_interface.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,8 +79,6 @@ class SPA : public AtmosphereProcess

// IO structure to read in data for standard grids (keep it around to avoid re-creating PIO decomps)
std::shared_ptr<AtmosphereInput> SPADataReader;
// Similar to above, but stores info to read data for IOP grid
std::shared_ptr<SPAFunc::IOPReader> SPAIOPDataReader;

// Structures to store the data used for interpolation
std::shared_ptr<AbstractRemapper> SPAHorizInterp;
Expand Down
60 changes: 11 additions & 49 deletions components/eamxx/src/physics/spa/spa_functions.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,6 @@ struct SPAFunctions

using gid_type = AbstractGrid::gid_type;

using iop_data_ptr_type = std::shared_ptr<control::IOPDataManager>;

template <typename S>
using view_1d = typename KT::template view_1d<S>;
template <typename S>
Expand Down Expand Up @@ -127,33 +125,6 @@ struct SPAFunctions
SPAData data; // All spa fields
}; // SPAInput

struct IOPReader {
IOPReader (iop_data_ptr_type& iop_,
const std::string file_name_,
const std::vector<Field>& io_fields_,
const std::shared_ptr<const AbstractGrid>& io_grid_)
: iop_data_manager(iop_), file_name(file_name_)
{
field_mgr = std::make_shared<FieldManager>(io_grid_);
for (auto& f : io_fields_) {
field_mgr->add_field(f);
field_names.push_back(f.name());
}

// Set IO info for this grid and file in IOP object
iop_data_manager->setup_io_info(file_name, io_grid_);
}

void read_variables(const int time_index, const util::TimeStamp& ts) {
iop_data_manager->read_fields_from_file_for_iop(file_name, field_names, ts, field_mgr, time_index);
}

iop_data_ptr_type iop_data_manager;
std::string file_name;
std::vector<std::string> field_names;
std::shared_ptr<FieldManager> field_mgr;
};

// The output is really just SPAData, but for clarity it might
// help to see a SPAOutput along a SPAInput in functions signatures
using SPAOutput = SPAData;
Expand All @@ -166,19 +137,13 @@ struct SPAFunctions
const std::shared_ptr<const AbstractGrid>& model_grid,
const std::string& spa_data_file,
const std::string& map_file,
const bool use_iop = false);
const std::shared_ptr<control::IOPDataManager>& iop_data_mgr = nullptr);

static std::shared_ptr<AtmosphereInput>
create_spa_data_reader (
const std::shared_ptr<AbstractRemapper>& horiz_remapper,
const std::string& spa_data_file);

static std::shared_ptr<IOPReader>
create_spa_data_reader (
iop_data_ptr_type& iop_data_manager,
const std::shared_ptr<AbstractRemapper>& horiz_remapper,
const std::string& spa_data_file);

static void spa_main(
const SPATimeState& time_state,
const view_2d<const Spack>& p_tgt,
Expand All @@ -189,21 +154,18 @@ struct SPAFunctions
const SPAOutput& data_out);

static void update_spa_data_from_file(
std::shared_ptr<AtmosphereInput>& scorpio_reader,
std::shared_ptr<IOPReader>& iop_reader,
const util::TimeStamp& ts,
const int time_index, // zero-based
AbstractRemapper& spa_horiz_interp,
SPAInput& spa_input);
AtmosphereInput& scorpio_reader,
const int time_index, // zero-based
AbstractRemapper& spa_horiz_interp,
SPAInput& spa_input);

static void update_spa_timestate(
std::shared_ptr<AtmosphereInput>& scorpio_reader,
std::shared_ptr<IOPReader>& iop_reader,
const util::TimeStamp& ts,
AbstractRemapper& spa_horiz_interp,
SPATimeState& time_state,
SPAInput& spa_beg,
SPAInput& spa_end);
AtmosphereInput& scorpio_reader,
const util::TimeStamp& ts,
AbstractRemapper& spa_horiz_interp,
SPATimeState& time_state,
SPAInput& spa_beg,
SPAInput& spa_end);

// The following three are called during spa_main
static void perform_time_interpolation (
Expand Down
Loading
Loading