nifty_utilities.py 6.16 KB
Newer Older
Ultima's avatar
Ultima committed
1
2
3
# -*- coding: utf-8 -*-

import numpy as np
4
5
6
7
8
9
from keepers import about
from itertools import product


def get_slice_list(shape, axes):
    """
theos's avatar
theos committed
10
11
    Helper function which generates slice list(s) to traverse over all
    combinations of axes, other than the selected axes.
Jait Dixit's avatar
Jait Dixit committed
12
13
14
15

    Parameters
    ----------
    shape: tuple
theos's avatar
theos committed
16
        Shape of the data array to traverse over.
Jait Dixit's avatar
Jait Dixit committed
17
    axes: tuple
theos's avatar
theos committed
18
        Axes which should not be iterated over.
Jait Dixit's avatar
Jait Dixit committed
19
20
21
22
23
24
25
26
27
28
29
30

    Yields
    -------
    list
        The next list of indices and/or slice objects for each dimension.

    Raises
    ------
    ValueError
        If shape is empty.
    ValueError
        If axes(axis) does not match shape.
31
    """
theos's avatar
theos committed
32

33
34
35
36
37
    if not shape:
        raise ValueError(about._errors.cstring("ERROR: shape cannot be None."))

    if not all(axis < len(shape) for axis in axes):
        raise ValueError(
Jait Dixit's avatar
Jait Dixit committed
38
            about._errors.cstring("ERROR: axes(axis) does not match shape.")
39
40
41
42
43
44
45
46
47
48
49
50
51
        )

    axes_select = [0 if x in axes else 1 for x, y in enumerate(shape)]

    axes_iterables = [range(y) for x, y in enumerate(shape) if x not in axes]

    for index in product(*axes_iterables):
        it_iter = iter(index)
        slice_list = [
            next(it_iter)
            if axis else slice(None, None) for axis in axes_select
        ]
        yield slice_list
Ultima's avatar
Ultima committed
52

Ultima's avatar
Ultima committed
53
54
55

def hermitianize_gaussian(x):
    # make the point inversions
Ultima's avatar
Ultima committed
56
57
    flipped_x = _hermitianize_inverter(x)
    flipped_x = flipped_x.conjugate()
Ultima's avatar
Ultima committed
58
    # check if x was already hermitian
Ultima's avatar
Ultima committed
59
60
    if (x == flipped_x).all():
        return x
Ultima's avatar
Ultima committed
61
62
    # average x and flipped_x.
    # Correct the variance by multiplying sqrt(0.5)
Ultima's avatar
Ultima committed
63
    x = (x + flipped_x) * np.sqrt(0.5)
Ultima's avatar
Ultima committed
64
65
66
    # The fixed points of the point inversion must not be avaraged.
    # Hence one must multiply them again with sqrt(0.5)
    # -> Get the middle index of the array
Ultima's avatar
Ultima committed
67
68
    mid_index = np.array(x.shape, dtype=np.int)//2
    dimensions = mid_index.size
Ultima's avatar
Ultima committed
69
70
    # Use ndindex to iterate over all combinations of zeros and the
    # mid_index in order to correct all fixed points.
Ultima's avatar
Ultima committed
71
72
73
74
75
76
77
    for i in np.ndindex((2,)*dimensions):
        temp_index = tuple(i*mid_index)
        x[temp_index] *= np.sqrt(0.5)
    try:
        x.hermitian = True
    except(AttributeError):
        pass
78

Ultima's avatar
Ultima committed
79
80
81
    return x


Ultima's avatar
Ultima committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def hermitianize(x):
    # make the point inversions
    flipped_x = _hermitianize_inverter(x)
    flipped_x = flipped_x.conjugate()
    # check if x was already hermitian
    if (x == flipped_x).all():
        return x
    # average x and flipped_x.
    # Correct the variance by multiplying sqrt(0.5)
    x = (x + flipped_x) / 2.
    try:
        x.hermitian = True
    except(AttributeError):
        pass

    return x

Ultima's avatar
Ultima committed
99
100

def _hermitianize_inverter(x):
Ultima's avatar
Ultima committed
101
    # calculate the number of dimensions the input array has
Ultima's avatar
Ultima committed
102
    dimensions = len(x.shape)
Ultima's avatar
Ultima committed
103
104
105
    # prepare the slicing object which will be used for mirroring
    slice_primitive = [slice(None), ]*dimensions
    # copy the input data
Ultima's avatar
Ultima committed
106
    y = x.copy()
Ultima's avatar
Ultima committed
107
    # flip in every direction
Ultima's avatar
Ultima committed
108
109
    for i in xrange(dimensions):
        slice_picker = slice_primitive[:]
Ultima's avatar
Ultima committed
110
        slice_picker[i] = slice(1, None, None)
ultimanet's avatar
ultimanet committed
111
112
        slice_picker = tuple(slice_picker)

Ultima's avatar
Ultima committed
113
114
        slice_inverter = slice_primitive[:]
        slice_inverter[i] = slice(None, 0, -1)
ultimanet's avatar
ultimanet committed
115
        slice_inverter = tuple(slice_inverter)
Ultima's avatar
Ultima committed
116
117

        try:
118
119
            y.set_data(to_key=slice_picker, data=y,
                       from_key=slice_inverter)
Ultima's avatar
Ultima committed
120
121
122
        except(AttributeError):
            y[slice_picker] = y[slice_inverter]
    return y
123
124


125
def direct_vdot(x, y):
Ultima's avatar
Ultima committed
126
    # the input could be fields. Try to extract the data
Ultima's avatar
Ultima committed
127
128
129
130
131
132
133
134
    try:
        x = x.get_val()
    except(AttributeError):
        pass
    try:
        y = y.get_val()
    except(AttributeError):
        pass
Ultima's avatar
Ultima committed
135
    # try to make a direct vdot
Ultima's avatar
Ultima committed
136
137
138
139
    try:
        return x.vdot(y)
    except(AttributeError):
        pass
140

Ultima's avatar
Ultima committed
141
142
143
    try:
        return y.vdot(x)
    except(AttributeError):
144
145
        pass

Ultima's avatar
Ultima committed
146
    # fallback to numpy
147
148
    return np.vdot(x, y)

Ultima's avatar
Ultima committed
149
150

def convert_nested_list_to_object_array(x):
Ultima's avatar
Ultima committed
151
152
153
    # if x is a nested_list full of ndarrays all having the same size,
    # np.shape returns the shape of the ndarrays, too, i.e. too many
    # dimensions
Ultima's avatar
Ultima committed
154
    possible_shape = np.shape(x)
Ultima's avatar
Ultima committed
155
    # Check if possible_shape goes too deep.
156
    dimension_counter = 0
Ultima's avatar
Ultima committed
157
158
    current_extract = x
    for i in xrange(len(possible_shape)):
159
160
        if not isinstance(current_extract, list) and \
                not isinstance(current_extract, tuple):
Ultima's avatar
Ultima committed
161
162
163
164
            break
        current_extract = current_extract[0]
        dimension_counter += 1
    real_shape = possible_shape[:dimension_counter]
Ultima's avatar
Ultima committed
165
    # if the numpy array was not encapsulated at all, return x directly
Ultima's avatar
Ultima committed
166
167
    if real_shape == ():
        return x
Ultima's avatar
Ultima committed
168
169
    # Prepare the carrier-object
    carrier = np.empty(real_shape, dtype=np.object)
Ultima's avatar
Ultima committed
170
171
    for i in xrange(np.prod(real_shape)):
        ii = np.unravel_index(i, real_shape)
172
        try:
Ultima's avatar
Ultima committed
173
174
175
            carrier[ii] = x[ii]
        except(TypeError):
            extracted = x
176
            for j in xrange(len(ii)):
Ultima's avatar
Ultima committed
177
                extracted = extracted[ii[j]]
178
            carrier[ii] = extracted
Ultima's avatar
Ultima committed
179
180
181
182
183
184
185
186
187
188
189
190
191
192
    return carrier


def field_map(ishape, function, *args):
    if ishape == ():
        return function(*args)
    else:
        if args == ():
            result = np.empty(ishape, dtype=np.object)
            for i in xrange(np.prod(ishape)):
                ii = np.unravel_index(i, ishape)
                result[ii] = function()
            return result
        else:
Ultima's avatar
Ultima committed
193
194
195
196
            # define a helper function in order to clip the get-indices
            # to be suitable for the foreign arrays in args.
            # This allows you to do operations, like adding to fields
            # with ishape (3,4,3) and (3,4,1)
Ultima's avatar
Ultima committed
197
198
199
200
201
202
203
            def get_clipped(w, ind):
                w_shape = np.array(np.shape(w))
                get_tuple = tuple(np.clip(ind, 0, w_shape-1))
                return w[get_tuple]
            result = np.empty_like(args[0])
            for i in xrange(np.prod(result.shape)):
                ii = np.unravel_index(i, result.shape)
204
205
206
207
208
                result[ii] = function(
                    *map(
                        lambda z: get_clipped(z, ii), args
                    )
                )
Ultima's avatar
Ultima committed
209
210
                # result[ii] = function(*map(lambda z: z[ii], args))
            return result