Allow reading of states and increments for structured columns (#52)

Author: mo-jonasganderton

src/monio/Monio.cc

@@ -60,9 +60,9 @@ void monio::Monio::readState(atlas::FieldSet& localFieldSet,
           atlas::Field globalField = utilsatlas::getGlobalField(localField);
           if (mpiCommunicator_.rank() == mpiRankOwner_) {
             auto& functionSpace = globalField.functionspace();
-            auto& grid = atlas::functionspace::NodeColumns(functionSpace).mesh().grid();
+            auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace);
             // Initialise file
-            int variableConvention = initialiseFile(grid.name(), filePath, true);
+            int variableConvention = initialiseFile(grid, filePath, true);
             // 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());
@@ -120,10 +120,10 @@ void monio::Monio::readIncrements(atlas::FieldSet& localFieldSet,
           atlas::Field globalField = utilsatlas::getGlobalField(localField);
           if (mpiCommunicator_.rank() == mpiRankOwner_) {
             auto& functionSpace = globalField.functionspace();
-            auto& grid = atlas::functionspace::NodeColumns(functionSpace).mesh().grid();
+            auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace);
 
             // Initialise file
-            int variableConvention = initialiseFile(grid.name(), filePath);
+            int variableConvention = initialiseFile(grid, filePath);
             // 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());
@@ -172,7 +172,7 @@ void monio::Monio::writeIncrements(const atlas::FieldSet& localFieldSet,
   if (filePath.length() != 0) {
     try {
       auto& functionSpace = localFieldSet[0].functionspace();
-      auto& grid = atlas::functionspace::NodeColumns(functionSpace).mesh().grid();
+      auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace);
       FileData fileData = getFileData(grid.name());
       cleanFileData(fileData);  // Remove metadata required for reading, but not for writing.
       if (isLfricConvention == false) {
@@ -235,7 +235,7 @@ void monio::Monio::writeState(const atlas::FieldSet& localFieldSet,
   if (filePath.length() != 0) {
     try {
       auto& functionSpace = localFieldSet[0].functionspace();
-      auto& grid = atlas::functionspace::NodeColumns(functionSpace).mesh().grid();
+      auto grid = utilsatlas::getGridFromFunctionSpace(functionSpace);
       FileData fileData = getFileData(grid.name());
       cleanFileData(fileData);  // Remove metadata required for reading, but not for writing.
       if (isLfricConvention == false) {
@@ -390,7 +390,7 @@ 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::CubedSphereGrid& grid) {
+void monio::Monio::createLfricAtlasMap(FileData& fileData, const atlas::Grid& grid) {
   oops::Log::trace() << "Monio::createLfricAtlasMap()" << std::endl;
   if (mpiCommunicator_.rank() == mpiRankOwner_) {
     if (fileData.getLfricAtlasMap().size() == 0) {

src/monio/Monio.h

@@ -89,7 +89,7 @@ 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::CubedSphereGrid& grid);
+  void createLfricAtlasMap(FileData& fileData, const atlas::Grid& grid);
 
   /// \brief Creates and stores date-times from a state file.
   void createDateTimes(FileData& fileData,

src/monio/UtilsAtlas.cc

@@ -165,6 +165,19 @@ atlas::FieldSet getGlobalFieldSet(const atlas::FieldSet& fieldSet) {
   }
 }
 
+atlas::Grid getGridFromFunctionSpace(const atlas::FunctionSpace& functionSpace) {
+  auto nc = atlas::functionspace::NodeColumns(functionSpace);
+  auto sc = atlas::functionspace::StructuredColumns(functionSpace);
+  if (nc) {
+    return nc.mesh().grid();
+  } else if (sc) {
+    return sc.grid();
+  } else {
+    utils::throwException("Monio::readState()> FunctionSpace not an accepted type. "
+                          "Accepted types: NodeColumns, StructuredColumns");
+  }
+}
+
 atlas::idx_t getHorizontalSize(const atlas::Field& field) {
   atlas::Field ghostField = field.functionspace().ghost();
   atlas::idx_t size = 0;

src/monio/UtilsAtlas.h

@@ -43,6 +43,8 @@ namespace utilsatlas {
 
   atlas::Field getGlobalField(const atlas::Field& field);
 
+  atlas::Grid getGridFromFunctionSpace(const atlas::FunctionSpace& functionSpace);
+
   atlas::idx_t getHorizontalSize(const atlas::Field& field);  // Just 2D size. Any field.
   atlas::idx_t getGlobalDataSize(const atlas::Field& field);  // Full 3D size of global field.