From 36b116fba270442906feb035d3b629827d213fd1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Franz=20P=C3=B6schel?= <franz.poeschel@gmail.com>
Date: Tue, 6 Feb 2024 12:27:23 +0100
Subject: [PATCH] Python test

---
 examples/5_write_parallel.py           | 21 +++++++--
 src/binding/python/RecordComponent.cpp | 61 +++++++++++++++++++-------
 2 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/examples/5_write_parallel.py b/examples/5_write_parallel.py
index ace0cd6e63..9ad07e8554 100644
--- a/examples/5_write_parallel.py
+++ b/examples/5_write_parallel.py
@@ -14,6 +14,13 @@
 import numpy as np
 import openpmd_api as io
 
+try:
+    import adios2
+    from packaging import version
+    USE_JOINED_DIMENSION = \
+        version.parse(adios2.__version__) >= version.parse('2.9.0')
+except ImportError:
+    USE_JOINED_DIMENSION = False
 if __name__ == "__main__":
     # also works with any other MPI communicator
     comm = MPI.COMM_WORLD
@@ -29,7 +36,9 @@
 
     # open file for writing
     series = io.Series(
-        "../samples/5_parallel_write_py.h5",
+        "../samples/5_parallel_write_py.bp"
+            if USE_JOINED_DIMENSION
+            else "../samples/5_parallel_write_py.bp",
         io.Access.create,
         comm
     )
@@ -51,7 +60,9 @@
         meshes["mymesh"]
 
     # example 1D domain decomposition in first index
-    global_extent = [comm.size * 10, 300]
+    global_extent = [io.Dataset.JOINED_DIMENSION, 300] \
+        if USE_JOINED_DIMENSION else [comm.size * 10, 300]
+
     dataset = io.Dataset(local_data.dtype, global_extent)
 
     if 0 == comm.rank:
@@ -64,7 +75,11 @@
               "mymesh in iteration 1")
 
     # example shows a 1D domain decomposition in first index
-    mymesh[comm.rank*10:(comm.rank+1)*10, :] = local_data
+
+    if USE_JOINED_DIMENSION:
+        mymesh.store_chunk(local_data, [], [10, 300])
+    else:
+        mymesh[comm.rank*10:(comm.rank+1)*10, :] = local_data
     if 0 == comm.rank:
         print("Registered a single chunk per MPI rank containing its "
               "contribution, ready to write content to disk")
diff --git a/src/binding/python/RecordComponent.cpp b/src/binding/python/RecordComponent.cpp
index 37ad9a7cff..7399184383 100644
--- a/src/binding/python/RecordComponent.cpp
+++ b/src/binding/python/RecordComponent.cpp
@@ -265,23 +265,52 @@ inline void store_chunk(
                         "in record component (") +
             std::to_string(r_shape.size()) + std::string("D)"));
 
-    for (auto d = 0; d < a.ndim(); ++d)
+    if (auto joined_dim = r.joinedDimension(); joined_dim.has_value())
     {
-        // selection causes overflow of r
-        if (offset.at(d) + extent.at(d) > r_shape.at(d))
-            throw py::index_error(
-                std::string("slice ") + std::to_string(offset.at(d)) +
-                std::string(":") + std::to_string(extent.at(d)) +
-                std::string(" is out of bounds for axis ") + std::to_string(d) +
-                std::string(" with size ") + std::to_string(r_shape.at(d)));
-        // underflow of selection in r for given a
-        if (s_shape.at(d) != std::uint64_t(a.shape()[d]))
-            throw py::index_error(
-                std::string("size of chunk (") + std::to_string(a.shape()[d]) +
-                std::string(") for axis ") + std::to_string(d) +
-                std::string(" does not match selection ") +
-                std::string("size in record component (") +
-                std::to_string(s_extent.at(d)) + std::string(")"));
+        for (py::ssize_t d = 0; d < a.ndim(); ++d)
+        {
+            // selection causes overflow of r
+            if (d != py::ssize_t(*joined_dim) && extent.at(d) != r_shape.at(d))
+                throw py::index_error(
+                    std::string("selection for axis ") + std::to_string(d) +
+                    " of record component with joined dimension " +
+                    std::to_string(*joined_dim) +
+                    " must be equivalent to its global extent " +
+                    std::to_string(extent.at(d)) + ", but was " +
+                    std::to_string(r_shape.at(d)) + ".");
+            // underflow of selection in r for given a
+            if (s_shape.at(d) != std::uint64_t(a.shape()[d]))
+                throw py::index_error(
+                    std::string("size of chunk (") +
+                    std::to_string(a.shape()[d]) + std::string(") for axis ") +
+                    std::to_string(d) +
+                    std::string(" does not match selection ") +
+                    std::string("size in record component (") +
+                    std::to_string(s_extent.at(d)) + std::string(")"));
+        }
+    }
+    else
+    {
+        for (auto d = 0; d < a.ndim(); ++d)
+        {
+            // selection causes overflow of r
+            if (offset.at(d) + extent.at(d) > r_shape.at(d))
+                throw py::index_error(
+                    std::string("slice ") + std::to_string(offset.at(d)) +
+                    std::string(":") + std::to_string(extent.at(d)) +
+                    std::string(" is out of bounds for axis ") +
+                    std::to_string(d) + std::string(" with size ") +
+                    std::to_string(r_shape.at(d)));
+            // underflow of selection in r for given a
+            if (s_shape.at(d) != std::uint64_t(a.shape()[d]))
+                throw py::index_error(
+                    std::string("size of chunk (") +
+                    std::to_string(a.shape()[d]) + std::string(") for axis ") +
+                    std::to_string(d) +
+                    std::string(" does not match selection ") +
+                    std::string("size in record component (") +
+                    std::to_string(s_extent.at(d)) + std::string(")"));
+        }
     }
 
     check_buffer_is_contiguous(a);