From 1b4ace3e2bd99d12e831afc0217c6c2072c39cb3 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 13 Jan 2020 10:49:39 +0100
Subject: [PATCH] Add DataHandling.add_arrays
This adds a function for lazy people who want to write
```python
dh = create_data_handling((20, 30))
dh.add_arrays('x, y(9), z')
```
instead of
```python
dh = create_data_handling((20, 30))
dh.add_array('x')
dh.add_array('y', values_per_cell=9)
dh.add_array('z')
```
because in most cases you want to use more than one array.
---
.../datahandling/datahandling_interface.py | 20 +++++++++++++++++++
pystencils/field.py | 20 ++++++++++---------
pystencils_tests/test_datahandling.py | 14 ++++++++++++-
3 files changed, 44 insertions(+), 10 deletions(-)
diff --git a/pystencils/datahandling/datahandling_interface.py b/pystencils/datahandling/datahandling_interface.py
index af1a6ba1f..2cc71b2b1 100644
--- a/pystencils/datahandling/datahandling_interface.py
+++ b/pystencils/datahandling/datahandling_interface.py
@@ -62,6 +62,26 @@ class DataHandling(ABC):
pystencils field, that can be used to formulate symbolic kernels
"""
+ def add_arrays(self, description: str):
+ """Adds multiple arrays using a string description similar to :func:`pystencils.fields`
+
+
+ >>> from pystencils.datahandling import create_data_handling
+ >>> dh = create_data_handling((20, 30))
+ >>> dh.add_arrays('x, y(9)')
+ >>> print(dh.fields)
+ {'x': x: double[20,30], 'y': y(9): double[20,30]}
+ >>> assert dh.fields['x'].shape = (20, 30)
+ >>> assert dh.fields['y'].index_shape = (9,)
+
+ Args:
+ description (str): String description of the fields to add
+ """
+ from pystencils.field import _parse_part1
+
+ for name, indices in _parse_part1(description):
+ self.add_array(name, values_per_cell=indices)
+
@abstractmethod
def has_data(self, name):
"""Returns true if a field or custom data element with this name was added."""
diff --git a/pystencils/field.py b/pystencils/field.py
index 69adfcb8a..0bf80b211 100644
--- a/pystencils/field.py
+++ b/pystencils/field.py
@@ -1058,15 +1058,17 @@ type_description_regex = re.compile(r"""
""", re.VERBOSE | re.IGNORECASE)
-def _parse_description(description):
- def parse_part1(d):
+def _parse_part1(d):
+ result = field_description_regex.match(d)
+ while result:
+ name, index_str = result.group(1), result.group(2)
+ index = tuple(int(e) for e in index_str.split(",")) if index_str else ()
+ yield name, index
+ d = d[result.end():]
result = field_description_regex.match(d)
- while result:
- name, index_str = result.group(1), result.group(2)
- index = tuple(int(e) for e in index_str.split(",")) if index_str else ()
- yield name, index
- d = d[result.end():]
- result = field_description_regex.match(d)
+
+
+def _parse_description(description):
def parse_part2(d):
result = type_description_regex.match(d)
@@ -1091,7 +1093,7 @@ def _parse_description(description):
else:
field_description, field_info = description, 'float64[2D]'
- fields_info = [e for e in parse_part1(field_description)]
+ fields_info = [e for e in _parse_part1(field_description)]
if not field_info:
raise ValueError("Could not parse field description")
diff --git a/pystencils_tests/test_datahandling.py b/pystencils_tests/test_datahandling.py
index 1af1a68e3..2f4d87a40 100644
--- a/pystencils_tests/test_datahandling.py
+++ b/pystencils_tests/test_datahandling.py
@@ -13,7 +13,6 @@ except ImportError:
pytest = unittest.mock.MagicMock()
-
def basic_iteration(dh):
dh.add_array('basic_iter_test_gl_default')
dh.add_array('basic_iter_test_gl_3', ghost_layers=3)
@@ -227,3 +226,16 @@ def test_vtk_output():
for domain_shape in [(4, 5), (3, 4, 5)]:
dh = create_data_handling(domain_size=domain_shape, periodicity=True)
vtk_output(dh)
+
+
+def test_add_arrays():
+ domain_shape = (3, 4, 5)
+ field_description = 'x, y(9)'
+
+ dh = create_data_handling(domain_size=domain_shape, default_ghost_layers=0, default_layout='numpy')
+ dh.add_arrays(field_description)
+
+ x, y = ps.fields(field_description + ': [3,4,5]')
+
+ assert x == dh.fields['x']
+ assert y == dh.fields['y']
--
GitLab