python/hl: DoRead() now handles the step selection case, too.

Essentially, all the work is now done in DoRead, which may optionally include
a Selection and a StepSelection (by default, stepsCount == 0, which means no
StepSelection, return data from current step).

bindings/Python/py11File.cpp

@@ -230,31 +230,14 @@ pybind11::array File::Read(const std::string &name, const Dims &start,
         std::copy(value.begin(), value.end(), pyArray.mutable_data());
         return pyArray;
     }
-#define declare_type(T)                                                        \
-    else if (type == helper::GetType<T>())                                     \
-    {                                                                          \
-        return DoRead<T>(name, start, count, blockID);                         \
-    }
-    ADIOS2_FOREACH_NUMPY_TYPE_1ARG(declare_type)
-#undef declare_type
 
-    throw std::invalid_argument(
-        "ERROR: adios2 file read variable " + name +
-        ", type can't be mapped to a numpy type, in call to read\n");
+    return Read(name, start, count, 0, 0, blockID);
 }
 
 pybind11::array File::Read(const std::string &name, const Dims &start,
                            const Dims &count, const size_t stepStart,
                            const size_t stepCount, const size_t blockID)
 {
-    // shape of the returned numpy array
-    Dims shapePy(count.size() + 1);
-    shapePy[0] = stepCount;
-    for (auto i = 1; i < shapePy.size(); ++i)
-    {
-        shapePy[i] = count[i - 1];
-    }
-
     const std::string type = m_Stream->m_IO->InquireVariableType(name);
 
     if (type.empty())
@@ -263,11 +246,7 @@ pybind11::array File::Read(const std::string &name, const Dims &start,
 #define declare_type(T)                                                        \
     else if (type == helper::GetType<T>())                                     \
     {                                                                          \
-        pybind11::array_t<T> pyArray(shapePy);                                 \
-        m_Stream->Read<T>(name, pyArray.mutable_data(),                        \
-                          Box<Dims>(start, count),                             \
-                          Box<size_t>(stepStart, stepCount), blockID);         \
-        return pyArray;                                                        \
+        return DoRead<T>(name, start, count, stepStart, stepCount, blockID);   \
     }
     ADIOS2_FOREACH_NUMPY_TYPE_1ARG(declare_type)
 #undef declare_type

bindings/Python/py11File.h

@@ -124,7 +124,8 @@ class File
 
     template <class T>
     pybind11::array DoRead(const std::string &name, const Dims &start,
-                           const Dims &count, const size_t blockID);
+                           const Dims &count, const size_t stepStart,
+                           const size_t stepCount, const size_t blockID);
 };
 
 } // end namespace py11

bindings/Python/py11File.tcc

@@ -20,53 +20,78 @@ namespace py11
 
 template <class T>
 pybind11::array File::DoRead(const std::string &name, const Dims &_start,
-                             const Dims &_count, const size_t blockID)
+                             const Dims &_count, const size_t stepStart,
+                             const size_t stepCount, const size_t blockID)
 {
     core::Variable<T> &variable = *m_Stream->m_IO->InquireVariable<T>(name);
     Dims &shape = variable.m_Shape;
-
     Dims start = _start;
-    if (start.empty())
-    {
-        // default start to be (0, 0, ...)
-        start = Dims(shape.size());
-    }
-
     Dims count = _count;
-    if (count.empty())
+
+    if (variable.m_ShapeID == ShapeID::GlobalValue)
     {
-        if (variable.m_ShapeID == ShapeID::GlobalArray)
+        if (!(_start.empty() && _count.empty()))
         {
-            // default count is everything (shape of whole array)
-            count = shape;
+            throw std::invalid_argument("when reading a scalar, start and "
+                                        "count cannot be specified.\n");
         }
-        else if (variable.m_ShapeID == ShapeID::LocalArray)
+        if (stepCount == 0)
         {
-            variable.SetBlockSelection(blockID);
-            count = variable.Count();
+            // for compatiblity, return 1-d arrays rather than 0-d
+            count = Dims{1};
         }
     }
 
-    if (variable.m_ShapeID == ShapeID::GlobalValue)
+    if (variable.m_ShapeID == ShapeID::LocalArray)
     {
-        count = Dims{1};
+        variable.SetBlockSelection(blockID);
     }
-    pybind11::array_t<T> pyArray(count);
-    if (variable.m_ShapeID == ShapeID::GlobalValue)
+    else
     {
-        if (!(_start.empty() && _count.empty()))
+        if (blockID != 0)
         {
-            throw std::invalid_argument("when reading a scalar, start and "
-                                        "count cannot be specified.\n");
+            throw std::invalid_argument(
+                "blockId can only be specified when reading LocalArrays.");
         }
-        m_Stream->Read<T>(name, pyArray.mutable_data(), blockID);
     }
-    else
+
+    if (start.empty())
+    {
+        // default start to be (0, 0, ...)
+        start = Dims(shape.size());
+    }
+
+    if (count.empty())
+    {
+        // does the right thing for global and local arrays
+        count = variable.Count();
+    }
+
+    // make numpy array, shape is count, possibly with extra dim for step added
+    Dims shapePy;
+    shapePy.reserve((stepCount > 0 ? 1 : 0) + count.size());
+    if (stepCount > 0)
+    {
+        shapePy.emplace_back(stepCount);
+    }
+    std::copy(count.begin(), count.end(), std::back_inserter(shapePy));
+
+    pybind11::array_t<T> pyArray(shapePy);
+
+    // set selection if requested
+    if (!start.empty() && !count.empty())
     {
-        m_Stream->Read<T>(name, pyArray.mutable_data(),
-                          Box<Dims>(std::move(start), std::move(count)),
-                          blockID);
+        variable.SetSelection(Box<Dims>(std::move(start), std::move(count)));
     }
+
+    // set step selection if requested
+    if (stepCount > 0)
+    {
+        variable.SetStepSelection({stepStart, stepCount});
+    }
+
+    m_Stream->Read(name, pyArray.mutable_data(), blockID);
+
     return pyArray;
 }
 

testing/adios2/bindings/python/TestHighLevelAPI.py

@@ -17,7 +17,8 @@
 n_blocks = 4
 filename = 'TestHighLevelAPI.bp'
 
-# The following are the test data, with an additional dimension for timestep prepended
+# The following are the test data, with an additional dimension for timestep
+# prepended
 
 global_values = np.arange(n_times)
 global_arrays = np.arange(n_times * 2 * 16).reshape(n_times, 2, 16)
@@ -30,15 +31,17 @@ def setUpModule():
     with adios2.open(filename, 'w') as fh:
         for t in range(n_times):
             fh.write('global_value', np.array(global_values[t]))
-            fh.write('global_array', global_arrays[t], global_arrays[t].shape, (
-                0, 0), global_arrays[t].shape)
-            # We're kinda faking a local array, since the blocks all written from one proc
+            fh.write('global_array', global_arrays[t], global_arrays[t].shape,
+                     (0, 0), global_arrays[t].shape)
+            # We're kinda faking a local array, since the blocks all written
+            # from one proc
             for b in range(n_blocks):
                 fh.write('local_value', np.array(
                     local_values[t][b]), True)
             for b in range(n_blocks):
                 fh.write(
-                    'local_array', local_arrays[t][b], (), (), local_arrays[t][b].shape)
+                    'local_array', local_arrays[t][b], (), (),
+                    local_arrays[t][b].shape)
             fh.end_step()
 
 
@@ -49,8 +52,7 @@ def test_GlobalValue1d(self):
             for fh_step in fh:
                 t = fh_step.current_step()
                 val = fh_step.read('global_value')
-                self.assertTrue(np.array_equal(
-                    val, np.array([global_values[t]])))
+                self.assertTrue(val == global_values[t])
 
     def test_GlobalArray(self):
         with adios2.open(filename, 'r') as fh:
@@ -126,7 +128,6 @@ def test_LocalArrayDefault(self):
 
 
 class TestReadStepSelection(unittest.TestCase):
-    @unittest.expectedFailure
     def test_GlobalValue(self):
         with adios2.open(filename, 'r') as fh:
             val = fh.read("global_value", (), (), 1, 2)
@@ -142,7 +143,6 @@ def test_LocalValue(self):
             val = fh.read("local_value", (1,), (3,), 1, 2)
             self.assertTrue(np.array_equal(val, local_values[1:3, 1:4]))
 
-    @unittest.expectedFailure
     def test_LocalValueDefault(self):
         with adios2.open(filename, 'r') as fh:
             val = fh.read("local_value", (), (), 1, 2)