Exchange halos with XIOS

core/src/ModelMetadata.cpp

@@ -15,9 +15,8 @@
 #include "include/ParallelNetcdfFile.hpp"
 #ifdef USE_XIOS
 #include "include/Xios.hpp"
-#else
-#include "include/Halo.hpp"
 #endif
+#include "include/Halo.hpp"
 #endif
 #include "include/gridNames.hpp"
 #include <array>
@@ -311,12 +310,8 @@ void ModelMetadata::setDimensionsFromFile(const std::string& filename)
                 localLength = dim.getSize();
                 start = 0;
             }
-#if USE_XIOS
-            ModelArray::setDimension(dimType, dim.getSize(), localLength, start);
-#else
             ModelArray::setDimension(
                 dimType, dim.getSize(), localLength + 2 * Halo::haloWidth, start);
-#endif
 #else
             ModelArray::setDimension(dimType, dim.getSize());
 #endif

core/src/Xios.cpp

@@ -19,6 +19,7 @@
 
 #include "StructureModule/include/ParametricGrid.hpp"
 #include "include/Finalizer.hpp"
+#include "include/Halo.hpp"
 #include "include/ModelMPI.hpp"
 #include "include/ModelMetadata.hpp"
 #include "include/ParallelNetcdfFile.hpp"
@@ -1122,6 +1123,24 @@ void Xios::setDiagnosticFieldType(const std::string& fieldId, const ModelArray::
     setFieldType(fieldId, fieldType, DIAGNOSTIC);
 }
 
+/*!
+ * Get the rank-local degrees of freedom for a given ModelArray.
+ *
+ * @param modelarray The ModelArray to get DoF information for
+ * @return MultiDim vector containing the number of DoFs in each direction
+ */
+const ModelArray::MultiDim Xios::getFieldLocalDoFs(const ModelArray& modelarray)
+{
+    const ModelArray::Type& type = modelarray.getType();
+    const std::string& domainId = domainIds[type];
+    xios::CDomain* domain = getDomain(domainId);
+    int ni;
+    cxios_get_domain_ni(domain, &ni);
+    int nj;
+    cxios_get_domain_nj(domain, &nj);
+    return std::vector<size_t> { (size_t)ni, (size_t)nj };
+}
+
 /*!
  * @brief   Do an initial read of input files to deduce field dimensions.
  *
@@ -1633,7 +1652,6 @@ void Xios::write(const std::string& fieldId, const ModelArray& modelarray)
     if (modelarray.nDimensions() != 2) {
         throw std::invalid_argument("Only ModelArrays of dimension 2 are supported");
     }
-    auto& dims = modelarray.dimensions();
     const ModelArray::Type& type = modelarray.getType();
     domainWritten[domainIds[type]] = true;
 
@@ -1644,19 +1662,28 @@ void Xios::write(const std::string& fieldId, const ModelArray& modelarray)
             "Xios::write: field '" + fieldId + "' does not have the expected type");
     }
 
+    // Create a halo object based on the input array
+    Halo halo(modelarray);
+
+    // Copy data into a temporary array without the halo
+    ModelArray::DataType tempData;
+    tempData.resize(halo.getInnerSize(), modelarray.nComponents());
+    halo.getInnerBlock(modelarray.data(), tempData);
+    auto& dims = halo.getInnerShape();
+
     // Write out according to field type
     if ((type == ModelArray::Type::H) || (type == ModelArray::Type::CG)) {
         cxios_write_data_k82(
-            fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0], dims[1], -1);
+            fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1], -1);
     } else if (type == ModelArray::Type::VERTEX) {
-        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::NCOORDS), -1);
+        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::NCOORDS), -1);
     } else if (type == ModelArray::Type::DG) {
-        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::DG), -1);
+        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::DG), -1);
     } else if (type == ModelArray::Type::DGSTRESS) {
-        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::DGSTRESS), -1);
+        cxios_write_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::DGSTRESS), -1);
     } else {
         throw std::invalid_argument(
             "Only HFields, VertexFields, DGFields, DGSFields, and CGFields are supported");
@@ -1681,6 +1708,9 @@ void Xios::read(const std::string& fieldId, ModelArray& modelarray)
     const ModelArray::Type& type = modelarray.getType();
     const ModelArray::Type& expectedType = getFieldType(fieldId);
 
+    // Create a halo object based on the input array
+    Halo halo(modelarray);
+
     // Account for fields to be read in as HField but converted to DGField
     if (inputFieldsToConvert.count(fieldId)) {
         if (expectedType != ModelArray::Type::H) {
@@ -1691,38 +1721,58 @@ void Xios::read(const std::string& fieldId, ModelArray& modelarray)
             throw std::runtime_error(
                 "Xios::read: field '" + fieldId + "' was expected to be converted to a DGField");
         }
+
+        // Create a HField to read into
         HField inputarray(ModelArray::Type::H);
-        auto& dims = inputarray.dimensions();
-        cxios_read_data_k82(
-            fieldId.c_str(), fieldId.length(), inputarray.getData(), dims[0], dims[1]);
+
+        // Read into a temporary array without the halo
+        ModelArray::DataType tempData;
+        tempData.resize(halo.getInnerSize(), inputarray.nComponents());
+        auto& dims = halo.getInnerShape();
+        cxios_read_data_k82(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1]);
+
+        // Set the inner block of the array based on what was read and exchange halos
+        halo.setInnerBlock(tempData, inputarray.getDataRef());
+        halo.exchangeHalos(inputarray.getDataRef());
+
+        // Set the DGField based on the contents of the HField
         modelarray = 0;
         // FIXME: Conversion with overloaded '=' operator is known to be problematic
         modelarray = inputarray;
         return;
     }
 
     // Other field types should not need converting
-    auto& dims = modelarray.dimensions();
     if (type != expectedType) {
         throw std::runtime_error(
             "Xios::read: field '" + fieldId + "' does not have the expected type");
     }
-    if ((type == ModelArray::Type::H) || (type == ModelArray::Type::CG)) {
-        cxios_read_data_k82(
-            fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0], dims[1]);
+
+    // Read into a temporary array without the halo
+    ModelArray::DataType tempData;
+    tempData.resize(halo.getInnerSize(), modelarray.nComponents());
+    auto& dims = halo.getInnerShape();
+    if (type == ModelArray::Type::H) {
+        cxios_read_data_k82(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1]);
+    } else if (type == ModelArray::Type::CG) {
+        cxios_read_data_k82(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1]);
     } else if (type == ModelArray::Type::VERTEX) {
-        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::NCOORDS));
+        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::NCOORDS));
     } else if (type == ModelArray::Type::DG) {
-        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::DG));
+        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::DG));
     } else if (type == ModelArray::Type::DGSTRESS) {
-        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), modelarray.getData(), dims[0],
-            dims[1], ModelArray::size(ModelArray::Dimension::DGSTRESS));
+        cxios_read_data_k83(fieldId.c_str(), fieldId.length(), tempData.data(), dims[0], dims[1],
+            ModelArray::size(ModelArray::Dimension::DGSTRESS));
     } else {
         throw std::invalid_argument(
             "Only HFields, VertexFields, DGFields, DGSFields, and CGFields are supported");
     }
+
+    // Set the inner block of the array based on what was read and exchange halos
+    halo.setInnerBlock(tempData, modelarray.getDataRef());
+    halo.exchangeHalos(modelarray.getDataRef());
 }
 }
 

core/src/include/Halo.hpp

@@ -53,6 +53,18 @@ class Halo {
         initializeHaloMetadata();
     }
 
+    /*!
+     * @brief Constructs a halo object from constant ModelArray
+     * @param ma ModelArray object to create halo from
+     */
+    Halo(const ModelArray& ma)
+    {
+        m_numComps = ma.nComponents();
+        isVertex = ma.getType() == ModelArray::Type::VERTEX;
+        setSpatialDims();
+        initializeHaloMetadata();
+    }
+
     /*!
      * @brief Constructs a halo object from DGVectorHolder
      * @param dgvh DGVectorHolder object to create halo from
@@ -434,6 +446,11 @@ class Halo {
      */
     size_t getInnerSize() { return m_innerNx * m_innerNy; }
 
+    /*!
+     * @brief Returns shape of the inner array
+     */
+    const MultiDim getInnerShape() const { return std::vector({ m_innerNx, m_innerNy }); }
+
     /*!
      * @brief Get inner block from source array and copy into target
      *

core/src/include/Xios.hpp

@@ -80,6 +80,7 @@ class Xios : public Configured<Xios> {
         = { sstName, sssName, mldName, uName, vName, sshName };
     void setPrognosticFieldType(const std::string& fieldId, const ModelArray::Type& type);
     void setDiagnosticFieldType(const std::string& fieldId, const ModelArray::Type& type);
+    const ModelArray::MultiDim getFieldLocalDoFs(const ModelArray& modelarray);
 
     enum {
         ENABLED_KEY,

core/src/include/xios_c_interface.hpp

@@ -83,15 +83,17 @@ void cxios_set_domain_ni(xios::CDomain* domain_hdl, int ni);
 void cxios_set_domain_nj(xios::CDomain* domain_hdl, int nj);
 void cxios_set_domain_ibegin(xios::CDomain* domain_hdl, int ibegin);
 void cxios_set_domain_jbegin(xios::CDomain* domain_hdl, int jbegin);
+void cxios_set_domain_lonvalue_1d(xios::CDomain* domain_hdl, double* data, const int* ni);
+void cxios_set_domain_latvalue_1d(xios::CDomain* domain_hdl, double* data, const int* nj);
+int cxios_get_domain_ni(xios::CDomain* domain_hdl, int* ni);
+int cxios_get_domain_nj(xios::CDomain* domain_hdl, int* nj);
 bool cxios_is_defined_domain_type(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_dim_i_name(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_dim_j_name(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_lat_name(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_lon_name(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_ni_glo(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_nj_glo(xios::CDomain* domain_hdl);
-void cxios_set_domain_lonvalue_1d(xios::CDomain* domain_hdl, double* data, const int* ni);
-void cxios_set_domain_latvalue_1d(xios::CDomain* domain_hdl, double* data, const int* nj);
 bool cxios_is_defined_domain_ni(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_nj(xios::CDomain* domain_hdl);
 bool cxios_is_defined_domain_ibegin(xios::CDomain* domain_hdl);

core/test/CMakeLists.txt

@@ -223,6 +223,8 @@ if(ENABLE_MPI)
                 "${MODEL_INCLUDE_DIR}"
                 "${XIOS_INCLUDE_LIST}"
                 "${ModulesRoot}/StructureModule"
+                "../../dynamics/src"
+                "../../dynamics/src/include"
         )
         target_link_libraries(testXiosReadForcing_MPI2 PRIVATE nextsimlib doctest::doctest)
 
@@ -239,6 +241,9 @@ if(ENABLE_MPI)
                 "${MODEL_INCLUDE_DIR}"
                 "${XIOS_INCLUDE_LIST}"
                 "${ModulesRoot}/StructureModule"
+                "${ModulesRoot}/StructureModule"
+                "../../dynamics/src"
+                "../../dynamics/src/include"
         )
         target_link_libraries(testXiosReadDiagnostic_MPI3 PRIVATE nextsimlib doctest::doctest)
 

core/test/XiosReadDiagnostic_test.cpp

@@ -6,7 +6,9 @@
 #include <doctest/extensions/doctest_mpi.h>
 #undef INFO
 
+#include "cgVector.hpp"
 #include "include/Finalizer.hpp"
+#include "include/Halo.hpp"
 #include "include/Model.hpp"
 #include "include/ModelMPI.hpp"
 #include "include/NextsimModule.hpp"
@@ -70,19 +72,24 @@ MPI_TEST_CASE("TestXiosReadDiagnostic", 3)
     Xios& xiosHandler = Xios::getInstance();
     REQUIRE(xiosHandler.getCalendarStep() == 0);
 
-    // Deduce the local lengths of the two dimensions
-    const size_t nx = ModelArray::size(ModelArray::Dimension::X);
-    const size_t ny = ModelArray::size(ModelArray::Dimension::Y);
-
     // Read restarts from file and check they take the expected values
     // NOTE: The ParametricGrid is created and the XIOS context definition is closed in the call to
     //       StructureFactory::stateFromFile()
     for (const auto [fieldName, modelarray] :
         StructureFactory::stateFromFile(diagnosticFilename).data) {
         REQUIRE(fieldName == hsnowName);
+        // Extract the inner block of the ModelArray
+        Halo halo(modelarray);
+        ModelArray::DataType tempData;
+        tempData.resize(halo.getInnerSize(), modelarray.nComponents());
+        halo.getInnerBlock(modelarray.data(), tempData);
+        // Check that the inner block has the expected values
+        auto& ndofs = xiosHandler.getFieldLocalDoFs(modelarray);
+        auto nx = ndofs[0];
+        auto ny = ndofs[1];
         for (size_t j = 0; j < ny; ++j) {
             for (size_t i = 0; i < nx; ++i) {
-                REQUIRE(modelarray(i, j) == doctest::Approx(0.15));
+                REQUIRE(tempData(i, j) == doctest::Approx(0.15));
             }
         }
     }

core/test/XiosReadForcing_test.cpp

@@ -7,7 +7,9 @@
 #undef INFO
 
 #include "StructureModule/include/ParametricGrid.hpp"
+#include "cgVector.hpp"
 #include "include/Finalizer.hpp"
+#include "include/Halo.hpp"
 #include "include/Model.hpp"
 #include "include/ModelMPI.hpp"
 #include "include/NextsimModule.hpp"
@@ -83,8 +85,6 @@ MPI_TEST_CASE("TestXiosReadForcing", 2)
     REQUIRE(xiosHandler.getCalendarStep() == 0);
 
     // Deduce the local lengths of the two dimensions
-    const size_t nx = ModelArray::size(ModelArray::Dimension::X);
-    const size_t ny = ModelArray::size(ModelArray::Dimension::Y);
     REQUIRE(ModelArray::nComponents(ModelArray::Type::DG) == DGCOMP);
     REQUIRE(ModelArray::size(ModelArray::Dimension::DG) == DGCOMP);
 
@@ -101,9 +101,18 @@ MPI_TEST_CASE("TestXiosReadForcing", 2)
         // Read ERA5 forcings from file and check they take the expected values
         for (const auto& [fieldName, modelarray] :
             pio->readForcingTimeStatic(era5ForcingFieldNames, time, era5ForcingFilename).data) {
+            // Extract the inner block of the ModelArray
+            Halo halo(modelarray);
+            ModelArray::DataType tempData;
+            tempData.resize(halo.getInnerSize(), modelarray.nComponents());
+            halo.getInnerBlock(modelarray.data(), tempData);
+            // Check that the inner block has the expected values
+            auto& ndofs = xiosHandler.getFieldLocalDoFs(modelarray);
+            auto nx = ndofs[0];
+            auto ny = ndofs[1];
             for (size_t j = 0; j < ny; ++j) {
                 for (size_t i = 0; i < nx; ++i) {
-                    REQUIRE(modelarray(i, j) == doctest::Approx(0.1 * ts));
+                    REQUIRE(tempData(i, j) == doctest::Approx(0.1 * ts));
                 }
             }
         }
@@ -114,9 +123,18 @@ MPI_TEST_CASE("TestXiosReadForcing", 2)
             for (const auto& [fieldName, modelarray] :
                 pio->readForcingTimeStatic(topazForcingFieldNames, time, topazForcingFilename)
                     .data) {
+                // Extract the inner block of the ModelArray
+                Halo halo(modelarray);
+                ModelArray::DataType tempData;
+                tempData.resize(halo.getInnerSize(), modelarray.nComponents());
+                halo.getInnerBlock(modelarray.data(), tempData);
+                // Check that the inner block has the expected values
+                auto& ndofs = xiosHandler.getFieldLocalDoFs(modelarray);
+                auto nx = ndofs[0];
+                auto ny = ndofs[1];
                 for (size_t j = 0; j < ny; ++j) {
                     for (size_t i = 0; i < nx; ++i) {
-                        REQUIRE(modelarray(i, j) == doctest::Approx(ts + 1));
+                        REQUIRE(tempData(i, j) == doctest::Approx(ts + 1));
                     }
                 }
             }

core/test/XiosWriteDiagnostic_test.cpp

@@ -64,8 +64,6 @@ MPI_TEST_CASE("TestXiosWriteDiagnostic", 2)
     Xios& xiosHandler = Xios::getInstance();
 
     // Set ModelArray dimensions
-    const size_t nx = ModelArray::size(ModelArray::Dimension::X);
-    const size_t ny = ModelArray::size(ModelArray::Dimension::Y);
     REQUIRE(ModelArray::nComponents(ModelArray::Type::DG) == DGCOMP);
     REQUIRE(ModelArray::size(ModelArray::Dimension::DG) == DGCOMP);
 
@@ -92,11 +90,15 @@ MPI_TEST_CASE("TestXiosWriteDiagnostic", 2)
     // Simulate 4 iterations (timesteps)
     ModelMetadata& metadata = ModelMetadata::getInstance();
     const Duration& timestep = metadata.stepLength();
+    const size_t nx = ModelArray::size(ModelArray::Dimension::X);
+    const size_t ny = ModelArray::size(ModelArray::Dimension::Y);
     int rank;
     MPI_Comm_rank(test_comm, &rank);
     for (int ts = 0; ts <= 4; ts++) {
 
         // Update diagnostics
+        // NOTE: This example writes the (spatially constant) value into the halo region as well as
+        //       the interior. It doesn't matter because the halo values get discarded.
         for (size_t j = 0; j < ny; ++j) {
             for (size_t i = 0; i < nx; ++i) {
                 hsnow(i, j) = 0.1 * ts;