Merge pull request #4491 from vicentebolea/python-stream-numpy

python: stream.write accept np array w/o extra args

Author: vicentebolea

python/adios2/stream.py

@@ -1,6 +1,6 @@
 """License:
-  Distributed under the OSI-approved Apache License, Version 2.0.  See
-  accompanying file Copyright.txt for details.
+Distributed under the OSI-approved Apache License, Version 2.0.  See
+accompanying file Copyright.txt for details.
 """
 
 from functools import singledispatchmethod
@@ -296,14 +296,17 @@ def _(self, name, content, shape=[], start=[], count=[], operations=None):
 
         if not variable:
             # Sequence variables
-            if isinstance(content, np.ndarray):
-                variable = self._io.define_variable(name, content, shape, start, count)
-            elif isinstance(content, list):
-                if shape == [] and count == []:
-                    shape = [len(content)]
-                    count = shape
-                    start = [0]
+            if isinstance(content, (list, np.ndarray)):
+                if isinstance(content, list):
+                    content = np.asarray(content)
+
+                # If shape, start, and count is not specified, use the numpy array's shape
+                if shape == [] and start == [] and count == []:
+                    shape = list(content.shape)
+                    start = [0] * content.ndim
+                    count = shape[:]
                 variable = self._io.define_variable(name, content, shape, start, count)
+
             # Scalar variables
             elif isinstance(content, str):
                 variable = self._io.define_variable(name, content)

testing/adios2/python/TestStream.py

@@ -1,5 +1,6 @@
 from adios2 import Stream, LocalValueDim
 from random import randint
+import numpy as np
 
 import unittest
 
@@ -22,14 +23,16 @@ def test_basic(self):
                 s.write("Wind", [5], shape=[LocalValueDim])
                 # Local Array
                 s.write("Coords", [38, -46], [], [], [2])
+                s.write("humidity", np.random.rand(3, 1))
 
         with Stream("pythonstreamtest.bp", "r") as s:
             for _ in s.steps():
                 for var_name in s.available_variables():
                     print(f"var:{var_name}\t{s.read(var_name)}")
-                    self.assertEqual(s.read("Wind", block_id=0), 5)
-                    self.assertEqual(s.read("Coords", block_id=0)[0], 38)
-                    self.assertEqual(s.read("Coords", block_id=0)[1], -46)
+                self.assertEqual(s.read("Wind", block_id=0), 5)
+                self.assertEqual(s.read("Coords", block_id=0)[0], 38)
+                self.assertEqual(s.read("Coords", block_id=0)[1], -46)
+                self.assertEqual(s.read("humidity", block_id=0).ndim, 2)
 
     def test_start_count(self):
         with Stream("pythonstreamtest.bp", "w") as s: