field.py 17.8 KB
Newer Older
1
2
3
4
5
from itertools import chain
import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
from sympy.tensor import IndexedBase
6
from pystencils.types import TypedSymbol, createType
7
8


Michael Kuron's avatar
Michael Kuron committed
9
class Field(object):
10
11
12
13
    """
    With fields one can formulate stencil-like update rules on structured grids.
    This Field class knows about the dimension, memory layout (strides) and optionally about the size of an array.

Martin Bauer's avatar
Martin Bauer committed
14
15
16
17
    Creating Fields:

        To create a field use one of the static create* members. There are two options:

18
        1. create a kernel with fixed loop sizes i.e. the shape of the array is already known. This is usually the
Martin Bauer's avatar
Martin Bauer committed
19
           case if just-in-time compilation directly from Python is done. (see :func:`Field.createFromNumpyArray`)
20
        2. create a more general kernel that works for variable array sizes. This can be used to create kernels
Martin Bauer's avatar
Martin Bauer committed
21
           beforehand for a library. (see :func:`Field.createGeneric`)
22
23
24
25
26
27

    Dimensions:
        A field has spatial and index dimensions, where the spatial dimensions come first.
        The interpretation is that the field has multiple cells in (usually) two or three dimensional space which are
        looped over. Additionally  N values are stored per cell. In this case spatialDimensions is two or three,
        and indexDimensions equals N. If you want to store a matrix on each point in a two dimensional grid, there
Martin Bauer's avatar
Martin Bauer committed
28
        are four dimensions, two spatial and two index dimensions: ``len(arr.shape) == spatialDims + indexDims``
29
30
31
32

    Indexing:
        When accessing (indexing) a field the result is a FieldAccess which is derived from sympy Symbol.
        First specify the spatial offsets in [], then in case indexDimension>0 the indices in ()
Martin Bauer's avatar
Martin Bauer committed
33
        e.g. ``f[-1,0,0](7)``
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

    Example without index dimensions:
        >>> a = np.zeros([10, 10])
        >>> f = Field.createFromNumpyArray("f", a, indexDimensions=0)
        >>> jacobi = ( f[-1,0] + f[1,0] + f[0,-1] + f[0,1] ) / 4

    Example with index dimensions: LBM D2Q9 stream pull
        >>> stencil = np.array([[0,0], [0,1], [0,-1]])
        >>> src = Field.createGeneric("src", spatialDimensions=2, indexDimensions=1)
        >>> dst = Field.createGeneric("dst", spatialDimensions=2, indexDimensions=1)
        >>> for i, offset in enumerate(stencil):
        ...     sp.Eq(dst[0,0](i), src[-offset](i))
        Eq(dst_C^0, src_C^0)
        Eq(dst_C^1, src_S^1)
        Eq(dst_C^2, src_N^2)
    """
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

    @staticmethod
    def createGeneric(fieldName, spatialDimensions, dtype=np.float64, indexDimensions=0, layout='numpy'):
        """
        Creates a generic field where the field size is not fixed i.e. can be called with arrays of different sizes

        :param fieldName: symbolic name for the field
        :param dtype: numpy data type of the array the kernel is called with later
        :param spatialDimensions: see documentation of Field
        :param indexDimensions: see documentation of Field
        :param layout: tuple specifying the loop ordering of the spatial dimensions e.g. (2, 1, 0 ) means that
                       the outer loop loops over dimension 2, the second outer over dimension 1, and the inner loop
                       over dimension 0. Also allowed: the strings 'numpy' (0,1,..d) or 'reverseNumpy' (d, ..., 1, 0)
        """
        if isinstance(layout, str) and (layout == 'numpy' or layout.lower() == 'c'):
            layout = tuple(range(spatialDimensions))
        elif isinstance(layout, str) and (layout == 'reverseNumpy' or layout.lower() == 'f'):
            layout = tuple(reversed(range(spatialDimensions)))
        if len(layout) != spatialDimensions:
            raise ValueError("Layout")
        shapeSymbol = IndexedBase(TypedSymbol(Field.SHAPE_PREFIX + fieldName, Field.SHAPE_DTYPE), shape=(1,))
        strideSymbol = IndexedBase(TypedSymbol(Field.STRIDE_PREFIX + fieldName, Field.STRIDE_DTYPE), shape=(1,))
        totalDimensions = spatialDimensions + indexDimensions
        shape = tuple([shapeSymbol[i] for i in range(totalDimensions)])
        strides = tuple([strideSymbol[i] for i in range(totalDimensions)])
        return Field(fieldName, dtype, layout, shape, strides)

77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
    @staticmethod
    def createFromNumpyArray(fieldName, npArray, indexDimensions=0):
        """
        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.
        :param fieldName: symbolic name for the field
        :param npArray: numpy array
        :param indexDimensions: see documentation of Field
        """
        spatialDimensions = len(npArray.shape) - indexDimensions
        if spatialDimensions < 1:
            raise ValueError("Too many index dimensions. At least one spatial dimension required")

        fullLayout = getLayoutFromNumpyArray(npArray)
        spatialLayout = tuple([i for i in fullLayout if i < spatialDimensions])
        assert len(spatialLayout) == spatialDimensions

        strides = tuple([s // np.dtype(npArray.dtype).itemsize for s in npArray.strides])
95
        shape = tuple(int(s) for s in npArray.shape)
96
97
98
99

        return Field(fieldName, npArray.dtype, spatialLayout, shape, strides)

    @staticmethod
100
    def createFixedSize(fieldName, shape, indexDimensions=0, dtype=np.float64, layout='numpy'):
101
        """
102
        Creates a field with fixed sizes i.e. can be called only wity arrays of the same size and layout
103

104
        :param fieldName: symbolic name for the field
105
106
        :param shape: overall shape of the array
        :param indexDimensions: how many of the trailing dimensions are interpreted as index (as opposed to spatial)
107
        :param dtype: numpy data type of the array the kernel is called with later
108
        :param layout: see createGeneric
109
        """
110
111
112
        spatialDimensions = len(shape) - indexDimensions
        assert spatialDimensions >= 1

113
        if isinstance(layout, str) and (layout == 'numpy' or layout.lower() == 'c'):
114
            layout = tuple(range(spatialDimensions))
115
        elif isinstance(layout, str) and (layout == 'reverseNumpy' or layout.lower() == 'f'):
Martin Bauer's avatar
Martin Bauer committed
116
            layout = tuple(reversed(range(spatialDimensions)))
117
118
119
120

        shape = tuple(int(s) for s in shape)
        strides = computeStrides(shape, layout)
        return Field(fieldName, dtype, layout[:spatialDimensions], shape, strides)
121
122
123
124

    def __init__(self, fieldName, dtype, layout, shape, strides):
        """Do not use directly. Use static create* methods"""
        self._fieldName = fieldName
125
        self._dtype = createType(dtype)
126
        self._layout = normalizeLayout(layout)
127
128
        self.shape = shape
        self.strides = strides
129
130
131
        # index fields are currently only used for boundary handling
        # the coordinates are not the loop counters in that case, but are read from this index field
        self.isIndexField = False
132
133
134
135
136
137
138

    @property
    def spatialDimensions(self):
        return len(self._layout)

    @property
    def indexDimensions(self):
139
        return len(self.shape) - len(self._layout)
140
141
142
143
144
145
146
147
148
149
150

    @property
    def layout(self):
        return self._layout

    @property
    def name(self):
        return self._fieldName

    @property
    def spatialShape(self):
151
        return self.shape[:self.spatialDimensions]
152

153
154
155
156
157
158
159
160
    @property
    def hasFixedShape(self):
        try:
            [int(i) for i in self.shape]
            return True
        except TypeError:
            return False

161
162
    @property
    def indexShape(self):
163
        return self.shape[self.spatialDimensions:]
164
165
166

    @property
    def spatialStrides(self):
167
        return self.strides[:self.spatialDimensions]
168
169
170

    @property
    def indexStrides(self):
171
        return self.strides[self.spatialDimensions:]
172
173
174
175
176
177
178
179

    @property
    def dtype(self):
        return self._dtype

    def __repr__(self):
        return self._fieldName

180
181
182
183
184
    def neighbor(self, coordId, offset):
        offsetList = [0] * self.spatialDimensions
        offsetList[coordId] = offset
        return Field.Access(self, tuple(offsetList))

185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
    def __getitem__(self, offset):
        if type(offset) is np.ndarray:
            offset = tuple(offset)
        if type(offset) is str:
            offset = tuple(directionStringToOffset(offset, self.spatialDimensions))
        if type(offset) is not tuple:
            offset = (offset,)
        if len(offset) != self.spatialDimensions:
            raise ValueError("Wrong number of spatial indices: "
                             "Got %d, expected %d" % (len(offset), self.spatialDimensions))
        return Field.Access(self, offset)

    def __call__(self, *args, **kwargs):
        center = tuple([0]*self.spatialDimensions)
        return Field.Access(self, center)(*args, **kwargs)

    def __hash__(self):
202
        return hash((self._layout, self.shape, self.strides, self._dtype, self._fieldName))
203
204
205
206
207
208
209
210
211
212
213

    def __eq__(self, other):
        selfTuple = (self.shape, self.strides, self.name, self.dtype)
        otherTuple = (other.shape, other.strides, other.name, other.dtype)
        return selfTuple == otherTuple

    PREFIX = "f"
    STRIDE_PREFIX = PREFIX + "stride_"
    SHAPE_PREFIX = PREFIX + "shape_"
    STRIDE_DTYPE = "const int *"
    SHAPE_DTYPE = "const int *"
214
    DATA_PREFIX = PREFIX + "d_"
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255

    class Access(sp.Symbol):
        def __new__(cls, name, *args, **kwargs):
            obj = Field.Access.__xnew_cached_(cls, name, *args, **kwargs)
            return obj

        def __new_stage2__(self, field, offsets=(0, 0, 0), idx=None):
            fieldName = field.name
            offsetsAndIndex = chain(offsets, idx) if idx is not None else offsets
            constantOffsets = not any([isinstance(o, sp.Basic) for o in offsetsAndIndex])

            if not idx:
                idx = tuple([0] * field.indexDimensions)

            if constantOffsets:
                offsetName = offsetToDirectionString(offsets)

                if field.indexDimensions == 0:
                    obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName)
                elif field.indexDimensions == 1:
                    obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName + "^" + str(idx[0]))
                else:
                    idxStr = ",".join([str(e) for e in idx])
                    obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName + "^" + idxStr)

            else:
                offsetName = "%0.10X" % (abs(hash(tuple(offsetsAndIndex))))
                obj = super(Field.Access, self).__xnew__(self, fieldName + "_" + offsetName)

            obj._field = field
            obj._offsets = []
            for o in offsets:
                if isinstance(o, sp.Basic):
                    obj._offsets.append(o)
                else:
                    obj._offsets.append(int(o))
            obj._offsetName = offsetName
            obj._index = idx

            return obj

256
        def __getnewargs__(self):
257
            return self.field, self.offsets, self.index
258

259
260
261
262
263
264
265
266
267
268
269
270
271
272
        __xnew__ = staticmethod(__new_stage2__)
        __xnew_cached_ = staticmethod(cacheit(__new_stage2__))

        def __call__(self, *idx):
            if self._index != tuple([0]*self.field.indexDimensions):
                print(self._index, tuple([0]*self.field.indexDimensions))
                raise ValueError("Indexing an already indexed Field.Access")

            idx = tuple(idx)
            if len(idx) != self.field.indexDimensions and idx != (0,):
                raise ValueError("Wrong number of indices: "
                                 "Got %d, expected %d" % (len(idx), self.field.indexDimensions))
            return Field.Access(self.field, self._offsets, idx)

Martin Bauer's avatar
Martin Bauer committed
273
274
275
        def __getitem__(self, *idx):
            return self.__call__(*idx)

Martin Bauer's avatar
Martin Bauer committed
276
277
278
279
280
        def __iter__(self):
            """This is necessary to work with parts of sympy that test if an object is iterable (e.g. simplify).
            The __getitem__ would make it iterable"""
            raise TypeError("Field access is not iterable")

281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
        @property
        def field(self):
            return self._field

        @property
        def offsets(self):
            return self._offsets

        @property
        def requiredGhostLayers(self):
            return int(np.max(np.abs(self._offsets)))

        @property
        def nrOfCoordinates(self):
            return len(self._offsets)

        @property
        def offsetName(self):
            return self._offsetName

        @property
        def index(self):
            return self._index

        def _hashable_content(self):
            superClassContents = list(super(Field.Access, self)._hashable_content())
Martin Bauer's avatar
Martin Bauer committed
307
            t = tuple(superClassContents + [hash(self._field), self._index] + self._offsets)
308
            return t
Martin Bauer's avatar
Martin Bauer committed
309
310
311
312
313
314
315
316
317


def extractCommonSubexpressions(equations):
    """
    Uses sympy to find common subexpressions in equations and returns
    them in a topologically sorted order, ready for evaluation.
    Usually called before list of equations is passed to :func:`createKernel`
    """
    replacements, newEq = sp.cse(equations)
Martin Bauer's avatar
Martin Bauer committed
318
319
320
321
322
    # Workaround for older sympy versions: here subexpressions (temporary = True) are extracted
    # which leads to problems in Piecewise functions which have to a default case indicated by True
    symbolsEqualToTrue = {r[0]: True for r in replacements if r[1] is sp.true}

    replacementEqs = [sp.Eq(*r) for r in replacements if r[1] is not sp.true]
Martin Bauer's avatar
Martin Bauer committed
323
324
    equations = replacementEqs + newEq
    topologicallySortedPairs = sp.cse_main.reps_toposort([[e.lhs, e.rhs] for e in equations])
Martin Bauer's avatar
Martin Bauer committed
325
    equations = [sp.Eq(a[0], a[1].subs(symbolsEqualToTrue)) for a in topologicallySortedPairs]
Martin Bauer's avatar
Martin Bauer committed
326
327
328
    return equations


Martin Bauer's avatar
Martin Bauer committed
329
def getLayoutFromNumpyArray(arr, indexDimensionIds=[]):
Martin Bauer's avatar
Martin Bauer committed
330
331
332
333
    """
    Returns a list indicating the memory layout (linearization order) of the numpy array.
    Example:
    >>> getLayoutFromNumpyArray(np.zeros([3,3,3]))
Martin Bauer's avatar
Martin Bauer committed
334
    (0, 1, 2)
Martin Bauer's avatar
Martin Bauer committed
335
336
337
338
339

    In this example the loop over the zeroth coordinate should be the outermost loop,
    followed by the first and second. Elements arr[x,y,0] and arr[x,y,1] are adjacent in memory.
    Normally constructed numpy arrays have this order, however by stride tricks or other frameworks, arrays
    with different memory layout can be created.
Martin Bauer's avatar
Martin Bauer committed
340
341

    The indexDimensionIds parameter leaves specifies which coordinates should not be
Martin Bauer's avatar
Martin Bauer committed
342
343
    """
    coordinates = list(range(len(arr.shape)))
Martin Bauer's avatar
Martin Bauer committed
344
    relevantStrides = [stride for i, stride in enumerate(arr.strides) if i not in indexDimensionIds]
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
    result = [x for (y, x) in sorted(zip(relevantStrides, coordinates), key=lambda pair: pair[0], reverse=True)]
    return normalizeLayout(result)


def normalizeLayout(layout):
    """Takes a layout tuple and subtracts the minimum from all entries"""
    minEntry = min(layout)
    return tuple(i - minEntry for i in layout)


def computeStrides(shape, layout):
    """
    Computes strides assuming no padding exists
    :param shape: shape (size) of array
    :param layout: layout specification as tuple
    :return: strides in elements, not in bytes
    """
    layout = list(reversed(layout))
    N = len(shape)
    assert len(layout) == N
    assert len(set(layout)) == N
    strides = [0] * N
    product = 1
    for i in range(N):
        j = layout.index(i)
        strides[j] = product
        product *= shape[j]
    return tuple(strides)
Martin Bauer's avatar
Martin Bauer committed
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465


def offsetComponentToDirectionString(coordinateId, value):
    """
    Translates numerical offset to string notation.
    x offsets are labeled with east 'E' and 'W',
    y offsets with north 'N' and 'S' and
    z offsets with top 'T' and bottom 'B'
    If the absolute value of the offset is bigger than 1, this number is prefixed.
    :param coordinateId: integer 0, 1 or 2 standing for x,y and z
    :param value: integer offset

    Example:
    >>> offsetComponentToDirectionString(0, 1)
    'E'
    >>> offsetComponentToDirectionString(1, 2)
    '2N'
    """
    nameComponents = (('W', 'E'),  # west, east
                      ('S', 'N'),  # south, north
                      ('B', 'T'),  # bottom, top
                      )
    if value == 0:
        result = ""
    elif value < 0:
        result = nameComponents[coordinateId][0]
    else:
        result = nameComponents[coordinateId][1]
    if abs(value) > 1:
        result = "%d%s" % (abs(value), result)
    return result


def offsetToDirectionString(offsetTuple):
    """
    Translates numerical offset to string notation.
    For details see :func:`offsetComponentToDirectionString`
    :param offsetTuple: 3-tuple with x,y,z offset

    Example:
    >>> offsetToDirectionString([1, -1, 0])
    'SE'
    >>> offsetToDirectionString(([-3, 0, -2]))
    '2B3W'
    """
    names = ["", "", ""]
    for i in range(len(offsetTuple)):
        names[i] = offsetComponentToDirectionString(i, offsetTuple[i])
    name = "".join(reversed(names))
    if name == "":
        name = "C"
    return name


def directionStringToOffset(directionStr, dim=3):
    """
    Reverse mapping of :func:`offsetToDirectionString`
    :param directionStr: string representation of offset
    :param dim: dimension of offset, i.e the length of the returned list

    >>> directionStringToOffset('NW', dim=3)
    array([-1,  1,  0])
    >>> directionStringToOffset('NW', dim=2)
    array([-1,  1])
    >>> directionStringToOffset(offsetToDirectionString([3,-2,1]))
    array([ 3, -2,  1])
    """
    offsetMap = {
        'C': np.array([0, 0, 0]),

        'W': np.array([-1, 0, 0]),
        'E': np.array([1, 0, 0]),

        'S': np.array([0, -1, 0]),
        'N': np.array([0, 1, 0]),

        'B': np.array([0, 0, -1]),
        'T': np.array([0, 0, 1]),
    }
    offset = np.array([0, 0, 0])

    while len(directionStr) > 0:
        factor = 1
        firstNonDigit = 0
        while directionStr[firstNonDigit].isdigit():
            firstNonDigit += 1
        if firstNonDigit > 0:
            factor = int(directionStr[:firstNonDigit])
            directionStr = directionStr[firstNonDigit:]
        curOffset = offsetMap[directionStr[0]]
        offset += factor * curOffset
        directionStr = directionStr[1:]
    return offset[:dim]