grid.py 8.42 KB
Newer Older
1
import numpy as np
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
2
import functools
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
3
import tfields.lib.util
4
5


6
7
8
9
10
11
12
13
14
15
16
17
def ensure_complex(*base_vectors):
    # ensure, that the third entry in base_vector of type tuple becomes a complex type
    base_vectors = list(base_vectors)
    for i, vector in enumerate(base_vectors):
        if isinstance(vector, tuple):
            if len(vector) == 3:
                vector = list(vector)
                vector[2] = complex(vector[2])
                base_vectors[i] = tuple(vector)
    return base_vectors


Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
18
def _to_base_vectors(*base_vectors):
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    """
    Transform tuples to arrays with np.mgrid
    Args:
        tuple of lenght 3 with complex third entry -> start, stop, n_steps
    Returns:
        list if np.array for each base
    """
    base_vectors = list(base_vectors)
    # transform tuples to arrays with mgrid
    for i in range(len(base_vectors)):
        if isinstance(base_vectors[i], tuple):
            base_vectors[i] = np.mgrid[slice(*base_vectors[i])]
        if isinstance(base_vectors[i], list):
            base_vectors[i] = np.array(base_vectors[i])
    return base_vectors


36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def igrid(*base_vectors, **kwargs):
    """
    Args:
        *base_vectors (tuple, list or np.array): base vectors spaning the grid
            behaviour for different input types:
                tuple: will be transformed to slices and given to np.mgrid
                list or np.array: arbitrary base vectors
        **kwargs
            iter_order (list): order in which the iteration will be done.
                Frequency rises with position in list. default is [0, 1, 2]
                iteration will be done like::
                      
                for v0 in base_vectors[iter_order[0]]:
                    for v1 in base_vectors[iter_order[1]]:
                        for v2 in base_vectors[iter_order[2]]:
                            ...

    Examples:
        Initilaize using the mgrid notation
        >>> import tfields
56
        >>> import numpy as np
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
        >>> assert np.array_equal(tfields.igrid((0, 1, 2j),
        ...                                     (3, 4, 2j),
        ...                                     (6, 7, 2j)),
        ...                       [[ 0.,  3.,  6.],
        ...                        [ 0.,  3.,  7.],
        ...                        [ 0.,  4.,  6.],
        ...                        [ 0.,  4.,  7.],
        ...                        [ 1.,  3.,  6.],
        ...                        [ 1.,  3.,  7.],
        ...                        [ 1.,  4.,  6.],
        ...                        [ 1.,  4.,  7.]])

        >>> assert np.array_equal(tfields.igrid([3, 4], np.linspace(0, 1, 2),
        ...                                     (6, 7, 2),
        ...                                     iter_order=[1, 0, 2]),
        ...                       [[ 3.,  0.,  6.],
        ...                        [ 3.,  0.,  7.],
        ...                        [ 4.,  0.,  6.],
        ...                        [ 4.,  0.,  7.],
        ...                        [ 3.,  1.,  6.],
        ...                        [ 3.,  1.,  7.],
        ...                        [ 4.,  1.,  6.],
        ...                        [ 4.,  1.,  7.]])
        >>> assert np.array_equal(tfields.igrid(np.linspace(0, 1, 2),
        ...                                     np.linspace(3, 4, 2),
        ...                                     np.linspace(6, 7, 2),
        ...                                     iter_order=[2, 0, 1]),
        ...                       [[ 0.,  3.,  6.],
        ...                        [ 0.,  4.,  6.],
        ...                        [ 1.,  3.,  6.],
        ...                        [ 1.,  4.,  6.],
        ...                        [ 0.,  3.,  7.],
        ...                        [ 0.,  4.,  7.],
        ...                        [ 1.,  3.,  7.],
        ...                        [ 1.,  4.,  7.]])
92
93
94

    """
    iter_order = kwargs.pop('iter_order', np.arange(len(base_vectors)))
95
    base_vectors = ensure_complex(*base_vectors)
96
97

    # create the grid
98
99
100
101
102
    if all([isinstance(val, tuple) for val in base_vectors]):
        base_vectors = [slice(*base_vectors[index]) for index in iter_order]
        obj = np.mgrid[base_vectors]
        obj = obj.reshape(obj.shape[0], -1).T
    else:
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
103
        base_vectors = _to_base_vectors(*base_vectors)
104

Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
105
106
        obj = np.empty(shape=(functools.reduce(lambda x, y: x * y,
                                               map(len, base_vectors)),
107
108
109
110
111
112
113
114
115
116
117
118
119
120
                              len(base_vectors)))

        def loop_rec(y, n_max, i=0, n=None, *vals):
            if n is None:
                n = n_max
            if n > 0:
                for val in y[n_max - n]:
                    i = loop_rec(y, n_max, i, n - 1, val, *vals)
            else:
                obj[i] = list(reversed(vals))
                i += 1
            return i

        loop_rec([base_vectors[i] for i in iter_order], len(base_vectors))
121

122
123
124
125
    swap_indices = compare_permutations(iter_order, np.arange(len(base_vectors)))
    swap_columns(obj, *swap_indices)
    return obj

126

127
def base_vectors(array, rtol=None, atol=None):
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
    """
    describe the array in terms of base vectors
    Inverse function of igrid

    Examples:
        >>> import tfields
        >>> grid = tfields.igrid((3, 5, 5j))
        >>> tfields.lib.grid.base_vectors(grid[:, 0])
        (3.0, 5.0, 5j)
        >>> grid2 = tfields.igrid((3, 5, 5j),
        ...                       (1, 2, 2j))
        >>> grid_circle = tfields.igrid(*tfields.lib.grid.base_vectors(grid2))
        >>> assert tfields.Tensors(grid_circle).equal(grid2)

    """
    if len(array.shape) == 1:
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
144
        values = set(array)
145
146
        if rtol is not None and atol is not None:
            duplicates = set([])
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
147
            for v1, v2 in tfields.lib.util.pairwise(sorted(values)):
148
149
150
151
152
153
154
155
                if np.isclose(v1, v2, rtol=rtol, atol=atol):
                    duplicates.add(v2)
            values = values.difference(duplicates)
            # round to given absolute precision
            n_digits = int(abs(np.log10(atol))) + 1
            values = {round(v, n_digits) for v in values}
        elif rtol is not None or atol is not None:
            raise ValueError("rtol and atol arguments only come in pairs.")
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
156
157
158
159
160
161
162
        spacing = complex(0, len(values))
        vmin = min(values)
        vmax = max(values)
        return (vmin, vmax, spacing)
    elif len(array.shape) == 2:
        bases = []
        for i in range(array.shape[1]):
163
            bases.append(base_vectors(array[:, i], rtol=rtol, atol=atol))
Daniel Boeckenhoff's avatar
Daniel Boeckenhoff committed
164
165
166
167
168
        return bases
    else:
        raise NotImplementedError("Description yet only till rank 1")


169
170
171
172
173
174
175
def swap_columns(array, *index_tuples):
    """
    Args:
        array (list or array)
        expects tuples with indices to swap as arguments
    Examples:
        >>> import tfields
176
177
        >>> l = np.array([[3, 2, 1, 0], [6, 5, 4, 3]])
        >>> tfields.lib.grid.swap_columns(l, (1, 2), (0, 3))
178
        >>> l
179
180
        array([[0, 1, 2, 3],
               [3, 4, 5, 6]])
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201

    """
    # test swap_indices type
    for si in index_tuples:
        if hasattr(si, '__iter__'):
            if len(si) == 2:
                continue
        raise TypeError("swap_indices must be tuple but is {}"
                        .format(si))
    for i, j in index_tuples:
        array[:, [i, j]] = array[:, [j, i]]


def swap_rows(array, *args):
    """
    Args:
        array (list)
        expects tuples with indices to swap as arguments
    Examples:
        >>> import tfields
        >>> l = [[3,3,3], [2,2,2], [1,1,1], [0, 0, 0]]
202
        >>> tfields.lib.grid.swap_rows(l, (1, 2), (0, 3))
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
        >>> l
        [[0, 0, 0], [1, 1, 1], [2, 2, 2], [3, 3, 3]]

    """
    for i, j in args:
        array[i], array[j] = array[j], array[i]


def compare_permutations(permut1, permut2):
    """
    Return what you need to switch in order to make permut1 become permut2
    Examples:
        >>> import tfields
        >>> a = [1, 2, 0, 4, 3]
        >>> b = [0, 1, 2, 3, 4]
218
        >>> si = tfields.lib.grid.compare_permutations(a, b)
219
220
        >>> si
        [(0, 2), (1, 2), (3, 4)]
221
        >>> tfields.lib.grid.swap_rows(a, *si)
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        >>> a
        [0, 1, 2, 3, 4]
        >>> a == b
        True

    """
    swap_indices = []
    permut1 = list(permut1)
    i = 0
    while i < len(permut2):
        if not permut1[i] == permut2[i]:
            j = permut1.index(permut2[i])
            swap_rows(permut1, (i, j))
            swap_indices.append((i, j))
        i += 1
    return swap_indices


if __name__ == '__main__':
    import doctest
    doctest.testmod()