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

import numpy as np
4
5
6
7
8
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
36
37
38
39
from keepers import about
from itertools import product


def get_slice_list(shape, axes):
    """
    Yields a slice list which can be passed directly to d2o object or
    numpy.ndarray to extract that slice of data.

    :shape: Tuple
        Shape of the target d2o object or numpy.ndarray
    :axes: Tuple
        Axes which are to be extracted
    :returns: List
        List of indices and/or slice objects

    """
    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(
            about._errors.cstring("ERROR: axes(axis) do not match shape.")
        )

    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
40

Ultima's avatar
Ultima committed
41
42
43

def hermitianize_gaussian(x):
    # make the point inversions
Ultima's avatar
Ultima committed
44
45
    flipped_x = _hermitianize_inverter(x)
    flipped_x = flipped_x.conjugate()
Ultima's avatar
Ultima committed
46
    # check if x was already hermitian
Ultima's avatar
Ultima committed
47
48
    if (x == flipped_x).all():
        return x
Ultima's avatar
Ultima committed
49
50
    # average x and flipped_x.
    # Correct the variance by multiplying sqrt(0.5)
Ultima's avatar
Ultima committed
51
    x = (x + flipped_x) * np.sqrt(0.5)
Ultima's avatar
Ultima committed
52
53
54
    # 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
55
56
    mid_index = np.array(x.shape, dtype=np.int)//2
    dimensions = mid_index.size
Ultima's avatar
Ultima committed
57
58
    # 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
59
60
61
62
63
64
65
    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
66

Ultima's avatar
Ultima committed
67
68
69
    return x


Ultima's avatar
Ultima committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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
87
88

def _hermitianize_inverter(x):
Ultima's avatar
Ultima committed
89
    # calculate the number of dimensions the input array has
Ultima's avatar
Ultima committed
90
    dimensions = len(x.shape)
Ultima's avatar
Ultima committed
91
92
93
    # prepare the slicing object which will be used for mirroring
    slice_primitive = [slice(None), ]*dimensions
    # copy the input data
Ultima's avatar
Ultima committed
94
    y = x.copy()
Ultima's avatar
Ultima committed
95
    # flip in every direction
Ultima's avatar
Ultima committed
96
97
    for i in xrange(dimensions):
        slice_picker = slice_primitive[:]
Ultima's avatar
Ultima committed
98
        slice_picker[i] = slice(1, None, None)
ultimanet's avatar
ultimanet committed
99
100
        slice_picker = tuple(slice_picker)

Ultima's avatar
Ultima committed
101
102
        slice_inverter = slice_primitive[:]
        slice_inverter[i] = slice(None, 0, -1)
ultimanet's avatar
ultimanet committed
103
        slice_inverter = tuple(slice_inverter)
Ultima's avatar
Ultima committed
104
105

        try:
106
107
            y.set_data(to_key=slice_picker, data=y,
                       from_key=slice_inverter)
Ultima's avatar
Ultima committed
108
109
110
        except(AttributeError):
            y[slice_picker] = y[slice_inverter]
    return y
111
112


113
def direct_vdot(x, y):
Ultima's avatar
Ultima committed
114
    # the input could be fields. Try to extract the data
Ultima's avatar
Ultima committed
115
116
117
118
119
120
121
122
    try:
        x = x.get_val()
    except(AttributeError):
        pass
    try:
        y = y.get_val()
    except(AttributeError):
        pass
Ultima's avatar
Ultima committed
123
    # try to make a direct vdot
Ultima's avatar
Ultima committed
124
125
126
127
    try:
        return x.vdot(y)
    except(AttributeError):
        pass
128

Ultima's avatar
Ultima committed
129
130
131
    try:
        return y.vdot(x)
    except(AttributeError):
132
133
        pass

Ultima's avatar
Ultima committed
134
    # fallback to numpy
135
136
    return np.vdot(x, y)

Ultima's avatar
Ultima committed
137
138

def convert_nested_list_to_object_array(x):
Ultima's avatar
Ultima committed
139
140
141
    # 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
142
    possible_shape = np.shape(x)
Ultima's avatar
Ultima committed
143
    # Check if possible_shape goes too deep.
144
    dimension_counter = 0
Ultima's avatar
Ultima committed
145
146
    current_extract = x
    for i in xrange(len(possible_shape)):
147
148
        if not isinstance(current_extract, list) and \
                not isinstance(current_extract, tuple):
Ultima's avatar
Ultima committed
149
150
151
152
            break
        current_extract = current_extract[0]
        dimension_counter += 1
    real_shape = possible_shape[:dimension_counter]
Ultima's avatar
Ultima committed
153
    # if the numpy array was not encapsulated at all, return x directly
Ultima's avatar
Ultima committed
154
155
    if real_shape == ():
        return x
Ultima's avatar
Ultima committed
156
157
    # Prepare the carrier-object
    carrier = np.empty(real_shape, dtype=np.object)
Ultima's avatar
Ultima committed
158
159
    for i in xrange(np.prod(real_shape)):
        ii = np.unravel_index(i, real_shape)
160
        try:
Ultima's avatar
Ultima committed
161
162
163
            carrier[ii] = x[ii]
        except(TypeError):
            extracted = x
164
            for j in xrange(len(ii)):
Ultima's avatar
Ultima committed
165
                extracted = extracted[ii[j]]
166
            carrier[ii] = extracted
Ultima's avatar
Ultima committed
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    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
181
182
183
184
            # 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
185
186
187
188
189
190
191
            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)
192
193
194
195
196
                result[ii] = function(
                    *map(
                        lambda z: get_clipped(z, ii), args
                    )
                )
Ultima's avatar
Ultima committed
197
198
                # result[ii] = function(*map(lambda z: z[ii], args))
            return result