field.py 24.1 KB
Newer Older
1
from enum import Enum
2
3
4
5
6
from itertools import chain
import numpy as np
import sympy as sp
from sympy.core.cache import cacheit
from sympy.tensor import IndexedBase
7
from pystencils.data_types import TypedSymbol, createType, createCompositeTypeFromString, StructType
8
from pystencils.sympyextensions import isIntegerSequence
9
10


11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class FieldType(Enum):
    # generic fields
    GENERIC = 0
    # 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
    INDEXED = 1
    # communication buffer, used for (un)packing data in communication.
    BUFFER = 2

    @staticmethod
    def isGeneric(field):
        assert isinstance(field, Field)
        return field.fieldType == FieldType.GENERIC

    @staticmethod
    def isIndexed(field):
        assert isinstance(field, Field)
        return field.fieldType == FieldType.INDEXED

    @staticmethod
    def isBuffer(field):
        assert isinstance(field, Field)
        return field.fieldType == FieldType.BUFFER


Michael Kuron's avatar
Michael Kuron committed
36
class Field(object):
37
38
39
40
    """
    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
41
42
43
44
    Creating Fields:

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

45
        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
46
           case if just-in-time compilation directly from Python is done. (see :func:`Field.createFromNumpyArray`)
47
        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
48
           beforehand for a library. (see :func:`Field.createGeneric`)
49
50
51
52
53
54

    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
55
        are four dimensions, two spatial and two index dimensions: ``len(arr.shape) == spatialDims + indexDims``
56
57
58
59

    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
60
        e.g. ``f[-1,0,0](7)``
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

    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)
    """
77
78

    @staticmethod
79
    def createGeneric(fieldName, spatialDimensions, dtype=np.float64, indexDimensions=0, layout='numpy',
80
                      indexShape=None, fieldType=FieldType.GENERIC):
81
82
83
84
85
86
87
88
89
90
        """
        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)
91
92
        :param indexShape: optional shape of the index dimensions i.e. maximum values allowed for each index dimension,
                           has to be a list or tuple
93
        """
94
        if isinstance(layout, str):
95
            layout = spatialLayoutStringToTuple(layout, dim=spatialDimensions)
96
97
98
        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
99
100
101
102
103
        if indexShape is None or len(indexShape) == 0:
            shape = tuple([shapeSymbol[i] for i in range(totalDimensions)])
        else:
            shape = tuple([shapeSymbol[i] for i in range(spatialDimensions)] + list(indexShape))

104
        strides = tuple([strideSymbol[i] for i in range(totalDimensions)])
105
106
107
108
109
110
111
112

        npDataType = np.dtype(dtype)
        if npDataType.fields is not None:
            if indexDimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

113
        return Field(fieldName, fieldType, dtype, layout, shape, strides)
114

115
116
117
118
119
120
121
122
123
124
125
126
127
    @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")

128
        fullLayout = getLayoutOfArray(npArray)
129
130
131
132
        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])
133
        shape = tuple(int(s) for s in npArray.shape)
134

135
136
137
138
139
140
141
        npDataType = np.dtype(npArray.dtype)
        if npDataType.fields is not None:
            if indexDimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

142
        return Field(fieldName, FieldType.GENERIC, npArray.dtype, spatialLayout, shape, strides)
143
144

    @staticmethod
145
    def createFixedSize(fieldName, shape, indexDimensions=0, dtype=np.float64, layout='numpy', strides=None):
146
        """
147
        Creates a field with fixed sizes i.e. can be called only with arrays of the same size and layout
148

149
        :param fieldName: symbolic name for the field
150
151
        :param shape: overall shape of the array
        :param indexDimensions: how many of the trailing dimensions are interpreted as index (as opposed to spatial)
152
        :param dtype: numpy data type of the array the kernel is called with later
153
        :param layout: full layout of array, not only spatial dimensions
154
        :param strides: strides in bytes or None to automatically compute them from shape (assuming no padding)
155
        """
156
157
158
        spatialDimensions = len(shape) - indexDimensions
        assert spatialDimensions >= 1

159
160
        if isinstance(layout, str):
            layout = layoutStringToTuple(layout, spatialDimensions + indexDimensions)
161
162

        shape = tuple(int(s) for s in shape)
163
164
165
166
167
        if strides is None:
            strides = computeStrides(shape, layout)
        else:
            assert len(strides) == len(shape)
            strides = tuple([s // np.dtype(dtype).itemsize for s in strides])
168
169
170
171
172
173
174
175

        npDataType = np.dtype(dtype)
        if npDataType.fields is not None:
            if indexDimensions != 0:
                raise ValueError("Structured arrays/fields are not allowed to have an index dimension")
            shape += (1,)
            strides += (1,)

176
177
178
        spatialLayout = list(layout)
        for i in range(spatialDimensions, len(layout)):
            spatialLayout.remove(i)
179
        return Field(fieldName, FieldType.GENERIC, dtype, tuple(spatialLayout), shape, strides)
180

181
    def __init__(self, fieldName, fieldType, dtype, layout, shape, strides):
182
183
        """Do not use directly. Use static create* methods"""
        self._fieldName = fieldName
184
185
        assert isinstance(fieldType, FieldType)
        self.fieldType = fieldType
186
        self._dtype = createType(dtype)
187
        self._layout = normalizeLayout(layout)
188
189
        self.shape = shape
        self.strides = strides
Martin Bauer's avatar
Martin Bauer committed
190
        self.latexName = None
191

192
    def newFieldWithDifferentName(self, newName):
193
        return Field(newName, self.fieldType, self._dtype, self._layout, self.shape, self.strides)
194

195
196
197
198
199
200
    @property
    def spatialDimensions(self):
        return len(self._layout)

    @property
    def indexDimensions(self):
201
        return len(self.shape) - len(self._layout)
202
203
204
205
206
207
208
209
210
211
212

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

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

    @property
    def spatialShape(self):
213
        return self.shape[:self.spatialDimensions]
214

215
216
217
218
    @property
    def indexShape(self):
        return self.shape[self.spatialDimensions:]

219
220
    @property
    def hasFixedShape(self):
221
        return isIntegerSequence(self.shape)
222

223
224
    @property
    def indexShape(self):
225
        return self.shape[self.spatialDimensions:]
226

227
228
229
230
    @property
    def hasFixedIndexShape(self):
        return isIntegerSequence(self.indexShape)

231
232
    @property
    def spatialStrides(self):
233
        return self.strides[:self.spatialDimensions]
234
235
236

    @property
    def indexStrides(self):
237
        return self.strides[self.spatialDimensions:]
238
239
240
241
242
243
244
245

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

    def __repr__(self):
        return self._fieldName

246
247
248
249
250
    def neighbor(self, coordId, offset):
        offsetList = [0] * self.spatialDimensions
        offsetList[coordId] = offset
        return Field.Access(self, tuple(offsetList))

251
    def neighbors(self, stencil):
252
        return [self.__getitem__(s) for s in stencil]
253

254
255
256
257
258
259
260
261
262
263
264
265
266
    @property
    def vecCenter(self):
        indexShape = self.indexShape
        if len(indexShape) == 0:
            return self.center
        elif len(indexShape) == 1:
            return sp.Matrix([self(i) for i in range(indexShape[0])])
        elif len(indexShape) == 2:
            def cb(*args):
                r = self.__call__(*args)
                return r
            return sp.Matrix(*indexShape, cb)

267
    @property
268
269
270
271
    def center(self):
        center = tuple([0] * self.spatialDimensions)
        return Field.Access(self, center)

272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
    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):
289
        return hash((self._layout, self.shape, self.strides, self._dtype, self.fieldType, self._fieldName))
290
291

    def __eq__(self, other):
292
293
        selfTuple = (self.shape, self.strides, self.name, self.dtype, self.fieldType)
        otherTuple = (other.shape, other.strides, other.name, other.dtype, other.fieldType)
294
295
        return selfTuple == otherTuple

296

297
298
299
    PREFIX = "f"
    STRIDE_PREFIX = PREFIX + "stride_"
    SHAPE_PREFIX = PREFIX + "shape_"
Martin Bauer's avatar
Martin Bauer committed
300
301
    STRIDE_DTYPE = createCompositeTypeFromString("const int *")
    SHAPE_DTYPE = createCompositeTypeFromString("const int *")
302
    DATA_PREFIX = PREFIX + "d_"
303
304
305
306
307
308
309
310
311

    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
312
            constantOffsets = not any([isinstance(o, sp.Basic) and not o.is_Integer for o in offsetsAndIndex])
313
314
315
316
317
318
319

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

            if constantOffsets:
                offsetName = offsetToDirectionString(offsets)
                if field.indexDimensions == 0:
320
                    superscript = None
321
                elif field.indexDimensions == 1:
322
                    superscript = str(idx[0])
323
324
                else:
                    idxStr = ",".join([str(e) for e in idx])
325
                    superscript = idxStr
326
327
328
329
                if field.hasFixedIndexShape and not isinstance(field.dtype, StructType):
                    for i, bound in zip(idx, field.indexShape):
                        if i >= bound:
                            raise ValueError("Field index out of bounds")
330
331
            else:
                offsetName = "%0.10X" % (abs(hash(tuple(offsetsAndIndex))))
332
                superscript = None
333

334
335
336
337
338
            symbolName = "%s_%s" % (fieldName, offsetName)
            if superscript is not None:
                symbolName += "^" + superscript

            obj = super(Field.Access, self).__xnew__(self, symbolName)
339
340
341
342
343
344
345
346
            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
347
            obj._superscript = superscript
348
349
350
351
            obj._index = idx

            return obj

352
        def __getnewargs__(self):
353
            return self.field, self.offsets, self.index
354

355
356
357
358
359
360
361
362
363
364
365
366
367
        __xnew__ = staticmethod(__new_stage2__)
        __xnew_cached_ = staticmethod(cacheit(__new_stage2__))

        def __call__(self, *idx):
            if 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
368
369
370
        def __getitem__(self, *idx):
            return self.__call__(*idx)

Martin Bauer's avatar
Martin Bauer committed
371
372
373
374
375
        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")

376
377
378
379
380
381
382
383
        @property
        def field(self):
            return self._field

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

384
385
386
387
        @offsets.setter
        def offsets(self, value):
            self._offsets = value

388
389
390
391
392
393
394
395
396
397
398
399
        @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

400
        def _latex(self, arg):
Martin Bauer's avatar
Martin Bauer committed
401
            n = self._field.latexName if self._field.latexName else self._field.name
402
            if self._superscript:
Martin Bauer's avatar
Martin Bauer committed
403
                return "{{%s}_{%s}^{%s}}" % (n, self._offsetName, self._superscript)
404
            else:
Martin Bauer's avatar
Martin Bauer committed
405
                return "{{%s}_{%s}}" % (n, self._offsetName)
406

407
408
409
410
        @property
        def index(self):
            return self._index

411
412
413
        def getNeighbor(self, *offsets):
            return Field.Access(self.field, offsets, self.index)

414
415
416
417
418
        def neighbor(self, coordId, offset):
            offsetList = list(self.offsets)
            offsetList[coordId] += offset
            return Field.Access(self.field, tuple(offsetList), self.index)

419
420
421
        def getShifted(self, *shift):
            return Field.Access(self.field, tuple(a + b for a, b in zip(shift, self.offsets)), self.index)

422
423
        def _hashable_content(self):
            superClassContents = list(super(Field.Access, self)._hashable_content())
Martin Bauer's avatar
Martin Bauer committed
424
            t = tuple(superClassContents + [hash(self._field), self._index] + self._offsets)
425
            return t
Martin Bauer's avatar
Martin Bauer committed
426
427
428
429
430
431
432
433
434


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
435
436
437
438
439
    # 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
440
441
    equations = replacementEqs + newEq
    topologicallySortedPairs = sp.cse_main.reps_toposort([[e.lhs, e.rhs] for e in equations])
Martin Bauer's avatar
Martin Bauer committed
442
    equations = [sp.Eq(a[0], a[1].subs(symbolsEqualToTrue)) for a in topologicallySortedPairs]
Martin Bauer's avatar
Martin Bauer committed
443
444
445
    return equations


446
447
448
449
450
451
452
def getLayoutFromStrides(strides, indexDimensionIds=[]):
    coordinates = list(range(len(strides)))
    relevantStrides = [stride for i, stride in enumerate(strides) if i not in indexDimensionIds]
    result = [x for (y, x) in sorted(zip(relevantStrides, coordinates), key=lambda pair: pair[0], reverse=True)]
    return normalizeLayout(result)


453
def getLayoutOfArray(arr, indexDimensionIds=[]):
Martin Bauer's avatar
Martin Bauer committed
454
455
456
    """
    Returns a list indicating the memory layout (linearization order) of the numpy array.
    Example:
457
    >>> getLayoutOfArray(np.zeros([3,3,3]))
Martin Bauer's avatar
Martin Bauer committed
458
    (0, 1, 2)
Martin Bauer's avatar
Martin Bauer committed
459
460
461
462
463

    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
464
465

    The indexDimensionIds parameter leaves specifies which coordinates should not be
Martin Bauer's avatar
Martin Bauer committed
466
    """
467
    return getLayoutFromStrides(arr.strides, indexDimensionIds)
468
469


470
def createNumpyArrayWithLayout(shape, layout, **kwargs):
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
    """
    Creates a numpy array with
    :param shape: shape of the resulting array
    :param layout: layout as tuple, where the coordinates are ordered from slow to fast
    >>> res = createNumpyArrayWithLayout(shape=(2, 3, 4, 5), layout=(3, 2, 0, 1))
    >>> res.shape
    (2, 3, 4, 5)
    >>> getLayoutOfArray(res)
    (3, 2, 0, 1)
    """
    assert set(layout) == set(range(len(shape))), "Wrong layout descriptor"
    currentLayout = list(range(len(shape)))
    swaps = []
    for i in range(len(layout)):
        if currentLayout[i] != layout[i]:
            indexToSwapWith = currentLayout.index(layout[i])
            swaps.append((i, indexToSwapWith))
            currentLayout[i], currentLayout[indexToSwapWith] = currentLayout[indexToSwapWith], currentLayout[i]
    assert tuple(currentLayout) == tuple(layout)

    shape = list(shape)
    for a, b in swaps:
        shape[a], shape[b] = shape[b], shape[a]

495
    res = np.empty(shape, order='c', **kwargs)
496
497
498
499
500
    for a, b in reversed(swaps):
        res = res.swapaxes(a, b)
    return res


501
502
503
504
def spatialLayoutStringToTuple(layoutStr, dim):
    if layoutStr in ('fzyx', 'zyxf'):
        assert dim <= 3
        return tuple(reversed(range(dim)))
505

Martin Bauer's avatar
Martin Bauer committed
506
    if layoutStr in ('fzyx', 'f', 'reverseNumpy', 'SoA'):
507
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
508
    elif layoutStr in ('c', 'numpy', 'AoS'):
509
        return tuple(range(dim))
510
511
512
513
    raise ValueError("Unknown layout descriptor " + layoutStr)


def layoutStringToTuple(layoutStr, dim):
514
515
    layoutStr = layoutStr.lower()
    if layoutStr == 'fzyx' or layoutStr == 'soa':
516
517
        assert dim <= 4
        return tuple(reversed(range(dim)))
518
    elif layoutStr == 'zyxf' or layoutStr == 'aos':
519
520
        assert dim <= 4
        return tuple(reversed(range(dim - 1))) + (dim-1,)
Martin Bauer's avatar
Martin Bauer committed
521
    elif layoutStr == 'f' or layoutStr == 'reversenumpy':
522
523
524
        return tuple(reversed(range(dim)))
    elif layoutStr == 'c' or layoutStr == 'numpy':
        return tuple(range(dim))
525
526
527
    raise ValueError("Unknown layout descriptor " + layoutStr)


528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
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
    """
    N = len(shape)
    assert len(layout) == N
    assert len(set(layout)) == N
    strides = [0] * N
    product = 1
546
    for j in reversed(layout):
547
548
549
        strides[j] = product
        product *= shape[j]
    return tuple(strides)
Martin Bauer's avatar
Martin Bauer committed
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642


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]
643
644
645


if __name__ == '__main__':
646
647
    f = Field.createGeneric('f', spatialDimensions=2, indexShape=(2,4))
    f(2, 0)
648
649
650
    fa = f[0, 1](4) ** 2
    print(fa)
    print(sp.latex(fa))