From 2258e7dd1883f7c1001f8e90cddabecb5c0e751b Mon Sep 17 00:00:00 2001
From: Michael Kuron <m.kuron@gmx.de>
Date: Sat, 23 Nov 2019 15:09:39 +0100
Subject: [PATCH] Field.staggered_vector_access and Field.center_vector for up
to 3 index dimensions
---
pystencils/field.py | 26 +++++++++++++++++++++-----
pystencils_tests/test_field.py | 7 +++++++
2 files changed, 28 insertions(+), 5 deletions(-)
diff --git a/pystencils/field.py b/pystencils/field.py
index fc6f93d60..eb962adbb 100644
--- a/pystencils/field.py
+++ b/pystencils/field.py
@@ -426,13 +426,15 @@ class Field(AbstractField):
index_shape = self.index_shape
if len(index_shape) == 0:
return sp.Matrix([self.center])
- if len(index_shape) == 1:
+ elif len(index_shape) == 1:
return sp.Matrix([self(i) for i in range(index_shape[0])])
elif len(index_shape) == 2:
- def cb(*args):
- r = self.__call__(*args)
- return r
- return sp.Matrix(*index_shape, cb)
+ return sp.Matrix([[self(i, j) for j in range(index_shape[1])] for i in range(index_shape[0])])
+ elif len(index_shape) == 3:
+ return sp.Matrix([[[self(i, j, k) for k in range(index_shape[2])]
+ for j in range(index_shape[1])] for i in range(index_shape[0])])
+ else:
+ raise NotImplementedError("center_vector is not implemented for more than 3 index dimensions")
@property
def center(self):
@@ -534,6 +536,20 @@ class Field(AbstractField):
return prefactor * Field.Access(self, offset, (idx, *index))
+ def staggered_vector_access(self, offset):
+ """Like staggered_access, but returns the entire vector/tensor stored at offset."""
+ assert FieldType.is_staggered(self)
+
+ if self.index_dimensions == 1:
+ return sp.Matrix([self.staggered_access(offset)])
+ elif self.index_dimensions == 2:
+ return sp.Matrix([self.staggered_access(offset, i) for i in range(self.index_shape[1])])
+ elif self.index_dimensions == 3:
+ return sp.Matrix([[self.staggered_access(offset, (i, k)) for k in range(self.index_shape[2])]
+ for i in range(self.index_shape[1])])
+ else:
+ raise NotImplementedError("center_vector is not implemented for more than 3 index dimensions")
+
@property
def staggered_stencil(self):
assert FieldType.is_staggered(self)
diff --git a/pystencils_tests/test_field.py b/pystencils_tests/test_field.py
index b1d9af430..ad23682a7 100644
--- a/pystencils_tests/test_field.py
+++ b/pystencils_tests/test_field.py
@@ -28,6 +28,9 @@ def test_field_basic():
assert neighbor.offsets == (-1, 1)
assert '_' in neighbor._latex('dummy')
+ f = Field.create_fixed_size('f', (8, 8, 2, 2, 2), index_dimensions=3)
+ assert f.center_vector == sp.Matrix([[[f(i, j, k) for k in range(2)] for j in range(2)] for i in range(2)])
+
f = Field.create_generic('f', spatial_dimensions=5, index_dimensions=2)
field_access = f[1, -1, 2, -3, 0](1, 0)
assert field_access.offsets == (1, -1, 2, -3, 0)
@@ -139,12 +142,16 @@ def test_staggered():
assert j1[0, 2](1) == j1.staggered_access((0, sp.Rational(3, 2)))
assert j1[0, 1](1) == j1.staggered_access("N")
assert j1[0, 0](1) == j1.staggered_access("S")
+ assert j1.staggered_vector_access("N") == sp.Matrix([j1.staggered_access("N")])
assert j2[0, 1](1, 1) == j2.staggered_access((0, sp.Rational(1, 2)), 1)
assert j2[0, 1](1, 1) == j2.staggered_access("N", 1)
+ assert j2.staggered_vector_access("N") == sp.Matrix([j2.staggered_access("N", 0), j2.staggered_access("N", 1)])
assert j3[0, 1](1, 1, 1) == j3.staggered_access((0, sp.Rational(1, 2)), (1, 1))
assert j3[0, 1](1, 1, 1) == j3.staggered_access("N", (1, 1))
+ assert j3.staggered_vector_access("N") == sp.Matrix([[j3.staggered_access("N", (i, j))
+ for j in range(2)] for i in range(2)])
# D2Q9
k = ps.fields('k(4) : double[2D]', field_type=FieldType.STAGGERED)
--
GitLab