nifty_mpi_data.py 109 KB
Newer Older
ultimanet's avatar
ultimanet committed
1
# -*- coding: utf-8 -*-
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
## NIFTY (Numerical Information Field Theory) has been developed at the
## Max-Planck-Institute for Astrophysics.
##
## Copyright (C) 2015 Max-Planck-Society
##
## Author: Theo Steininger
## Project homepage: <http://www.mpa-garching.mpg.de/ift/nifty/>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
## See the GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program. If not, see <http://www.gnu.org/licenses/>.


ultimanet's avatar
ultimanet committed
24
import numpy as np
Ultimanet's avatar
Ultimanet committed
25
from nifty_about import about
26
from weakref import WeakValueDictionary as weakdict
ultimanet's avatar
ultimanet committed
27

28
29
# initialize the 'FOUND-packages'-dictionary
FOUND = {}
ultimanet's avatar
ultimanet committed
30
try:
31
    from mpi4py import MPI
Ultima's avatar
Ultima committed
32
    FOUND['MPI'] = True
33
except(ImportError):
34
    import mpi_dummy as MPI
Ultima's avatar
Ultima committed
35
    FOUND['MPI'] = False
ultimanet's avatar
ultimanet committed
36
37
38

try:
    import pyfftw
Ultima's avatar
Ultima committed
39
    FOUND['pyfftw'] = True
40
except(ImportError):
Ultima's avatar
Ultima committed
41
    FOUND['pyfftw'] = False
ultimanet's avatar
ultimanet committed
42
43

try:
44
    import h5py
Ultima's avatar
Ultima committed
45
46
    FOUND['h5py'] = True
    FOUND['h5py_parallel'] = h5py.get_config().mpi
ultimanet's avatar
ultimanet committed
47
except(ImportError):
Ultima's avatar
Ultima committed
48
49
    FOUND['h5py'] = False
    FOUND['h5py_parallel'] = False
ultimanet's avatar
ultimanet committed
50
51


Ultima's avatar
Ultima committed
52
53
54
55
ALL_DISTRIBUTION_STRATEGIES = ['not', 'equal', 'fftw', 'freeform']
GLOBAL_DISTRIBUTION_STRATEGIES = ['not', 'equal', 'fftw']
LOCAL_DISTRIBUTION_STRATEGIES = ['freeform']
HDF5_DISTRIBUTION_STRATEGIES = ['equal', 'fftw']
56
57
58

COMM = MPI.COMM_WORLD

59

ultimanet's avatar
ultimanet committed
60
61
62
63
64
65
66
67
class distributed_data_object(object):
    """

        NIFTY class for distributed data

        Parameters
        ----------
        global_data : {tuple, list, numpy.ndarray} *at least 1-dimensional*
68
            Initial data which will be casted to a numpy.ndarray and then
ultimanet's avatar
ultimanet committed
69
70
71
72
73
74
            stored according to the distribution strategy. The global_data's
            shape overwrites global_shape.
        global_shape : tuple of ints, *optional*
            If no global_data is supplied, global_shape can be used to
            initialize an empty distributed_data_object
        dtype : type, *optional*
75
76
            If an explicit dtype is supplied, the given global_data will be
            casted to it.
ultimanet's avatar
ultimanet committed
77
        distribution_strategy : {'fftw' (default), 'not'}, *optional*
78
79
            Specifies the way, how global_data will be distributed to the
            individual nodes.
ultimanet's avatar
ultimanet committed
80
            'fftw' follows the distribution strategy of pyfftw.
81
82
            'not' does not distribute the data at all.

ultimanet's avatar
ultimanet committed
83
84
85
86
87
88
89
90
91
92

        Attributes
        ----------
        data : numpy.ndarray
            The numpy.ndarray in which the individual node's data is stored.
        dtype : type
            Data type of the data object.
        distribution_strategy : string
            Name of the used distribution_strategy
        distributor : distributor
93
94
            The distributor object which takes care of all distribution and
            consolidation of the data.
ultimanet's avatar
ultimanet committed
95
96
        shape : tuple of int
            The global shape of the data
97

ultimanet's avatar
ultimanet committed
98
99
        Raises
        ------
100
101
102
        TypeError :
            If the supplied distribution strategy is not known.

ultimanet's avatar
ultimanet committed
103
    """
104
    def __init__(self, global_data=None, global_shape=None, dtype=None,
Ultima's avatar
Ultima committed
105
                 local_data=None, local_shape=None,
106
                 distribution_strategy='fftw', hermitian=False,
107
108
109
110
111
                 alias=None, path=None, comm=MPI.COMM_WORLD,
                 copy=True, *args, **kwargs):

        # TODO: allow init with empty shape

112
113
114
115
        if isinstance(global_data, tuple) or isinstance(global_data, list):
            global_data = np.array(global_data, copy=False)
        if isinstance(local_data, tuple) or isinstance(local_data, list):
            local_data = np.array(local_data, copy=False)
116

117
118
        self.distributor = distributor_factory.get_distributor(
                                distribution_strategy = distribution_strategy,
119
                                comm = comm,
120
                                global_data = global_data,
121
                                global_shape = global_shape,
Ultima's avatar
Ultima committed
122
123
124
125
                                local_data = local_data,
                                local_shape = local_shape,
                                alias = alias,
                                path = path,
126
127
                                dtype = dtype,
                                **kwargs)
128

ultimanet's avatar
ultimanet committed
129
130
131
        self.distribution_strategy = distribution_strategy
        self.dtype = self.distributor.dtype
        self.shape = self.distributor.global_shape
Ultima's avatar
Ultima committed
132
133
        self.local_shape = self.distributor.local_shape
        self.comm = self.distributor.comm
134
135

        self.init_args = args
136
        self.init_kwargs = kwargs
137

Ultima's avatar
Ultima committed
138
139
140
141
        (self.data, self.hermitian) = self.distributor.initialize_data(
                                                     global_data = global_data,
                                                     local_data = local_data,
                                                     alias = alias,
Ultima's avatar
Ultima committed
142
                                                     path = path,
Ultima's avatar
Ultima committed
143
144
                                                     hermitian = hermitian,
                                                     copy = copy)
145
        self.index = d2o_librarian.register(self)
146

Ultimanet's avatar
Ultimanet committed
147
    def copy(self, dtype=None, distribution_strategy=None, **kwargs):
148
149
150
        temp_d2o = self.copy_empty(dtype=dtype,
                                   distribution_strategy=distribution_strategy,
                                   **kwargs)
Ultima's avatar
Ultima committed
151
        if distribution_strategy is None or \
152
                distribution_strategy == self.distribution_strategy:
Ultimanet's avatar
Ultimanet committed
153
154
            temp_d2o.set_local_data(self.get_local_data(), copy=True)
        else:
Ultima's avatar
Ultima committed
155
            temp_d2o.inject((slice(None),), self, (slice(None),))
156
        temp_d2o.hermitian = self.hermitian
157
        return temp_d2o
158
159

    def copy_empty(self, global_shape=None, local_shape=None, dtype=None,
160
                   distribution_strategy=None, **kwargs):
Ultima's avatar
Ultima committed
161
162
163
164
165
166
167
168
169
        if self.distribution_strategy == 'not' and \
            distribution_strategy in LOCAL_DISTRIBUTION_STRATEGIES and \
            local_shape == None:
            result = self.copy_empty(global_shape = global_shape,
                                     local_shape = local_shape,
                                     dtype = dtype,
                                     distribution_strategy = 'equal',
                                     **kwargs)
            return result.copy_empty(distribution_strategy = 'freeform')
170

Ultima's avatar
Ultima committed
171
        if global_shape is None:
172
            global_shape = self.shape
Ultima's avatar
Ultima committed
173
174
175
        if local_shape is None:
            local_shape = self.local_shape
        if dtype is None:
176
            dtype = self.dtype
Ultima's avatar
Ultima committed
177
        if distribution_strategy is None:
178
179
180
            distribution_strategy = self.distribution_strategy

        kwargs.update(self.init_kwargs)
181

182
        temp_d2o = distributed_data_object(global_shape=global_shape,
183
184
185
                               local_shape = local_shape,
                               dtype = dtype,
                               distribution_strategy = distribution_strategy,
Ultima's avatar
Ultima committed
186
                               comm = self.comm,
187
188
                               *self.init_args,
                               **kwargs)
189
        return temp_d2o
190

191
    def apply_scalar_function(self, function, inplace=False, dtype=None):
192
        remember_hermitianQ = self.hermitian
193
194

        if inplace == True:
Ultimanet's avatar
Ultimanet committed
195
            temp = self
Ultima's avatar
Ultima committed
196
            if dtype is not None and self.dtype != np.dtype(dtype):
197
198
199
                about.warnings.cprint(
                    "WARNING: Inplace dtype conversion is not possible!")

Ultimanet's avatar
Ultimanet committed
200
        else:
201
            temp = self.copy_empty(dtype=dtype)
Ultimanet's avatar
Ultimanet committed
202

Ultima's avatar
Ultima committed
203
        if np.prod(self.local_shape) != 0:
204
            try:
Ultima's avatar
Ultima committed
205
206
207
208
                temp.data[:] = function(self.data)
            except:
                temp.data[:] = np.vectorize(function)(self.data)
        else:
209
210
            # Noting to do here. The value-empty array
            # is also geometrically empty
Ultima's avatar
Ultima committed
211
            pass
212

213
214
215
216
        if function in (np.exp, np.log):
            temp.hermitian = remember_hermitianQ
        else:
            temp.hermitian = False
Ultimanet's avatar
Ultimanet committed
217
        return temp
218

Ultimanet's avatar
Ultimanet committed
219
220
221
    def apply_generator(self, generator):
        self.set_local_data(generator(self.distributor.local_shape))
        self.hermitian = False
222

ultimanet's avatar
ultimanet committed
223
224
    def __str__(self):
        return self.data.__str__()
225

ultimanet's avatar
ultimanet committed
226
227
    def __repr__(self):
        return '<distributed_data_object>\n'+self.data.__repr__()
228
229


230
    def _compare_helper(self, other, op):
231
        result = self.copy_empty(dtype = np.bool_)
232
233
        # Case 1: 'other' is a scalar
        # -> make point-wise comparison
Ultimanet's avatar
Ultimanet committed
234
        if np.isscalar(other):
235
236
            result.set_local_data(
                    getattr(self.get_local_data(copy = False), op)(other))
237
            return result
Ultimanet's avatar
Ultimanet committed
238

239
240
        # Case 2: 'other' is a numpy array or a distributed_data_object
        # -> extract the local data and make point-wise comparison
Ultimanet's avatar
Ultimanet committed
241
242
243
        elif isinstance(other, np.ndarray) or\
        isinstance(other, distributed_data_object):
            temp_data = self.distributor.extract_local_data(other)
244
245
            result.set_local_data(
                getattr(self.get_local_data(copy=False), op)(temp_data))
Ultimanet's avatar
Ultimanet committed
246
            return result
247
248

        # Case 3: 'other' is None
Ultima's avatar
Ultima committed
249
        elif other is None:
Ultimanet's avatar
Ultimanet committed
250
            return False
251
252
253

        # Case 4: 'other' is something different
        # -> make a numpy casting and make a recursive call
Ultimanet's avatar
Ultimanet committed
254
255
        else:
            temp_other = np.array(other)
256
            return getattr(self, op)(temp_other)
257

258
259
260

    def __ne__(self, other):
        return self._compare_helper(other, '__ne__')
261

262
263
    def __lt__(self, other):
        return self._compare_helper(other, '__lt__')
264

265
266
267
268
269
270
271
272
273
274
275
276
    def __le__(self, other):
        return self._compare_helper(other, '__le__')

    def __eq__(self, other):

        return self._compare_helper(other, '__eq__')
    def __ge__(self, other):
        return self._compare_helper(other, '__ge__')

    def __gt__(self, other):
        return self._compare_helper(other, '__gt__')

Ultimanet's avatar
Ultimanet committed
277
    def equal(self, other):
Ultimanet's avatar
Ultimanet committed
278
279
280
281
282
283
284
285
286
        if other is None:
            return False
        try:
            assert(self.dtype == other.dtype)
            assert(self.shape == other.shape)
            assert(self.init_args == other.init_args)
            assert(self.init_kwargs == other.init_kwargs)
            assert(self.distribution_strategy == other.distribution_strategy)
            assert(np.all(self.data == other.data))
Ultimanet's avatar
Ultimanet committed
287
        except(AssertionError, AttributeError):
Ultimanet's avatar
Ultimanet committed
288
289
290
291
            return False
        else:
            return True

292
293
294



295
    def __pos__(self):
296
        temp_d2o = self.copy_empty()
297
        temp_d2o.set_local_data(data = self.get_local_data(), copy = True)
298
        return temp_d2o
299

ultimanet's avatar
ultimanet committed
300
    def __neg__(self):
301
        temp_d2o = self.copy_empty()
302
        temp_d2o.set_local_data(data = self.get_local_data().__neg__(),
303
                                copy = True)
ultimanet's avatar
ultimanet committed
304
        return temp_d2o
305

306
    def __abs__(self):
307
        # translate complex dtypes
308
309
310
311
312
313
        if self.dtype == np.dtype('complex64'):
            new_dtype = np.dtype('float32')
        elif self.dtype == np.dtype('complex128'):
            new_dtype = np.dtype('float64')
        elif issubclass(self.dtype.type, np.complexfloating):
            new_dtype = np.dtype('float')
Ultimanet's avatar
Ultimanet committed
314
315
316
        else:
            new_dtype = self.dtype
        temp_d2o = self.copy_empty(dtype = new_dtype)
317
        temp_d2o.set_local_data(data = self.get_local_data().__abs__(),
318
                                copy = True)
319
        return temp_d2o
320

Ultima's avatar
Ultima committed
321
322
323
324
325
    def _builtin_helper(self, operator, other, inplace=False):
        if isinstance(other, distributed_data_object):
            other_is_real = other.isreal()
        else:
            other_is_real = np.isreal(other)
326
327

        # Case 1: other is not a scalar
Ultimanet's avatar
Ultimanet committed
328
        if not (np.isscalar(other) or np.shape(other) == (1,)):
329
            try:
330
                hermitian_Q = (other.hermitian and self.hermitian)
331
332
            except(AttributeError):
                hermitian_Q = False
333
            # extract the local data from the 'other' object
Ultimanet's avatar
Ultimanet committed
334
335
            temp_data = self.distributor.extract_local_data(other)
            temp_data = operator(temp_data)
336
337

        # Case 2: other is a real scalar -> preserve hermitianity
Ultima's avatar
Ultima committed
338
339
        elif other_is_real or (self.dtype not in (np.dtype('complex128'),
                                                  np.dtype('complex256'))):
340
            hermitian_Q = self.hermitian
ultimanet's avatar
ultimanet committed
341
            temp_data = operator(other)
342
        # Case 3: other is complex
343
344
        else:
            hermitian_Q = False
345
346
            temp_data = operator(other)
        # write the new data into a new distributed_data_object
347
348
349
        if inplace == True:
            temp_d2o = self
        else:
350
            # use common datatype for self and other
351
            new_dtype = np.dtype(np.find_common_type((self.dtype,),
352
                                                     (temp_data.dtype,)))
353
354
            temp_d2o = self.copy_empty(
                            dtype = new_dtype)
ultimanet's avatar
ultimanet committed
355
        temp_d2o.set_local_data(data=temp_data)
356
        temp_d2o.hermitian = hermitian_Q
ultimanet's avatar
ultimanet committed
357
        return temp_d2o
358

ultimanet's avatar
ultimanet committed
359
    def __add__(self, other):
Ultima's avatar
Ultima committed
360
        return self._builtin_helper(self.get_local_data().__add__, other)
ultimanet's avatar
ultimanet committed
361
362

    def __radd__(self, other):
Ultima's avatar
Ultima committed
363
        return self._builtin_helper(self.get_local_data().__radd__, other)
Ultimanet's avatar
Ultimanet committed
364
365

    def __iadd__(self, other):
366
        return self._builtin_helper(self.get_local_data().__iadd__,
367
368
                                               other,
                                               inplace = True)
Ultimanet's avatar
Ultimanet committed
369

ultimanet's avatar
ultimanet committed
370
    def __sub__(self, other):
Ultima's avatar
Ultima committed
371
        return self._builtin_helper(self.get_local_data().__sub__, other)
372

ultimanet's avatar
ultimanet committed
373
    def __rsub__(self, other):
Ultima's avatar
Ultima committed
374
        return self._builtin_helper(self.get_local_data().__rsub__, other)
375

ultimanet's avatar
ultimanet committed
376
    def __isub__(self, other):
377
        return self._builtin_helper(self.get_local_data().__isub__,
378
379
                                               other,
                                               inplace = True)
380

ultimanet's avatar
ultimanet committed
381
    def __div__(self, other):
Ultima's avatar
Ultima committed
382
        return self._builtin_helper(self.get_local_data().__div__, other)
383

384
385
    def __truediv__(self, other):
        return self.__div__(other)
386

ultimanet's avatar
ultimanet committed
387
    def __rdiv__(self, other):
Ultima's avatar
Ultima committed
388
        return self._builtin_helper(self.get_local_data().__rdiv__, other)
389

390
391
    def __rtruediv__(self, other):
        return self.__rdiv__(other)
ultimanet's avatar
ultimanet committed
392

Ultimanet's avatar
Ultimanet committed
393
    def __idiv__(self, other):
394
        return self._builtin_helper(self.get_local_data().__idiv__,
395
396
                                               other,
                                               inplace = True)
397
    def __itruediv__(self, other):
398
        return self.__idiv__(other)
399

ultimanet's avatar
ultimanet committed
400
    def __floordiv__(self, other):
401
402
        return self._builtin_helper(self.get_local_data().__floordiv__,
                                       other)
ultimanet's avatar
ultimanet committed
403
    def __rfloordiv__(self, other):
404
        return self._builtin_helper(self.get_local_data().__rfloordiv__,
Ultimanet's avatar
Ultimanet committed
405
406
                                       other)
    def __ifloordiv__(self, other):
Ultima's avatar
Ultima committed
407
        return self._builtin_helper(
408
409
                    self.get_local_data().__ifloordiv__, other,
                                               inplace = True)
410

ultimanet's avatar
ultimanet committed
411
    def __mul__(self, other):
Ultima's avatar
Ultima committed
412
        return self._builtin_helper(self.get_local_data().__mul__, other)
413

ultimanet's avatar
ultimanet committed
414
    def __rmul__(self, other):
Ultima's avatar
Ultima committed
415
        return self._builtin_helper(self.get_local_data().__rmul__, other)
ultimanet's avatar
ultimanet committed
416
417

    def __imul__(self, other):
418
        return self._builtin_helper(self.get_local_data().__imul__,
419
420
                                               other,
                                               inplace = True)
Ultimanet's avatar
Ultimanet committed
421

ultimanet's avatar
ultimanet committed
422
    def __pow__(self, other):
Ultima's avatar
Ultima committed
423
        return self._builtin_helper(self.get_local_data().__pow__, other)
424

ultimanet's avatar
ultimanet committed
425
    def __rpow__(self, other):
Ultima's avatar
Ultima committed
426
        return self._builtin_helper(self.get_local_data().__rpow__, other)
ultimanet's avatar
ultimanet committed
427
428

    def __ipow__(self, other):
429
        return self._builtin_helper(self.get_local_data().__ipow__,
430
431
                                               other,
                                               inplace = True)
Ultima's avatar
Ultima committed
432
433
434
    def __mod__(self, other):
        return self._builtin_helper(self.get_local_data().__mod__, other)
    def __rmod__(self, other):
435
        return self._builtin_helper(self.get_local_data().__rmod__, other)
Ultima's avatar
Ultima committed
436
    def __imod__(self, other):
437
        return self._builtin_helper(self.get_local_data().__imod__,
Ultima's avatar
Ultima committed
438
                                               other,
439
                                               inplace = True)
440
441
    def __len__(self):
        return self.shape[0]
442

443
    def get_dim(self):
444
        return np.prod(self.shape)
445

446
    def vdot(self, other):
447
        other = self.distributor.extract_local_data(other)
448
449
450
451
        local_vdot = np.vdot(self.get_local_data(), other)
        local_vdot_list = self.distributor._allgather(local_vdot)
        global_vdot = np.sum(local_vdot_list)
        return global_vdot
Ultimanet's avatar
Ultimanet committed
452

453
454


ultimanet's avatar
ultimanet committed
455
    def __getitem__(self, key):
Ultima's avatar
Ultima committed
456
        return self.get_data(key)
457
458


ultimanet's avatar
ultimanet committed
459
460
    def __setitem__(self, key, data):
        self.set_data(data, key)
461

462
    def _contraction_helper(self, function, **kwargs):
Ultima's avatar
Ultima committed
463
464
465
466
467
468
469
        if np.prod(self.data.shape) == 0:
            local = 0
            include = False
        else:
            local = function(self.data, **kwargs)
            include = True

470
        local_list = self.distributor._allgather(local)
Ultima's avatar
Ultima committed
471
472
473
474
475
476
        local_list = np.array(local_list, dtype = np.dtype(local_list[0]))
        include_list = np.array(self.distributor._allgather(include))
        work_list = local_list[include_list]
        if work_list.shape[0] == 0:
            raise ValueError("ERROR: Zero-size array to reduction operation "+
                             "which has no identity")
477
        else:
Ultima's avatar
Ultima committed
478
479
            result = function(work_list, axis=0)
            return result
480

481
    def amin(self, **kwargs):
482
        return self._contraction_helper(np.amin, **kwargs)
483
484

    def nanmin(self, **kwargs):
485
        return self._contraction_helper(np.nanmin, **kwargs)
486

487
    def amax(self, **kwargs):
488
        return self._contraction_helper(np.amax, **kwargs)
489

490
    def nanmax(self, **kwargs):
491
        return self._contraction_helper(np.nanmax, **kwargs)
492

493
494
495
496
    def sum(self, **kwargs):
        return self._contraction_helper(np.sum, **kwargs)

    def prod(self, **kwargs):
497
498
        return self._contraction_helper(np.prod, **kwargs)

499
    def mean(self, power=1):
500
        # compute the local means and the weights for the mean-mean.
Ultima's avatar
Ultima committed
501
502
503
504
505
506
        if np.prod(self.data.shape) == 0:
            local_mean = 0
            include = False
        else:
            local_mean = np.mean(self.data**power)
            include = True
507

508
        local_weight = np.prod(self.data.shape)
509
        # collect the local means and cast the result to a ndarray
Ultima's avatar
Ultima committed
510
511
        local_mean_list = self.distributor._allgather(local_mean)
        local_weight_list = self.distributor._allgather(local_weight)
512
513

        local_mean_list =np.array(local_mean_list,
Ultima's avatar
Ultima committed
514
                                  dtype = np.dtype(local_mean_list[0]))
515
516
        local_weight_list = np.array(local_weight_list)
        # extract the parts from the non-empty nodes
Ultima's avatar
Ultima committed
517
518
519
520
521
        include_list = np.array(self.distributor._allgather(include))
        work_mean_list = local_mean_list[include_list]
        work_weight_list = local_weight_list[include_list]
        if work_mean_list.shape[0] == 0:
            raise ValueError("ERROR:  Mean of empty slice.")
522
523
        else:
            # compute the denominator for the weighted mean-mean
Ultima's avatar
Ultima committed
524
            global_weight = np.sum(work_weight_list)
525
            # compute the numerator
Ultima's avatar
Ultima committed
526
527
528
            numerator = np.sum(work_mean_list * work_weight_list)
            global_mean = numerator/global_weight
            return global_mean
529
530
531
532
533

    def var(self):
        mean_of_the_square = self.mean(power=2)
        square_of_the_mean = self.mean()**2
        return mean_of_the_square - square_of_the_mean
534

535
536
    def std(self):
        return np.sqrt(self.var())
537

Ultima's avatar
Ultima committed
538
539
540
541
542
543
544
    def argmin(self):
        if np.prod(self.data.shape) == 0:
            local_argmin = np.nan
            local_argmin_value = np.nan
            globalized_local_argmin = np.nan
        else:
            local_argmin = np.argmin(self.data)
545
            local_argmin_value = self.data[np.unravel_index(local_argmin,
Ultima's avatar
Ultima committed
546
                                                            self.data.shape)]
547

Ultima's avatar
Ultima committed
548
            globalized_local_argmin = self.distributor.globalize_flat_index(
549
550
                                                                local_argmin)
        local_argmin_list = self.distributor._allgather((local_argmin_value,
551
552
                                                    globalized_local_argmin))
        local_argmin_list = np.array(local_argmin_list, dtype=[
Ultima's avatar
Ultima committed
553
                                        ('value', np.dtype('complex128')),
554
555
556
                                        ('index', np.dtype('float'))])
        local_argmin_list = np.sort(local_argmin_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
557
        return np.int(local_argmin_list[0][1])
558

Ultima's avatar
Ultima committed
559
560
561
562
563
564
565
    def argmax(self):
        if np.prod(self.data.shape) == 0:
            local_argmax = np.nan
            local_argmax_value = np.nan
            globalized_local_argmax = np.nan
        else:
            local_argmax = np.argmax(self.data)
566
            local_argmax_value = -self.data[np.unravel_index(local_argmax,
Ultima's avatar
Ultima committed
567
568
                                                            self.data.shape)]
            globalized_local_argmax = self.distributor.globalize_flat_index(
569
                                                                local_argmax)
570
        local_argmax_list = self.distributor._allgather((local_argmax_value,
571
572
                                                    globalized_local_argmax))
        local_argmax_list = np.array(local_argmax_list, dtype=[
Ultima's avatar
Ultima committed
573
                                        ('value', np.dtype('complex128')),
574
575
576
                                        ('index', np.dtype('float'))])
        local_argmax_list = np.sort(local_argmax_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
577
        return np.int(local_argmax_list[0][1])
578

579
580

    def argmin_nonflat(self):
Ultima's avatar
Ultima committed
581
        return np.unravel_index(self.argmin(), self.shape)
582

Ultima's avatar
Ultima committed
583
584
    def argmax_nonflat(self):
        return np.unravel_index(self.argmax(), self.shape)
585

586
587
588
589
590
591
    def conjugate(self):
        temp_d2o = self.copy_empty()
        temp_data = np.conj(self.get_local_data())
        temp_d2o.set_local_data(temp_data)
        return temp_d2o

592

593
    def conj(self):
594
595
        return self.conjugate()

596
    def median(self):
Ultimanet's avatar
Ultimanet committed
597
        about.warnings.cprint(\
598
599
600
            "WARNING: The current implementation of median is very expensive!")
        median = np.median(self.get_full_data())
        return median
601

602
    def iscomplex(self):
Ultima's avatar
Ultima committed
603
        temp_d2o = self.copy_empty(dtype=np.dtype('bool'))
604
605
        temp_d2o.set_local_data(np.iscomplex(self.data))
        return temp_d2o
606

607
    def isreal(self):
Ultima's avatar
Ultima committed
608
        temp_d2o = self.copy_empty(dtype=np.dtype('bool'))
609
610
        temp_d2o.set_local_data(np.isreal(self.data))
        return temp_d2o
611

612

613
614
615
616
617
618
619
620
    def all(self):
        local_all = np.all(self.get_local_data())
        global_all = self.distributor._allgather(local_all)
        return np.all(global_all)

    def any(self):
        local_any = np.any(self.get_local_data())
        global_any = self.distributor._allgather(local_any)
621
        return np.any(global_any)
622

623
624
625
626
627
    def unique(self):
        local_unique = np.unique(self.get_local_data())
        global_unique = self.distributor._allgather(local_unique)
        global_unique = np.concatenate(global_unique)
        return np.unique(global_unique)
628

629
    def bincount(self, weights = None, minlength = None):
630
631
        if self.dtype not in [np.dtype('int16'), np.dtype('int32'),
                np.dtype('int64'),  np.dtype('uint16'),
632
                np.dtype('uint32'), np.dtype('uint64')]:
633
            raise TypeError(about._errors.cstring(
634
635
                "ERROR: Distributed-data-object must be of integer datatype!"))

636
        minlength = max(self.amax()+1, minlength)
637

638
639
640
641
642
        if weights is not None:
            local_weights = self.distributor.extract_local_data(weights).\
                                                                    flatten()
        else:
            local_weights = None
643

644
645
646
        local_counts = np.bincount(self.get_local_data().flatten(),
                                  weights = local_weights,
                                  minlength = minlength)
Ultima's avatar
Ultima committed
647
648
649
650
651
652
        if self.distribution_strategy == 'not':
            return local_counts
        else:
            list_of_counts = self.distributor._allgather(local_counts)
            counts = np.sum(list_of_counts, axis = 0)
            return counts
653

Ultima's avatar
Ultima committed
654
655
656

    def where(self):
        return self.distributor.where(self.data)
657

658
    def set_local_data(self, data, hermitian=False, copy=True):
ultimanet's avatar
ultimanet committed
659
        """
660
            Stores data directly in the local data attribute. No distribution
ultimanet's avatar
ultimanet committed
661
662
663
664
665
            is done. The shape of the data must fit the local data attributes
            shape.

            Parameters
            ----------
666
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
667
                The data which should be stored in the local data attribute.
668

ultimanet's avatar
ultimanet committed
669
670
671
            Returns
            -------
            None
672

ultimanet's avatar
ultimanet committed
673
        """
Ultimanet's avatar
Ultimanet committed
674
        self.hermitian = hermitian
Ultima's avatar
Ultima committed
675
676
677
        if copy == True:
            self.data[:] = data
        else:
678
            self.data = np.array(data, dtype=self.dtype,
Ultima's avatar
Ultima committed
679
                            copy=False, order='C').reshape(self.local_shape)
680

Ultima's avatar
Ultima committed
681
    def set_data(self, data, to_key, from_key=None, local_keys=False,
682
                 hermitian=False, copy=True, **kwargs):
ultimanet's avatar
ultimanet committed
683
        """
684
            Stores the supplied data in the region which is specified by key.
ultimanet's avatar
ultimanet committed
685
            The data is distributed according to the distribution strategy. If
686
            the individual nodes get different key-arguments. Their data is
ultimanet's avatar
ultimanet committed
687
            processed one-by-one.
688

ultimanet's avatar
ultimanet committed
689
690
            Parameters
            ----------
691
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
692
693
                The data which should be distributed.
            key : int, slice, tuple of int or slice
694
695
696
                The key is the object which specifies the region, where data
                will be stored in.

ultimanet's avatar
ultimanet committed
697
698
699
            Returns
            -------
            None
700

ultimanet's avatar
ultimanet committed
701
        """
Ultimanet's avatar
Ultimanet committed
702
        self.hermitian = hermitian
Ultima's avatar
Ultima committed
703
        self.distributor.disperse_data(data = self.data,
704
                                       to_key = to_key,
Ultima's avatar
Ultima committed
705
                                       data_update = data,
706
                                       from_key = from_key,
Ultima's avatar
Ultima committed
707
                                       local_keys = local_keys,
Ultima's avatar
Ultima committed
708
709
                                       copy = copy,
                                       **kwargs)
710

711
    def set_full_data(self, data, hermitian=False, copy = True, **kwargs):
ultimanet's avatar
ultimanet committed
712
        """
713
            Distributes the supplied data to the nodes. The shape of data must
ultimanet's avatar
ultimanet committed
714
            match the shape of the distributed_data_object.
715

ultimanet's avatar
ultimanet committed
716
717
            Parameters
            ----------
718
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
719
                The data which should be distributed.
720

ultimanet's avatar
ultimanet committed
721
722
            Notes
            -----
723
            set_full_data(foo) is equivalent to set_data(foo,slice(None)) but
ultimanet's avatar
ultimanet committed
724
            faster.
725

ultimanet's avatar
ultimanet committed
726
727
728
            Returns
            -------
            None
729

ultimanet's avatar
ultimanet committed
730
        """
Ultimanet's avatar
Ultimanet committed
731
        self.hermitian = hermitian
732
        self.data = self.distributor.distribute_data(data=data, copy = copy,
733
                                                     **kwargs)
ultimanet's avatar
ultimanet committed
734

Ultimanet's avatar
Ultimanet committed
735
    def get_local_data(self, key=(slice(None),), copy=True):
ultimanet's avatar
ultimanet committed
736
        """
737
738
            Loads data directly from the local data attribute. No consolidation
            is done.
ultimanet's avatar
ultimanet committed
739
740
741
742

            Parameters
            ----------
            key : int, slice, tuple of int or slice
743
744
                The key which will be used to access the data.

ultimanet's avatar
ultimanet committed
745
746
747
            Returns
            -------
            self.data[key] : numpy.ndarray
748

Ultimanet's avatar
Ultimanet committed
749
        """
Ultimanet's avatar
Ultimanet committed
750
        if copy == True:
751
            return self.data[key]
Ultimanet's avatar
Ultimanet committed
752
753
        if copy == False:
            return self.data
754

755
    def get_data(self, key, local_keys=False, **kwargs):
ultimanet's avatar
ultimanet committed
756
        """
757
758
            Loads data from the region which is specified by key. The data is
            consolidated according to the distribution strategy. If the
ultimanet's avatar
ultimanet committed
759
            individual nodes get different key-arguments, they get individual
760
761
            data.

ultimanet's avatar
ultimanet committed
762
763
            Parameters
            ----------
764

ultimanet's avatar
ultimanet committed
765
            key : int, slice, tuple of int or slice
766
767
768
                The key is the object which specifies the region, where data
                will be loaded from.

ultimanet's avatar
ultimanet committed
769
770
771
            Returns
            -------
            global_data[key] : numpy.ndarray
772

ultimanet's avatar
ultimanet committed
773
        """
Ultima's avatar
Ultima committed
774
775
776
777
778
779
780
781
782
783
784
785
        if key is None:
            return self.copy()
        elif isinstance(key, slice):
            if key == slice(None):
                return self.copy()
        elif isinstance(key, tuple):
            try:
                if all(x == slice(None) for x in key):
                    return self.copy()
            except(ValueError):
                pass

786
787
788
        return self.distributor.collect_data(self.data,
                                             key,
                                             local_keys = local_keys,
789
                                             **kwargs)
790
791


ultimanet's avatar
ultimanet committed
792
793
    def get_full_data(self, target_rank='all'):
        """
794
795
            Fully consolidates the distributed data.

ultimanet's avatar
ultimanet committed
796
797
798
            Parameters
            ----------
            target_rank : 'all' (default), int *optional*
799
                If only one node should recieve the full data, it can be
ultimanet's avatar
ultimanet committed
800
                specified here.
801

ultimanet's avatar
ultimanet committed
802
803
            Notes
            -----
804
            get_full_data() is equivalent to get_data(slice(None)) but
ultimanet's avatar
ultimanet committed
805
            faster.
806

ultimanet's avatar
ultimanet committed
807
808
809
810
811
            Returns
            -------
            None
        """

812
        return self.distributor.consolidate_data(self.data,
813
                                                 target_rank = target_rank)
ultimanet's avatar
ultimanet committed
814

815
    def inject(self, to_key=(slice(None),), data=None,
Ultima's avatar
Ultima committed
816
817
               from_key=(slice(None),)):
        if data is None:
Ultimanet's avatar
Ultimanet committed
818
            return self
Ultima's avatar
Ultima committed
819
        self.distributor.inject(self.data, to_key, data, from_key)
820

821
    def flatten(self, inplace = False):
822
823
        flat_data = self.distributor.flatten(self.data, inplace = inplace)

Ultima's avatar
Ultima committed
824
825
        flat_global_shape = (np.prod(self.shape),)
        flat_local_shape = np.shape(flat_data)
826
827
828

        # Try to keep the distribution strategy. Therefore
        # create an empty copy of self which has the new shape
Ultima's avatar
Ultima committed
829
830
        temp_d2o = self.copy_empty(global_shape = flat_global_shape,
                                   local_shape = flat_local_shape)
831
        # Check if the local shapes match.
Ultima's avatar
Ultima committed
832
833
        if temp_d2o.local_shape == flat_local_shape:
            work_d2o = temp_d2o
834
        # if the shapes do not match, create a freeform d2o
Ultima's avatar
Ultima committed
835
836
837
838
        else:
            work_d2o = self.copy_empty(local_shape = flat_local_shape,
                            distribution_strategy = 'freeform')

839
        # Feed the work_d2o with the flat data
Ultima's avatar
Ultima committed
840
        work_d2o.set_local_data(data = flat_data,
841
842
                                copy = False)

843
        if inplace == True:
Ultima's avatar
Ultima committed
844
            self = work_d2o
845
846
            return self
        else:
Ultima's avatar
Ultima committed