nifty_utilities.py 6.42 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):
    """
Jait Dixit's avatar
Jait Dixit committed
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
    Helper function which generate slice list(s) to traverse over all
    combinations of axes, other than the selected axes. This is used in
    conjunction with outer-product spaces. The generated slice list can
    passed to the target outer-product space to extract the selected
    data. The target outer-product space can be either a d2o object or
    a numpy.ndarray.

    Parameters
    ----------
    shape: tuple
        Shape of the target d2o object or numpy.ndarray.
    axes: tuple
        Axes along which data needs to be extracted.

    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.
35
36
37
38
39
40
    """
    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
41
            about._errors.cstring("ERROR: axes(axis) does not match shape.")
42
43
44
45
46
47
48
49
50
51
52
53
54
        )

    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
55

Ultima's avatar
Ultima committed
56
57
58

def hermitianize_gaussian(x):
    # make the point inversions
Ultima's avatar
Ultima committed
59
60
    flipped_x = _hermitianize_inverter(x)
    flipped_x = flipped_x.conjugate()
Ultima's avatar
Ultima committed
61
    # check if x was already hermitian
Ultima's avatar
Ultima committed
62
63
    if (x == flipped_x).all():
        return x
Ultima's avatar
Ultima committed
64
65
    # average x and flipped_x.
    # Correct the variance by multiplying sqrt(0.5)
Ultima's avatar
Ultima committed
66
    x = (x + flipped_x) * np.sqrt(0.5)
Ultima's avatar
Ultima committed
67
68
69
    # 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
70
71
    mid_index = np.array(x.shape, dtype=np.int)//2
    dimensions = mid_index.size
Ultima's avatar
Ultima committed
72
73
    # 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
74
75
76
77
78
79
80
    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
81

Ultima's avatar
Ultima committed
82
83
84
    return x


Ultima's avatar
Ultima committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
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
102
103

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

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

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


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

Ultima's avatar
Ultima committed
144
145
146
    try:
        return y.vdot(x)
    except(AttributeError):
147
148
        pass

Ultima's avatar
Ultima committed
149
    # fallback to numpy
150
151
    return np.vdot(x, y)

Ultima's avatar
Ultima committed
152
153

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