diff --git a/src/monio/Constants.h b/src/monio/Constants.h index ab08b0f..f4adcdb 100644 --- a/src/monio/Constants.h +++ b/src/monio/Constants.h @@ -95,6 +95,11 @@ enum eDataTypes { eNumberOfDataTypes }; +enum eLfricAtlasMapMethods { + eTrivial, + eKDtree +}; + /// String/Views /////////////////////////////////////////////////////////////////////////////////// const std::string_view kTimeDimName = "time_counter"; diff --git a/src/monio/Monio.cc b/src/monio/Monio.cc index 8d3729a..b56fdda 100644 --- a/src/monio/Monio.cc +++ b/src/monio/Monio.cc @@ -44,9 +44,10 @@ monio::Monio::~Monio() { } void monio::Monio::readState(atlas::FieldSet& localFieldSet, - const std::vector& fieldMetadataVec, - const std::string& filePath, - const util::DateTime& dateTime) { + const std::vector& fieldMetadataVec, + const std::string& filePath, + const util::DateTime& dateTime, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod) { oops::Log::trace() << "Monio::readState()" << std::endl; if (localFieldSet.size() == 0) { Monio::get().closeFiles(); @@ -62,7 +63,8 @@ void monio::Monio::readState(atlas::FieldSet& localFieldSet, auto& functionSpace = globalField.functionspace(); auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace); // Initialise file - int variableConvention = initialiseFile(grid, filePath, true); + const int variableConvention = initialiseFile(grid, filePath, true, + lfricAtlasMapMethod); // getFileData returns a copy of FileData (with required LFRic mesh data), so read data // is discarded when FileData goes out-of-scope for reading subsequent fields. FileData fileData = getFileData(grid.name()); @@ -106,7 +108,8 @@ void monio::Monio::readState(atlas::FieldSet& localFieldSet, void monio::Monio::readIncrements(atlas::FieldSet& localFieldSet, const std::vector& fieldMetadataVec, - const std::string& filePath) { + const std::string& filePath, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod) { oops::Log::trace() << "Monio::readIncrements()" << std::endl; if (localFieldSet.size() == 0) { Monio::get().closeFiles(); @@ -123,7 +126,8 @@ void monio::Monio::readIncrements(atlas::FieldSet& localFieldSet, auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace); // Initialise file - int variableConvention = initialiseFile(grid, filePath); + const int variableConvention = initialiseFile(grid, filePath, false, + lfricAtlasMapMethod); // getFileData returns a copy of FileData (with required LFRic mesh data), so read data // is discarded when FileData goes out-of-scope for reading subsequent fields. FileData fileData = getFileData(grid.name()); @@ -329,7 +333,8 @@ void monio::Monio::closeFiles() { int monio::Monio::initialiseFile(const atlas::Grid& grid, const std::string& filePath, - bool doCreateDateTimes) { + const bool doCreateDateTimes, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod) { oops::Log::trace() << "Monio::initialiseFile()" << std::endl; int variableConvention = consts::eLfricConvention; // LFRic convention is default if (mpiCommunicator_.rank() == mpiRankOwner_) { @@ -343,7 +348,7 @@ int monio::Monio::initialiseFile(const atlas::Grid& grid, reader_.readFullDatum(fileData, std::string(consts::kVerticalFullName)); reader_.readFullDatum(fileData, std::string(consts::kVerticalHalfName)); // Process read data - createLfricAtlasMap(fileData, grid); + createLfricAtlasMap(fileData, grid, lfricAtlasMapMethod); if (doCreateDateTimes == true) { reader_.readFullDatum(fileData, std::string(consts::kTimeVarName)); createDateTimes(fileData, @@ -390,7 +395,9 @@ monio::FileData monio::Monio::getFileData(const std::string& gridName) { return FileData(); // This function is called by all PEs. A return is essential. } -void monio::Monio::createLfricAtlasMap(FileData& fileData, const atlas::Grid& grid) { +void monio::Monio::createLfricAtlasMap(FileData& fileData, + const atlas::Grid& grid, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod) { oops::Log::trace() << "Monio::createLfricAtlasMap()" << std::endl; if (mpiCommunicator_.rank() == mpiRankOwner_) { if (fileData.getLfricAtlasMap().size() == 0) { @@ -399,7 +406,8 @@ void monio::Monio::createLfricAtlasMap(FileData& fileData, const atlas::Grid& gr reader_.getCoordData(fileData, consts::kLfricCoordVarNames); std::vector lfricCoords = utilsatlas::getLfricCoords(coordData); std::vector atlasCoords = utilsatlas::getAtlasCoords(grid); - fileData.setLfricAtlasMap(utilsatlas::createLfricAtlasMap(atlasCoords, lfricCoords)); + fileData.setLfricAtlasMap(utilsatlas::createLfricAtlasMap(atlasCoords, lfricCoords, + lfricAtlasMapMethod)); } } } diff --git a/src/monio/Monio.h b/src/monio/Monio.h index 6ecb7c8..95322ca 100644 --- a/src/monio/Monio.h +++ b/src/monio/Monio.h @@ -16,6 +16,7 @@ #include "AtlasReader.h" #include "AtlasWriter.h" +#include "Constants.h" #include "FileData.h" #include "Reader.h" #include "Writer.h" @@ -41,12 +42,16 @@ class Monio { void readState(atlas::FieldSet& localFieldSet, const std::vector& fieldMetadataVec, const std::string& filePath, - const util::DateTime& dateTime); + const util::DateTime& dateTime, + const consts::eLfricAtlasMapMethods& lfricAtlasMap = + consts::eLfricAtlasMapMethods::eKDtree); /// \brief Reads files without a time component, i.e. increment files. void readIncrements(atlas::FieldSet& localFieldSet, const std::vector& fieldMetadataVec, - const std::string& filePath); + const std::string& filePath, + const consts::eLfricAtlasMapMethods& lfricAtlasMap = + consts::eLfricAtlasMapMethods::eKDtree); /// \brief Writes increment files. No time component but the variables can use JEDI or LFRic write /// names. @@ -73,7 +78,9 @@ class Monio { /// it's called from LFRic-Lite. int initialiseFile(const atlas::Grid& grid, const std::string& filePath, - bool doCreateDateTimes = false); + const bool doCreateDateTimes = false, + const consts::eLfricAtlasMapMethods& lfricAtlasMap = + consts::eLfricAtlasMapMethods::eKDtree); private: /// \brief Private class constructor to prevent instantiation outside of the singleton. @@ -89,7 +96,10 @@ class Monio { FileData getFileData(const std::string& gridName); /// \brief Creates and stores a map between Atlas and LFRic horizontal ordering. - void createLfricAtlasMap(FileData& fileData, const atlas::Grid& grid); + void createLfricAtlasMap(FileData& fileData, + const atlas::Grid& grid, + const consts::eLfricAtlasMapMethods& lfricAtlasMap = + consts::eLfricAtlasMapMethods::eKDtree); /// \brief Creates and stores date-times from a state file. void createDateTimes(FileData& fileData, diff --git a/src/monio/UtilsAtlas.cc b/src/monio/UtilsAtlas.cc index 2a784de..eed64e4 100644 --- a/src/monio/UtilsAtlas.cc +++ b/src/monio/UtilsAtlas.cc @@ -98,7 +98,8 @@ std::vector> convertLatLonToContainers } std::vector createLfricAtlasMap(const std::vector& atlasCoords, - const std::vector& lfricCoords) { + const std::vector& lfricCoords, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod) { std::vector lfricAtlasMap; // Essential check to ensure grid is configured to accommodate the data if (atlasCoords.size() != lfricCoords.size()) { @@ -108,22 +109,26 @@ std::vector createLfricAtlasMap(const std::vector& a } lfricAtlasMap.reserve(atlasCoords.size()); - // Make a kd-tree using atlasLonLat as the point, - // with element index i as payload std::vector indices(atlasCoords.size()); std::iota(begin(indices), end(indices), 0); - const atlas::Geometry unitSphere(1.0); - atlas::util::IndexKDTree tree(unitSphere); - tree.build(atlasCoords, indices); - - // Partitioning vector to create atlas::Distribution for mesh generation - std::vector partitioning(tree.size(), -1); - - // find atlas global indices for each element of modelLonLat - for (const auto& lfricCoord : lfricCoords) { - auto idx = tree.closestPoint(lfricCoord).payload(); - lfricAtlasMap.push_back(idx); + // Trivially map between the LFRic and Atlas indices. + if (lfricAtlasMapMethod == consts::eLfricAtlasMapMethods::eTrivial) { + lfricAtlasMap = indices; + } else if (lfricAtlasMapMethod == consts::eLfricAtlasMapMethods::eKDtree) { + // Make a kd-tree using atlasLonLat as the point, + // with element index i as payload + const atlas::Geometry unitSphere(1.0); + atlas::util::IndexKDTree tree(unitSphere); + tree.build(atlasCoords, indices); + + // find atlas global indices for each element of modelLonLat + for (const auto& lfricCoord : lfricCoords) { + const auto idx = tree.closestPoint(lfricCoord).payload(); + lfricAtlasMap.push_back(idx); + } + } else { + utils::throwException("utilsatlas::createLfricAtlasMap()> Invalid LFRic-Atlas map chosen..."); } return lfricAtlasMap; } diff --git a/src/monio/UtilsAtlas.h b/src/monio/UtilsAtlas.h index fcd3866..2b4c5ad 100644 --- a/src/monio/UtilsAtlas.h +++ b/src/monio/UtilsAtlas.h @@ -37,7 +37,8 @@ namespace utilsatlas { const std::vector& coordNames); std::vector createLfricAtlasMap(const std::vector& atlasCoords, - const std::vector& lfricCoords); + const std::vector& lfricCoords, + const consts::eLfricAtlasMapMethods& lfricAtlasMapMethod); atlas::FieldSet getGlobalFieldSet(const atlas::FieldSet& fieldSet);