Skip to content

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

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

Merged
merged 11 commits into from
Jan 29, 2025
Merged

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
11 commits
Select commit Hold shift + click to select a range
0329560
EAMxx: added a few utility methods to Field/FieldLayout
bartgol Jan 17, 2025
bb3303a
EAMxx: add an IOP (horiz) remapper
bartgol Jan 17, 2025
e07fe64
EAMxx: refactor SPA to use IOPRemapper and prune IOPReader
bartgol Jan 17, 2025
6b2d85d
EAMxx: fix setup of iop remapper test
bartgol Jan 20, 2025
148fdd2
EAMxx: fix issue with catch rng seed in unit tests
bartgol Jan 20, 2025
82e35dc
EAMxx: fix compilation error on CUDA
bartgol Jan 20, 2025
7fc987c
EAMxx: use IOPRemapper in IOPDataManager
bartgol Jan 20, 2025
3595856
EAMxx: fix linking error in SCMInput sources
bartgol Jan 21, 2025
dace0f3
EAMxx: fix warning in VerticalRemapper
bartgol Jan 21, 2025
a8ffa9c
EAMxx: fix bad usage of macro for IOP remapper
bartgol Jan 21, 2025
a3925a8
EAMxx: fix bug in SCMInput and cleanup scm reader test
bartgol Jan 27, 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