From 3aed489c6cbb56137260860c63c4d32b9bd8029b Mon Sep 17 00:00:00 2001
From: Norbert Podhorszki <pnorbert@ornl.gov>
Date: Fri, 26 Apr 2024 09:41:34 -0400
Subject: [PATCH] Add attributes to ADIOS2 output to "define" dimensions as
 names. We need this for xarray support that is designed for the NetCDF model,
 where dimensions are separately defined entities.

---
 include/bout/adios_object.hxx     |  8 +++++++-
 src/sys/options/options_adios.cxx | 24 ++++++++++++++++++------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/include/bout/adios_object.hxx b/include/bout/adios_object.hxx
index b14316f1ba..12da8ce4c7 100755
--- a/include/bout/adios_object.hxx
+++ b/include/bout/adios_object.hxx
@@ -57,11 +57,17 @@ public:
   }
 
   template <class T>
-  adios2::Variable<T> GetArrayVariable(const std::string& varname, adios2::Dims& shape) {
+  adios2::Variable<T> GetArrayVariable(const std::string& varname, adios2::Dims& shape,
+                                       const std::vector<std::string>& dimNames,
+                                       int rank) {
     adios2::Variable<T> v = io.InquireVariable<T>(varname);
     if (!v) {
       adios2::Dims start(shape.size());
       v = io.DefineVariable<T>(varname, shape, start, shape);
+      if (!rank && dimNames.size()) {
+        io.DefineAttribute<std::string>("__xarray_dimensions__", dimNames.data(),
+                                        dimNames.size(), varname, "/", true);
+      }
     } else {
       v.SetShape(shape);
     }
diff --git a/src/sys/options/options_adios.cxx b/src/sys/options/options_adios.cxx
index b3acbaada6..09797647ed 100644
--- a/src/sys/options/options_adios.cxx
+++ b/src/sys/options/options_adios.cxx
@@ -307,6 +307,12 @@ void OptionsADIOS::verifyTimesteps() const {
   return;
 }
 
+const std::vector<std::string> DIMS_NONE;
+const std::vector<std::string> DIMS_X = {"x"};
+const std::vector<std::string> DIMS_XY = {"x", "y"};
+const std::vector<std::string> DIMS_XZ = {"x", "z"};
+const std::vector<std::string> DIMS_XYZ = {"x", "y", "z"};
+
 /// Visit a variant type, and put the data into a NcVar
 struct ADIOSPutVarVisitor {
   ADIOSPutVarVisitor(const std::string& name, ADIOSStream& stream)
@@ -388,7 +394,8 @@ void ADIOSPutVarVisitor::operator()<Field2D>(const Field2D& value) {
   adios2::Dims memCount = {static_cast<size_t>(value.getNx()),
                            static_cast<size_t>(value.getNy())};
 
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_XY, BoutComm::rank());
   var.SetSelection({start, count});
   var.SetMemorySelection({memStart, memCount});
   stream.engine.Put<BoutReal>(var, &value(0, 0));
@@ -425,7 +432,8 @@ void ADIOSPutVarVisitor::operator()<Field3D>(const Field3D& value) {
                            static_cast<size_t>(value.getNy()),
                            static_cast<size_t>(value.getNz())};
 
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_XYZ, BoutComm::rank());
   var.SetSelection({start, count});
   var.SetMemorySelection({memStart, memCount});
   stream.engine.Put<BoutReal>(var, &value(0, 0, 0));
@@ -457,7 +465,8 @@ void ADIOSPutVarVisitor::operator()<FieldPerp>(const FieldPerp& value) {
   adios2::Dims memCount = {static_cast<size_t>(value.getNx()),
                            static_cast<size_t>(value.getNz())};
 
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_XZ, BoutComm::rank());
   var.SetSelection({start, count});
   var.SetMemorySelection({memStart, memCount});
   stream.engine.Put<BoutReal>(var, &value(0, 0));
@@ -469,7 +478,8 @@ void ADIOSPutVarVisitor::operator()<Array<BoutReal>>(const Array<BoutReal>& valu
   adios2::Dims shape = {(size_t)BoutComm::size(), (size_t)value.size()};
   adios2::Dims start = {(size_t)BoutComm::rank(), 0};
   adios2::Dims count = {1, shape[1]};
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_NONE, BoutComm::rank());
   var.SetSelection({start, count});
   stream.engine.Put<BoutReal>(var, value.begin());
 }
@@ -482,7 +492,8 @@ void ADIOSPutVarVisitor::operator()<Matrix<BoutReal>>(const Matrix<BoutReal>& va
                         (size_t)std::get<1>(s)};
   adios2::Dims start = {(size_t)BoutComm::rank(), 0, 0};
   adios2::Dims count = {1, shape[1], shape[2]};
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_NONE, BoutComm::rank());
   var.SetSelection({start, count});
   stream.engine.Put<BoutReal>(var, value.begin());
 }
@@ -495,7 +506,8 @@ void ADIOSPutVarVisitor::operator()<Tensor<BoutReal>>(const Tensor<BoutReal>& va
                         (size_t)std::get<1>(s), (size_t)std::get<2>(s)};
   adios2::Dims start = {(size_t)BoutComm::rank(), 0, 0, 0};
   adios2::Dims count = {1, shape[1], shape[2], shape[3]};
-  adios2::Variable<BoutReal> var = stream.GetArrayVariable<BoutReal>(varname, shape);
+  adios2::Variable<BoutReal> var =
+      stream.GetArrayVariable<BoutReal>(varname, shape, DIMS_NONE, BoutComm::rank());
   var.SetSelection({start, count});
   stream.engine.Put<BoutReal>(var, value.begin());
 }