Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Tom Harke
pystencils
Commits
23de3441
Commit
23de3441
authored
Nov 14, 2019
by
Michael Kuron
Browse files
Staggered field access
parent
f86c8502
Changes
7
Hide whitespace changes
Inline
Side-by-side
pystencils/datahandling/datahandling_interface.py
View file @
23de3441
...
...
@@ -36,7 +36,7 @@ class DataHandling(ABC):
@
abstractmethod
def
add_array
(
self
,
name
:
str
,
values_per_cell
,
dtype
=
np
.
float64
,
latex_name
:
Optional
[
str
]
=
None
,
ghost_layers
:
Optional
[
int
]
=
None
,
layout
:
Optional
[
str
]
=
None
,
cpu
:
bool
=
True
,
gpu
:
Optional
[
bool
]
=
None
,
alignment
=
False
)
->
Field
:
cpu
:
bool
=
True
,
gpu
:
Optional
[
bool
]
=
None
,
alignment
=
False
,
staggered
=
False
)
->
Field
:
"""Adds a (possibly distributed) array to the handling that can be accessed using the given name.
For each array a symbolic field is available via the 'fields' dictionary
...
...
pystencils/datahandling/parallel_datahandling.py
View file @
23de3441
...
...
@@ -90,7 +90,7 @@ class ParallelDataHandling(DataHandling):
self
.
_custom_data_names
.
append
(
name
)
def
add_array
(
self
,
name
,
values_per_cell
=
1
,
dtype
=
np
.
float64
,
latex_name
=
None
,
ghost_layers
=
None
,
layout
=
None
,
cpu
=
True
,
gpu
=
None
,
alignment
=
False
):
layout
=
None
,
cpu
=
True
,
gpu
=
None
,
alignment
=
False
,
staggered
=
False
):
if
ghost_layers
is
None
:
ghost_layers
=
self
.
default_ghost_layers
if
gpu
is
None
:
...
...
@@ -140,7 +140,8 @@ class ParallelDataHandling(DataHandling):
assert
all
(
f
.
name
!=
name
for
f
in
self
.
fields
.
values
()),
"Symbolic field with this name already exists"
self
.
fields
[
name
]
=
Field
.
create_generic
(
name
,
self
.
dim
,
dtype
,
index_dimensions
,
layout
,
index_shape
=
(
values_per_cell
,)
if
index_dimensions
>
0
else
None
)
index_shape
=
(
values_per_cell
,)
if
index_dimensions
>
0
else
None
,
staggered
=
staggered
)
self
.
fields
[
name
].
latex_name
=
latex_name
self
.
_field_name_to_cpu_data_name
[
name
]
=
name
if
gpu
:
...
...
pystencils/datahandling/serial_datahandling.py
View file @
23de3441
...
...
@@ -76,7 +76,7 @@ class SerialDataHandling(DataHandling):
return
self
.
_field_information
[
name
][
'values_per_cell'
]
def
add_array
(
self
,
name
,
values_per_cell
=
1
,
dtype
=
np
.
float64
,
latex_name
=
None
,
ghost_layers
=
None
,
layout
=
None
,
cpu
=
True
,
gpu
=
None
,
alignment
=
False
):
cpu
=
True
,
gpu
=
None
,
alignment
=
False
,
staggered
=
False
):
if
ghost_layers
is
None
:
ghost_layers
=
self
.
default_ghost_layers
if
layout
is
None
:
...
...
@@ -128,7 +128,8 @@ class SerialDataHandling(DataHandling):
self
.
gpu_arrays
[
name
]
=
gpuarray
.
to_gpu
(
cpu_arr
)
assert
all
(
f
.
name
!=
name
for
f
in
self
.
fields
.
values
()),
"Symbolic field with this name already exists"
self
.
fields
[
name
]
=
Field
.
create_from_numpy_array
(
name
,
cpu_arr
,
index_dimensions
=
index_dimensions
)
self
.
fields
[
name
]
=
Field
.
create_from_numpy_array
(
name
,
cpu_arr
,
index_dimensions
=
index_dimensions
,
staggered
=
staggered
)
self
.
fields
[
name
].
latex_name
=
latex_name
return
self
.
fields
[
name
]
...
...
pystencils/field.py
View file @
23de3441
...
...
@@ -21,7 +21,7 @@ from pystencils.sympyextensions import is_integer_sequence
__all__
=
[
'Field'
,
'fields'
,
'FieldType'
,
'AbstractField'
]
def
fields
(
description
=
None
,
index_dimensions
=
0
,
layout
=
None
,
**
kwargs
):
def
fields
(
description
=
None
,
index_dimensions
=
0
,
layout
=
None
,
staggered
=
False
,
**
kwargs
):
"""Creates pystencils fields from a string description.
Examples:
...
...
@@ -61,23 +61,25 @@ def fields(description=None, index_dimensions=0, layout=None, **kwargs):
arr
=
kwargs
[
field_name
]
idx_shape_of_arr
=
()
if
not
len
(
idx_shape
)
else
arr
.
shape
[
-
len
(
idx_shape
):]
assert
idx_shape_of_arr
==
idx_shape
f
=
Field
.
create_from_numpy_array
(
field_name
,
kwargs
[
field_name
],
index_dimensions
=
len
(
idx_shape
))
f
=
Field
.
create_from_numpy_array
(
field_name
,
kwargs
[
field_name
],
index_dimensions
=
len
(
idx_shape
),
staggered
=
staggered
)
elif
isinstance
(
shape
,
tuple
):
f
=
Field
.
create_fixed_size
(
field_name
,
shape
+
idx_shape
,
dtype
=
dtype
,
index_dimensions
=
len
(
idx_shape
),
layout
=
layout
)
index_dimensions
=
len
(
idx_shape
),
layout
=
layout
,
staggered
=
staggered
)
elif
isinstance
(
shape
,
int
):
f
=
Field
.
create_generic
(
field_name
,
spatial_dimensions
=
shape
,
dtype
=
dtype
,
index_shape
=
idx_shape
,
layout
=
layout
)
index_shape
=
idx_shape
,
layout
=
layout
,
staggered
=
staggered
)
elif
shape
is
None
:
f
=
Field
.
create_generic
(
field_name
,
spatial_dimensions
=
2
,
dtype
=
dtype
,
index_shape
=
idx_shape
,
layout
=
layout
)
index_shape
=
idx_shape
,
layout
=
layout
,
staggered
=
staggered
)
else
:
assert
False
result
.
append
(
f
)
else
:
assert
layout
is
None
,
"Layout can not be specified when creating Field from numpy array"
for
field_name
,
arr
in
kwargs
.
items
():
result
.
append
(
Field
.
create_from_numpy_array
(
field_name
,
arr
,
index_dimensions
=
index_dimensions
))
result
.
append
(
Field
.
create_from_numpy_array
(
field_name
,
arr
,
index_dimensions
=
index_dimensions
,
staggered
=
staggered
))
if
len
(
result
)
==
0
:
return
None
...
...
@@ -158,6 +160,14 @@ class Field(AbstractField):
First specify the spatial offsets in [], then in case index_dimension>0 the indices in ()
e.g. ``f[-1,0,0](7)``
Staggered Fields:
Staggered fields are used to store a value on a second grid shifted by half a cell with respect to the usual
grid.
The first index dimension is used to specify the position on the staggered grid (e.g. 0 means half-way to the
eastern neighbor, 1 is half-way to the northern neighbor, etc.), while additional indices can be used to store
multiple values at each position.
Example using no index dimensions:
>>> a = np.zeros([10, 10])
>>> f = Field.create_from_numpy_array("f", a, index_dimensions=0)
...
...
@@ -172,7 +182,7 @@ class Field(AbstractField):
@
staticmethod
def
create_generic
(
field_name
,
spatial_dimensions
,
dtype
=
np
.
float64
,
index_dimensions
=
0
,
layout
=
'numpy'
,
index_shape
=
None
,
field_type
=
FieldType
.
GENERIC
)
->
'Field'
:
index_shape
=
None
,
field_type
=
FieldType
.
GENERIC
,
staggered
=
False
)
->
'Field'
:
"""
Creates a generic field where the field size is not fixed i.e. can be called with arrays of different sizes
...
...
@@ -189,6 +199,7 @@ class Field(AbstractField):
field_type: besides the normal GENERIC fields, there are INDEXED fields that store indices of the domain
that should be iterated over, and BUFFER fields that are used to generate
communication packing/unpacking kernels
staggered: enables staggered access (with half-integer offsets) and corresponding printing
"""
if
index_shape
is
not
None
:
assert
index_dimensions
==
0
or
index_dimensions
==
len
(
index_shape
)
...
...
@@ -210,11 +221,14 @@ class Field(AbstractField):
raise
ValueError
(
"Structured arrays/fields are not allowed to have an index dimension"
)
shape
+=
(
1
,)
strides
+=
(
1
,)
if
staggered
and
index_dimensions
==
0
:
raise
ValueError
(
"A staggered field needs at least one index dimension"
)
return
Field
(
field_name
,
field_type
,
dtype
,
layout
,
shape
,
strides
)
return
Field
(
field_name
,
field_type
,
dtype
,
layout
,
shape
,
strides
,
staggered
)
@
staticmethod
def
create_from_numpy_array
(
field_name
:
str
,
array
:
np
.
ndarray
,
index_dimensions
:
int
=
0
)
->
'Field'
:
def
create_from_numpy_array
(
field_name
:
str
,
array
:
np
.
ndarray
,
index_dimensions
:
int
=
0
,
staggered
=
False
)
->
'Field'
:
"""Creates a field based on the layout, data type, and shape of a given numpy array.
Kernels created for these kind of fields can only be called with arrays of the same layout, shape and type.
...
...
@@ -223,6 +237,7 @@ class Field(AbstractField):
field_name: symbolic name for the field
array: numpy array
index_dimensions: see documentation of Field
staggered: enables staggered access (with half-integer offsets) and corresponding printing
"""
spatial_dimensions
=
len
(
array
.
shape
)
-
index_dimensions
if
spatial_dimensions
<
1
:
...
...
@@ -241,12 +256,15 @@ class Field(AbstractField):
raise
ValueError
(
"Structured arrays/fields are not allowed to have an index dimension"
)
shape
+=
(
1
,)
strides
+=
(
1
,)
if
staggered
and
index_dimensions
==
0
:
raise
ValueError
(
"A staggered field needs at least one index dimension"
)
return
Field
(
field_name
,
FieldType
.
GENERIC
,
array
.
dtype
,
spatial_layout
,
shape
,
strides
)
return
Field
(
field_name
,
FieldType
.
GENERIC
,
array
.
dtype
,
spatial_layout
,
shape
,
strides
,
staggered
)
@
staticmethod
def
create_fixed_size
(
field_name
:
str
,
shape
:
Tuple
[
int
,
...],
index_dimensions
:
int
=
0
,
dtype
=
np
.
float64
,
layout
:
str
=
'numpy'
,
strides
:
Optional
[
Sequence
[
int
]]
=
None
)
->
'Field'
:
dtype
=
np
.
float64
,
layout
:
str
=
'numpy'
,
strides
:
Optional
[
Sequence
[
int
]]
=
None
,
staggered
=
False
)
->
'Field'
:
"""
Creates a field with fixed sizes i.e. can be called only with arrays of the same size and layout
...
...
@@ -257,6 +275,7 @@ class Field(AbstractField):
dtype: numpy data type of the array the kernel is called with later
layout: full layout of array, not only spatial dimensions
strides: strides in bytes or None to automatically compute them from shape (assuming no padding)
staggered: enables staggered access (with half-integer offsets) and corresponding printing
"""
spatial_dimensions
=
len
(
shape
)
-
index_dimensions
assert
spatial_dimensions
>=
1
...
...
@@ -277,13 +296,15 @@ class Field(AbstractField):
raise
ValueError
(
"Structured arrays/fields are not allowed to have an index dimension"
)
shape
+=
(
1
,)
strides
+=
(
1
,)
if
staggered
and
index_dimensions
==
0
:
raise
ValueError
(
"A staggered field needs at least one index dimension"
)
spatial_layout
=
list
(
layout
)
for
i
in
range
(
spatial_dimensions
,
len
(
layout
)):
spatial_layout
.
remove
(
i
)
return
Field
(
field_name
,
FieldType
.
GENERIC
,
dtype
,
tuple
(
spatial_layout
),
shape
,
strides
)
return
Field
(
field_name
,
FieldType
.
GENERIC
,
dtype
,
tuple
(
spatial_layout
),
shape
,
strides
,
staggered
)
def
__init__
(
self
,
field_name
,
field_type
,
dtype
,
layout
,
shape
,
strides
):
def
__init__
(
self
,
field_name
,
field_type
,
dtype
,
layout
,
shape
,
strides
,
staggered
):
"""Do not use directly. Use static create* methods"""
self
.
_field_name
=
field_name
assert
isinstance
(
field_type
,
FieldType
)
...
...
@@ -298,6 +319,7 @@ class Field(AbstractField):
0
for
_
in
range
(
self
.
spatial_dimensions
)
))
# type: tuple[float,sp.Symbol]
self
.
coordinate_transform
=
sp
.
eye
(
self
.
spatial_dimensions
)
self
.
is_staggered
=
staggered
def
new_field_with_different_name
(
self
,
new_name
):
if
self
.
has_fixed_shape
:
...
...
@@ -427,6 +449,51 @@ class Field(AbstractField):
address_mode
,
allow_textures
=
allow_textures
).
at
(
offset
)
def
staggered_access
(
self
,
offset
,
index
=
None
):
"""If this field is a staggered field, it can be accessed using half-integer offsets.
For example, an offset of ``(0, sp.Rational(1,2))`` or ``"E"`` corresponds to the staggered point to the east
of the cell center, i.e. half-way to the eastern-next cell.
If the field stores more than one value per staggered point (e.g. a vector or a tensor), the index (integer or
tuple of integers) refers to which of these values to access.
"""
assert
self
.
is_staggered
if
type
(
offset
)
is
np
.
ndarray
:
offset
=
tuple
(
offset
)
if
type
(
offset
)
is
str
:
offset
=
tuple
(
direction_string_to_offset
(
offset
,
self
.
spatial_dimensions
))
offset
=
tuple
([
o
*
sp
.
Rational
(
1
,
2
)
for
o
in
offset
])
if
type
(
offset
)
is
not
tuple
:
offset
=
(
offset
,)
if
len
(
offset
)
!=
self
.
spatial_dimensions
:
raise
ValueError
(
"Wrong number of spatial indices: "
"Got %d, expected %d"
%
(
len
(
offset
),
self
.
spatial_dimensions
))
offset
=
list
(
offset
)
for
i
,
o
in
enumerate
(
offset
):
if
(
o
+
sp
.
Rational
(
1
,
2
)).
is_Integer
:
offset
[
i
]
+=
sp
.
Rational
(
1
,
2
)
idx
=
i
offset
=
tuple
(
offset
)
if
self
.
index_dimensions
==
1
:
# this field stores a scalar value at each staggered position
if
index
is
not
None
:
raise
ValueError
(
"Cannot specify an index for a scalar staggered field"
)
return
Field
.
Access
(
self
,
offset
,
(
idx
,))
else
:
# this field stores a vector or tensor at each staggered position
if
index
is
None
:
raise
ValueError
(
"Wrong number of indices: "
"Got %d, expected %d"
%
(
0
,
self
.
index_dimensions
-
1
))
if
type
(
index
)
is
np
.
ndarray
:
index
=
tuple
(
index
)
if
type
(
index
)
is
not
tuple
:
index
=
(
index
,)
if
self
.
index_dimensions
!=
len
(
index
)
+
1
:
raise
ValueError
(
"Wrong number of indices: "
"Got %d, expected %d"
%
(
len
(
index
),
self
.
index_dimensions
-
1
))
return
Field
.
Access
(
self
,
offset
,
(
idx
,
*
index
))
def
__call__
(
self
,
*
args
,
**
kwargs
):
center
=
tuple
([
0
]
*
self
.
spatial_dimensions
)
return
Field
.
Access
(
self
,
center
)(
*
args
,
**
kwargs
)
...
...
@@ -544,6 +611,7 @@ class Field(AbstractField):
obj
.
_indirect_addressing_fields
.
update
(
a
.
field
for
a
in
e
.
atoms
(
Field
.
Access
))
obj
.
_is_absolute_access
=
is_absolute_access
obj
.
is_staggered
=
field
.
is_staggered
return
obj
def
__getnewargs__
(
self
):
...
...
@@ -676,25 +744,45 @@ class Field(AbstractField):
def
_latex
(
self
,
_
):
n
=
self
.
_field
.
latex_name
if
self
.
_field
.
latex_name
else
self
.
_field
.
name
offset_str
=
","
.
join
([
sp
.
latex
(
o
)
for
o
in
self
.
offsets
])
if
self
.
is_staggered
:
offset_str
=
","
.
join
([
sp
.
latex
(
o
-
sp
.
Rational
(
int
(
i
==
self
.
index
[
0
]),
2
))
for
i
,
o
in
enumerate
(
self
.
offsets
)])
if
self
.
is_absolute_access
:
offset_str
=
"
\\
mathbf{}"
.
format
(
offset_str
)
elif
self
.
field
.
spatial_dimensions
>
1
:
offset_str
=
"({})"
.
format
(
offset_str
)
if
self
.
index
and
self
.
index
!=
(
0
,):
return
"{{%s}_{%s}^{%s}}"
%
(
n
,
offset_str
,
self
.
index
if
len
(
self
.
index
)
>
1
else
self
.
index
[
0
])
if
self
.
is_staggered
:
if
self
.
index
and
self
.
field
.
index_dimensions
>
1
:
return
"{{%s}_{%s}^{%s}}"
%
(
n
,
offset_str
,
self
.
index
[
1
:]
if
len
(
self
.
index
)
>
2
else
self
.
index
[
1
])
else
:
return
"{{%s}_{%s}}"
%
(
n
,
offset_str
)
else
:
return
"{{%s}_{%s}}"
%
(
n
,
offset_str
)
if
self
.
index
and
self
.
field
.
index_dimensions
>
0
:
return
"{{%s}_{%s}^{%s}}"
%
(
n
,
offset_str
,
self
.
index
if
len
(
self
.
index
)
>
1
else
self
.
index
[
0
])
else
:
return
"{{%s}_{%s}}"
%
(
n
,
offset_str
)
def
__str__
(
self
):
n
=
self
.
_field
.
latex_name
if
self
.
_field
.
latex_name
else
self
.
_field
.
name
offset_str
=
","
.
join
([
sp
.
latex
(
o
)
for
o
in
self
.
offsets
])
if
self
.
is_staggered
:
offset_str
=
","
.
join
([
sp
.
latex
(
o
-
sp
.
Rational
(
int
(
i
==
self
.
index
[
0
]),
2
))
for
i
,
o
in
enumerate
(
self
.
offsets
)])
if
self
.
is_absolute_access
:
offset_str
=
"[abs]{}"
.
format
(
offset_str
)
if
self
.
index
and
self
.
index
!=
(
0
,):
return
"%s[%s](%s)"
%
(
n
,
offset_str
,
self
.
index
if
len
(
self
.
index
)
>
1
else
self
.
index
[
0
])
if
self
.
is_staggered
:
if
self
.
index
and
self
.
field
.
index_dimensions
>
1
:
return
"%s[%s](%s)"
%
(
n
,
offset_str
,
self
.
index
[
1
:]
if
len
(
self
.
index
)
>
2
else
self
.
index
[
1
])
else
:
return
"%s[%s]"
%
(
n
,
offset_str
)
else
:
return
"%s[%s]"
%
(
n
,
offset_str
)
if
self
.
index
and
self
.
field
.
index_dimensions
>
0
:
return
"%s[%s](%s)"
%
(
n
,
offset_str
,
self
.
index
if
len
(
self
.
index
)
>
1
else
self
.
index
[
0
])
else
:
return
"%s[%s]"
%
(
n
,
offset_str
)
def
get_layout_from_strides
(
strides
:
Sequence
[
int
],
index_dimension_ids
:
Optional
[
List
[
int
]]
=
None
):
...
...
pystencils_tests/test_blocking_staggered.py
View file @
23de3441
...
...
@@ -4,7 +4,8 @@ import pystencils as ps
def
test_blocking_staggered
():
f
,
stag
=
ps
.
fields
(
"f, stag(3): double[3D]"
)
f
=
ps
.
fields
(
"f: double[3D]"
)
stag
=
ps
.
fields
(
"stag(3): double[3D]"
,
staggered
=
True
)
terms
=
[
f
[
0
,
0
,
0
]
-
f
[
-
1
,
0
,
0
],
f
[
0
,
0
,
0
]
-
f
[
0
,
-
1
,
0
],
...
...
pystencils_tests/test_field.py
View file @
23de3441
...
...
@@ -127,3 +127,17 @@ def test_itemsize():
assert
x
.
itemsize
==
4
assert
y
.
itemsize
==
8
assert
i
.
itemsize
==
2
def
test_staggered
():
j1
,
j2
,
j3
=
ps
.
fields
(
'j1(2), j2(2,2), j3(2,2,2) : double[2D]'
,
staggered
=
True
)
assert
j1
[
0
,
1
](
1
)
==
j1
.
staggered_access
((
0
,
sp
.
Rational
(
1
,
2
)))
assert
j1
[
0
,
1
](
1
)
==
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
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
))
pystencils_tests/test_loop_cutting.py
View file @
23de3441
...
...
@@ -34,9 +34,9 @@ def test_staggered_iteration():
s_arr_ref
=
s_arr
.
copy
()
fields_fixed
=
(
Field
.
create_from_numpy_array
(
'f'
,
f_arr
),
Field
.
create_from_numpy_array
(
's'
,
s_arr
,
index_dimensions
=
1
))
Field
.
create_from_numpy_array
(
's'
,
s_arr
,
index_dimensions
=
1
,
staggered
=
True
))
fields_var
=
(
Field
.
create_generic
(
'f'
,
2
),
Field
.
create_generic
(
's'
,
2
,
index_dimensions
=
1
))
Field
.
create_generic
(
's'
,
2
,
index_dimensions
=
1
,
staggered
=
True
))
for
f
,
s
in
[
fields_var
,
fields_fixed
]:
# --- Manual
...
...
@@ -70,7 +70,7 @@ def test_staggered_iteration_manual():
s_arr_ref
=
s_arr
.
copy
()
f
=
Field
.
create_from_numpy_array
(
'f'
,
f_arr
)
s
=
Field
.
create_from_numpy_array
(
's'
,
s_arr
,
index_dimensions
=
1
)
s
=
Field
.
create_from_numpy_array
(
's'
,
s_arr
,
index_dimensions
=
1
,
staggered
=
True
)
eqs
=
[]
...
...
@@ -107,7 +107,8 @@ def test_staggered_iteration_manual():
def
test_staggered_gpu
():
dim
=
2
f
,
s
=
ps
.
fields
(
"f, s({dim}): double[{dim}D]"
.
format
(
dim
=
dim
))
f
=
ps
.
fields
(
"f: double[{dim}D]"
.
format
(
dim
=
dim
))
s
=
ps
.
fields
(
"s({dim}): double[{dim}D]"
.
format
(
dim
=
dim
),
staggered
=
True
)
expressions
=
[(
f
[
0
,
0
]
+
f
[
-
1
,
0
])
/
2
,
(
f
[
0
,
0
]
+
f
[
0
,
-
1
])
/
2
]
kernel_ast
=
ps
.
create_staggered_kernel
(
s
,
expressions
,
target
=
'gpu'
,
gpu_exclusive_conditions
=
True
)
...
...
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment