field.py 24.9 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

8
from pystencils.assignment import Assignment
9
from pystencils.alignedarray import aligned_empty
10
from pystencils.data_types import TypedSymbol, createType, createCompositeTypeFromString, StructType
11
from pystencils.sympyextensions import isIntegerSequence
12
13


14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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
39
class Field(object):
40
41
42
43
    """
    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
44
45
46
47
    Creating Fields:

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

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

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

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

    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):
75
76
77
78
        ...     Assignment(dst[0,0](i), src[-offset](i))
        Assignment(dst_C^0, src_C^0)
        Assignment(dst_C^1, src_S^1)
        Assignment(dst_C^2, src_N^2)
79
    """
80
81

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

107
        strides = tuple([strideSymbol[i] for i in range(totalDimensions)])
108
109
110
111
112
113
114
115

        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,)

116
        return Field(fieldName, fieldType, dtype, layout, shape, strides)
117

118
119
120
121
122
123
124
125
126
127
128
129
130
    @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")

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

138
139
140
141
142
143
144
        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,)

145
        return Field(fieldName, FieldType.GENERIC, npArray.dtype, spatialLayout, shape, strides)
146
147

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

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

162
163
        if isinstance(layout, str):
            layout = layoutStringToTuple(layout, spatialDimensions + indexDimensions)
164
165

        shape = tuple(int(s) for s in shape)
166
167
168
169
170
        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])
171
172
173
174
175
176
177
178

        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,)

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

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

195
    def newFieldWithDifferentName(self, newName):
196
        return Field(newName, self.fieldType, self._dtype, self._layout, self.shape, self.strides)
197

198
199
200
201
202
203
    @property
    def spatialDimensions(self):
        return len(self._layout)

    @property
    def indexDimensions(self):
204
        return len(self.shape) - len(self._layout)
205
206
207
208
209
210
211
212
213
214
215

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

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

    @property
    def spatialShape(self):
216
        return self.shape[:self.spatialDimensions]
217

218
219
220
221
    @property
    def indexShape(self):
        return self.shape[self.spatialDimensions:]

222
223
    @property
    def hasFixedShape(self):
224
        return isIntegerSequence(self.shape)
225

226
227
    @property
    def indexShape(self):
228
        return self.shape[self.spatialDimensions:]
229

230
231
232
233
    @property
    def hasFixedIndexShape(self):
        return isIntegerSequence(self.indexShape)

234
235
    @property
    def spatialStrides(self):
236
        return self.strides[:self.spatialDimensions]
237
238
239

    @property
    def indexStrides(self):
240
        return self.strides[self.spatialDimensions:]
241
242
243
244
245
246
247
248

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

    def __repr__(self):
        return self._fieldName

249
250
251
252
253
    def neighbor(self, coordId, offset):
        offsetList = [0] * self.spatialDimensions
        offsetList[coordId] = offset
        return Field.Access(self, tuple(offsetList))

254
    def neighbors(self, stencil):
255
        return [self.__getitem__(s) for s in stencil]
256

257
258
259
260
261
262
263
264
265
266
267
268
269
    @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)

270
    @property
271
272
273
274
    def center(self):
        center = tuple([0] * self.spatialDimensions)
        return Field.Access(self, center)

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

    def __eq__(self, other):
295
296
        selfTuple = (self.shape, self.strides, self.name, self.dtype, self.fieldType)
        otherTuple = (other.shape, other.strides, other.name, other.dtype, other.fieldType)
297
298
        return selfTuple == otherTuple

299

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

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

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

            if constantOffsets:
                offsetName = offsetToDirectionString(offsets)
                if field.indexDimensions == 0:
323
                    superscript = None
324
                elif field.indexDimensions == 1:
325
                    superscript = str(idx[0])
326
327
                else:
                    idxStr = ",".join([str(e) for e in idx])
328
                    superscript = idxStr
329
330
331
332
                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")
333
334
            else:
                offsetName = "%0.10X" % (abs(hash(tuple(offsetsAndIndex))))
335
                superscript = None
336

337
338
339
340
341
            symbolName = "%s_%s" % (fieldName, offsetName)
            if superscript is not None:
                symbolName += "^" + superscript

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

            return obj

355
        def __getnewargs__(self):
356
            return self.field, self.offsets, self.index
357

358
359
360
361
362
363
364
365
        __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)
366
367
368
369
370

            if self.field.indexDimensions == 0 and idx == (0,):
                idx = ()

            if len(idx) != self.field.indexDimensions:
371
372
373
374
                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
375
376
377
        def __getitem__(self, *idx):
            return self.__call__(*idx)

Martin Bauer's avatar
Martin Bauer committed
378
379
380
381
382
        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")

383
384
385
386
387
388
389
390
        @property
        def field(self):
            return self._field

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

391
392
393
394
        @offsets.setter
        def offsets(self, value):
            self._offsets = value

395
396
397
398
399
400
401
402
403
404
405
406
        @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

407
        def _latex(self, arg):
Martin Bauer's avatar
Martin Bauer committed
408
            n = self._field.latexName if self._field.latexName else self._field.name
409
            if self._superscript:
Martin Bauer's avatar
Martin Bauer committed
410
                return "{{%s}_{%s}^{%s}}" % (n, self._offsetName, self._superscript)
411
            else:
Martin Bauer's avatar
Martin Bauer committed
412
                return "{{%s}_{%s}}" % (n, self._offsetName)
413

414
415
416
417
        @property
        def index(self):
            return self._index

418
419
420
        def getNeighbor(self, *offsets):
            return Field.Access(self.field, offsets, self.index)

421
422
423
424
425
        def neighbor(self, coordId, offset):
            offsetList = list(self.offsets)
            offsetList[coordId] += offset
            return Field.Access(self.field, tuple(offsetList), self.index)

426
427
428
        def getShifted(self, *shift):
            return Field.Access(self.field, tuple(a + b for a, b in zip(shift, self.offsets)), self.index)

429
430
        def _hashable_content(self):
            superClassContents = list(super(Field.Access, self)._hashable_content())
Martin Bauer's avatar
Martin Bauer committed
431
            t = tuple(superClassContents + [hash(self._field), self._index] + self._offsets)
432
            return t
Martin Bauer's avatar
Martin Bauer committed
433
434
435
436
437
438
439
440


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`
    """
441
    replacements, new_eq = sp.cse(equations)
Martin Bauer's avatar
Martin Bauer committed
442
443
    # 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
444
    symbols_equal_to_true = {r[0]: True for r in replacements if r[1] is sp.true}
Martin Bauer's avatar
Martin Bauer committed
445

446
447
448
449
    replacement_eqs = [Assignment(*r) for r in replacements if r[1] is not sp.true]
    equations = replacement_eqs + new_eq
    topologically_sorted_pairs = sp.cse_main.reps_toposort([[e.lhs, e.rhs] for e in equations])
    equations = [Assignment(a[0], a[1].subs(symbols_equal_to_true)) for a in topologically_sorted_pairs]
Martin Bauer's avatar
Martin Bauer committed
450
451
452
    return equations


453
454
def getLayoutFromStrides(strides, indexDimensionIds=[]):
    coordinates = list(range(len(strides)))
455
456
    relevant_strides = [stride for i, stride in enumerate(strides) if i not in indexDimensionIds]
    result = [x for (y, x) in sorted(zip(relevant_strides, coordinates), key=lambda pair: pair[0], reverse=True)]
457
458
459
    return normalizeLayout(result)


460
def getLayoutOfArray(arr, indexDimensionIds=[]):
Martin Bauer's avatar
Martin Bauer committed
461
462
463
    """
    Returns a list indicating the memory layout (linearization order) of the numpy array.
    Example:
464
    >>> getLayoutOfArray(np.zeros([3,3,3]))
Martin Bauer's avatar
Martin Bauer committed
465
    (0, 1, 2)
Martin Bauer's avatar
Martin Bauer committed
466
467
468
469
470

    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
471
472

    The indexDimensionIds parameter leaves specifies which coordinates should not be
Martin Bauer's avatar
Martin Bauer committed
473
    """
474
    return getLayoutFromStrides(arr.strides, indexDimensionIds)
475
476


477
def createNumpyArrayWithLayout(shape, layout, alignment=False, byteOffset=0, **kwargs):
478
479
480
481
    """
    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
482
483
484
    :param alignment: number of bytes to align the beginning and the innermost coordinate to, or False for no alignment
    :param byteOffset: only used when alignment is specified, align not beginning but address at this offset
                       mostly used to align first inner cell, not ghost cells
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
    >>> 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]

505
506
507
508
509
    if not alignment:
        res = np.empty(shape, order='c', **kwargs)
    else:
        if alignment is True:
            alignment = 8 * 4
Martin Bauer's avatar
Martin Bauer committed
510
        res = aligned_empty(shape, alignment, byteOffset=byteOffset, **kwargs)
511

512
513
514
515
516
    for a, b in reversed(swaps):
        res = res.swapaxes(a, b)
    return res


517
518
519
520
def spatialLayoutStringToTuple(layoutStr, dim):
    if layoutStr in ('fzyx', 'zyxf'):
        assert dim <= 3
        return tuple(reversed(range(dim)))
521

Martin Bauer's avatar
Martin Bauer committed
522
    if layoutStr in ('fzyx', 'f', 'reverseNumpy', 'SoA'):
523
        return tuple(reversed(range(dim)))
Martin Bauer's avatar
Martin Bauer committed
524
    elif layoutStr in ('c', 'numpy', 'AoS'):
525
        return tuple(range(dim))
526
527
528
529
    raise ValueError("Unknown layout descriptor " + layoutStr)


def layoutStringToTuple(layoutStr, dim):
530
531
    layoutStr = layoutStr.lower()
    if layoutStr == 'fzyx' or layoutStr == 'soa':
532
533
        assert dim <= 4
        return tuple(reversed(range(dim)))
534
    elif layoutStr == 'zyxf' or layoutStr == 'aos':
535
536
        assert dim <= 4
        return tuple(reversed(range(dim - 1))) + (dim-1,)
Martin Bauer's avatar
Martin Bauer committed
537
    elif layoutStr == 'f' or layoutStr == 'reversenumpy':
538
539
540
        return tuple(reversed(range(dim)))
    elif layoutStr == 'c' or layoutStr == 'numpy':
        return tuple(range(dim))
541
542
543
    raise ValueError("Unknown layout descriptor " + layoutStr)


544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
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
562
    for j in reversed(layout):
563
564
565
        strides[j] = product
        product *= shape[j]
    return tuple(strides)
Martin Bauer's avatar
Martin Bauer committed
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
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658


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]
659
660
661


if __name__ == '__main__':
662
663
    f = Field.createGeneric('f', spatialDimensions=2, indexShape=(2,4))
    f(2, 0)
664
665
666
    fa = f[0, 1](4) ** 2
    print(fa)
    print(sp.latex(fa))