nifty_mpi_data.py 112 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
# 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/>.
22
23


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

27
28
29
from keepers import about,\
                    global_configuration as gc,\
                    global_dependency_injector as gdi
ultimanet's avatar
ultimanet committed
30

31
32
33
MPI = gdi[gc['mpi_module']]
h5py = gdi.get('h5py')
pyfftw = gdi.get('pyfftw')
ultimanet's avatar
ultimanet committed
34

35
36
37
38
39
40
41
42
43
_maybe_fftw = ['fftw'] if ('pyfftw' in gdi) else []
STRATEGIES = {
                'all': ['not', 'equal', 'freeform'] + _maybe_fftw,
                'global': ['not', 'equal'] + _maybe_fftw,
                'local': ['freeform'],
                'slicing': ['equal', 'freeform'] + _maybe_fftw,
                'not': ['not'],
                'hdf5': ['equal'] + _maybe_fftw,
             }
44

45

ultimanet's avatar
ultimanet committed
46
47
48
49
50
51
52
53
class distributed_data_object(object):
    """

        NIFTY class for distributed data

        Parameters
        ----------
        global_data : {tuple, list, numpy.ndarray} *at least 1-dimensional*
54
            Initial data which will be casted to a numpy.ndarray and then
ultimanet's avatar
ultimanet committed
55
56
57
58
59
60
            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*
61
62
            If an explicit dtype is supplied, the given global_data will be
            casted to it.
ultimanet's avatar
ultimanet committed
63
        distribution_strategy : {'fftw' (default), 'not'}, *optional*
64
65
            Specifies the way, how global_data will be distributed to the
            individual nodes.
ultimanet's avatar
ultimanet committed
66
            'fftw' follows the distribution strategy of pyfftw.
67
68
            'not' does not distribute the data at all.

ultimanet's avatar
ultimanet committed
69
70
71
72
73
74
75
76
77
78

        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
79
80
            The distributor object which takes care of all distribution and
            consolidation of the data.
ultimanet's avatar
ultimanet committed
81
82
        shape : tuple of int
            The global shape of the data
83

ultimanet's avatar
ultimanet committed
84
85
        Raises
        ------
86
87
88
        TypeError :
            If the supplied distribution strategy is not known.

ultimanet's avatar
ultimanet committed
89
    """
90

91
    def __init__(self, global_data=None, global_shape=None, dtype=None,
Ultima's avatar
Ultima committed
92
                 local_data=None, local_shape=None,
93
                 distribution_strategy='fftw', hermitian=False,
94
95
96
97
98
                 alias=None, path=None, comm=MPI.COMM_WORLD,
                 copy=True, *args, **kwargs):

        # TODO: allow init with empty shape

99
100
101
102
        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)
103

104
        self.distributor = distributor_factory.get_distributor(
105
106
107
108
109
110
111
112
113
                                distribution_strategy=distribution_strategy,
                                comm=comm,
                                global_data=global_data,
                                global_shape=global_shape,
                                local_data=local_data,
                                local_shape=local_shape,
                                alias=alias,
                                path=path,
                                dtype=dtype,
114
                                **kwargs)
115

ultimanet's avatar
ultimanet committed
116
117
118
        self.distribution_strategy = distribution_strategy
        self.dtype = self.distributor.dtype
        self.shape = self.distributor.global_shape
Ultima's avatar
Ultima committed
119
120
        self.local_shape = self.distributor.local_shape
        self.comm = self.distributor.comm
121
122

        self.init_args = args
123
        self.init_kwargs = kwargs
124

Ultima's avatar
Ultima committed
125
        (self.data, self.hermitian) = self.distributor.initialize_data(
126
127
128
129
130
131
            global_data=global_data,
            local_data=local_data,
            alias=alias,
            path=path,
            hermitian=hermitian,
            copy=copy)
132
        self.index = d2o_librarian.register(self)
133

134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
    @property
    def real(self):
        new_data = self.get_local_data().real
        new_dtype = new_data.dtype
        new_d2o = self.copy_empty(dtype=new_dtype)
        new_d2o.set_local_data(data=new_data,
                               hermitian=self.hermitian)
        return new_d2o

    @property
    def imag(self):
        new_data = self.get_local_data().imag
        new_dtype = new_data.dtype
        new_d2o = self.copy_empty(dtype=new_dtype)
        new_d2o.set_local_data(data=new_data,
                               hermitian=self.hermitian)
        return new_d2o

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

    def copy_empty(self, global_shape=None, local_shape=None, dtype=None,
165
                   distribution_strategy=None, **kwargs):
Ultima's avatar
Ultima committed
166
        if self.distribution_strategy == 'not' and \
167
                distribution_strategy in STRATEGIES['local'] and \
168
169
170
171
172
                local_shape is None:
            result = self.copy_empty(global_shape=global_shape,
                                     local_shape=local_shape,
                                     dtype=dtype,
                                     distribution_strategy='equal',
Ultima's avatar
Ultima committed
173
                                     **kwargs)
174
            return result.copy_empty(distribution_strategy='freeform')
175

Ultima's avatar
Ultima committed
176
        if global_shape is None:
177
            global_shape = self.shape
Ultima's avatar
Ultima committed
178
179
180
        if local_shape is None:
            local_shape = self.local_shape
        if dtype is None:
181
            dtype = self.dtype
Ultima's avatar
Ultima committed
182
        if distribution_strategy is None:
183
184
185
            distribution_strategy = self.distribution_strategy

        kwargs.update(self.init_kwargs)
186

187
188
189
190
191
192
193
194
        temp_d2o = distributed_data_object(
                                   global_shape=global_shape,
                                   local_shape=local_shape,
                                   dtype=dtype,
                                   distribution_strategy=distribution_strategy,
                                   comm=self.comm,
                                   *self.init_args,
                                   **kwargs)
195
        return temp_d2o
196

197
    def apply_scalar_function(self, function, inplace=False, dtype=None):
198
        remember_hermitianQ = self.hermitian
199

200
        if inplace is True:
Ultimanet's avatar
Ultimanet committed
201
            temp = self
Ultima's avatar
Ultima committed
202
            if dtype is not None and self.dtype != np.dtype(dtype):
203
204
205
                about.warnings.cprint(
                    "WARNING: Inplace dtype conversion is not possible!")

Ultimanet's avatar
Ultimanet committed
206
        else:
207
            temp = self.copy_empty(dtype=dtype)
Ultimanet's avatar
Ultimanet committed
208

Ultima's avatar
Ultima committed
209
        if np.prod(self.local_shape) != 0:
210
            try:
Ultima's avatar
Ultima committed
211
212
                temp.data[:] = function(self.data)
            except:
213
214
                about.warnings.cprint(
                    "WARNING: Trying to use np.vectorize!")
Ultima's avatar
Ultima committed
215
216
                temp.data[:] = np.vectorize(function)(self.data)
        else:
217
218
            # Noting to do here. The value-empty array
            # is also geometrically empty
Ultima's avatar
Ultima committed
219
            pass
220

221
222
223
224
        if function in (np.exp, np.log):
            temp.hermitian = remember_hermitianQ
        else:
            temp.hermitian = False
Ultimanet's avatar
Ultimanet committed
225
        return temp
226

Ultimanet's avatar
Ultimanet committed
227
228
229
    def apply_generator(self, generator):
        self.set_local_data(generator(self.distributor.local_shape))
        self.hermitian = False
230

ultimanet's avatar
ultimanet committed
231
232
    def __str__(self):
        return self.data.__str__()
233

ultimanet's avatar
ultimanet committed
234
    def __repr__(self):
235
        return '<distributed_data_object>\n' + self.data.__repr__()
236

237
    def _compare_helper(self, other, op):
238
        result = self.copy_empty(dtype=np.bool_)
239
240
        # Case 1: 'other' is a scalar
        # -> make point-wise comparison
Ultimanet's avatar
Ultimanet committed
241
        if np.isscalar(other):
242
            result.set_local_data(
243
                getattr(self.get_local_data(copy=False), op)(other))
244
            return result
Ultimanet's avatar
Ultimanet committed
245

246
247
        # 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
248
        elif isinstance(other, np.ndarray) or\
249
                isinstance(other, distributed_data_object):
Ultimanet's avatar
Ultimanet committed
250
            temp_data = self.distributor.extract_local_data(other)
251
252
            result.set_local_data(
                getattr(self.get_local_data(copy=False), op)(temp_data))
Ultimanet's avatar
Ultimanet committed
253
            return result
254
255

        # Case 3: 'other' is None
Ultima's avatar
Ultima committed
256
        elif other is None:
Ultimanet's avatar
Ultimanet committed
257
            return False
258
259
260

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

265
266
    def __ne__(self, other):
        return self._compare_helper(other, '__ne__')
267

268
269
    def __lt__(self, other):
        return self._compare_helper(other, '__lt__')
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__')
277

278
279
280
281
282
283
    def __ge__(self, other):
        return self._compare_helper(other, '__ge__')

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

Ultima's avatar
Ultima committed
284
285
286
    def __iter__(self):
        return self.distributor.get_iter(self)

Ultimanet's avatar
Ultimanet committed
287
    def equal(self, other):
Ultimanet's avatar
Ultimanet committed
288
289
290
291
292
293
294
295
296
        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
297
        except(AssertionError, AttributeError):
Ultimanet's avatar
Ultimanet committed
298
299
300
301
            return False
        else:
            return True

302
    def __pos__(self):
303
        temp_d2o = self.copy_empty()
304
        temp_d2o.set_local_data(data=self.get_local_data(), copy=True)
305
        return temp_d2o
306

ultimanet's avatar
ultimanet committed
307
    def __neg__(self):
308
        temp_d2o = self.copy_empty()
309
310
        temp_d2o.set_local_data(data=self.get_local_data().__neg__(),
                                copy=True)
ultimanet's avatar
ultimanet committed
311
        return temp_d2o
312

313
    def __abs__(self):
314
        # translate complex dtypes
315
316
317
318
319
320
        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
321
322
        else:
            new_dtype = self.dtype
323
324
325
        temp_d2o = self.copy_empty(dtype=new_dtype)
        temp_d2o.set_local_data(data=self.get_local_data().__abs__(),
                                copy=True)
326
        return temp_d2o
327

Ultima's avatar
Ultima committed
328
329
330
331
332
    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)
333
334

        # Case 1: other is not a scalar
Ultimanet's avatar
Ultimanet committed
335
        if not (np.isscalar(other) or np.shape(other) == (1,)):
336
            try:
337
                hermitian_Q = (other.hermitian and self.hermitian)
338
339
            except(AttributeError):
                hermitian_Q = False
340
            # extract the local data from the 'other' object
Ultimanet's avatar
Ultimanet committed
341
342
            temp_data = self.distributor.extract_local_data(other)
            temp_data = operator(temp_data)
343
344

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

ultimanet's avatar
ultimanet committed
366
    def __add__(self, other):
Ultima's avatar
Ultima committed
367
        return self._builtin_helper(self.get_local_data().__add__, other)
ultimanet's avatar
ultimanet committed
368
369

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

    def __iadd__(self, other):
373
        return self._builtin_helper(self.get_local_data().__iadd__,
374
375
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
376

ultimanet's avatar
ultimanet committed
377
    def __sub__(self, other):
Ultima's avatar
Ultima committed
378
        return self._builtin_helper(self.get_local_data().__sub__, other)
379

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

ultimanet's avatar
ultimanet committed
383
    def __isub__(self, other):
384
        return self._builtin_helper(self.get_local_data().__isub__,
385
386
                                    other,
                                    inplace=True)
387

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

391
392
    def __truediv__(self, other):
        return self.__div__(other)
393

ultimanet's avatar
ultimanet committed
394
    def __rdiv__(self, other):
Ultima's avatar
Ultima committed
395
        return self._builtin_helper(self.get_local_data().__rdiv__, other)
396

397
398
    def __rtruediv__(self, other):
        return self.__rdiv__(other)
ultimanet's avatar
ultimanet committed
399

Ultimanet's avatar
Ultimanet committed
400
    def __idiv__(self, other):
401
        return self._builtin_helper(self.get_local_data().__idiv__,
402
403
404
                                    other,
                                    inplace=True)

405
    def __itruediv__(self, other):
406
        return self.__idiv__(other)
407

ultimanet's avatar
ultimanet committed
408
    def __floordiv__(self, other):
409
        return self._builtin_helper(self.get_local_data().__floordiv__,
410
411
                                    other)

ultimanet's avatar
ultimanet committed
412
    def __rfloordiv__(self, other):
413
        return self._builtin_helper(self.get_local_data().__rfloordiv__,
414
415
                                    other)

Ultimanet's avatar
Ultimanet committed
416
    def __ifloordiv__(self, other):
Ultima's avatar
Ultima committed
417
        return self._builtin_helper(
418
419
            self.get_local_data().__ifloordiv__, other,
            inplace=True)
420

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

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

    def __imul__(self, other):
428
        return self._builtin_helper(self.get_local_data().__imul__,
429
430
                                    other,
                                    inplace=True)
Ultimanet's avatar
Ultimanet committed
431

ultimanet's avatar
ultimanet committed
432
    def __pow__(self, other):
Ultima's avatar
Ultima committed
433
        return self._builtin_helper(self.get_local_data().__pow__, other)
434

ultimanet's avatar
ultimanet committed
435
    def __rpow__(self, other):
Ultima's avatar
Ultima committed
436
        return self._builtin_helper(self.get_local_data().__rpow__, other)
ultimanet's avatar
ultimanet committed
437
438

    def __ipow__(self, other):
439
        return self._builtin_helper(self.get_local_data().__ipow__,
440
441
442
                                    other,
                                    inplace=True)

Ultima's avatar
Ultima committed
443
444
    def __mod__(self, other):
        return self._builtin_helper(self.get_local_data().__mod__, other)
445

Ultima's avatar
Ultima committed
446
    def __rmod__(self, other):
447
        return self._builtin_helper(self.get_local_data().__rmod__, other)
448

Ultima's avatar
Ultima committed
449
    def __imod__(self, other):
450
        return self._builtin_helper(self.get_local_data().__imod__,
451
452
453
                                    other,
                                    inplace=True)

454
455
    def __len__(self):
        return self.shape[0]
456

457
    def get_dim(self):
458
        return np.prod(self.shape)
459

460
    def vdot(self, other):
461
        other = self.distributor.extract_local_data(other)
462
463
        local_vdot = np.vdot(self.get_local_data(), other)
        local_vdot_list = self.distributor._allgather(local_vdot)
464
465
        global_vdot = np.result_type(self.dtype,
                                     other.dtype).type(np.sum(local_vdot_list))
466
        return global_vdot
Ultimanet's avatar
Ultimanet committed
467

ultimanet's avatar
ultimanet committed
468
    def __getitem__(self, key):
Ultima's avatar
Ultima committed
469
        return self.get_data(key)
470

ultimanet's avatar
ultimanet committed
471
472
    def __setitem__(self, key, data):
        self.set_data(data, key)
473

474
    def _contraction_helper(self, function, **kwargs):
475
476
477
        if self.shape == (0,):
            raise ValueError("ERROR: Zero-size array to reduction operation " +
                             "which has no identity")
Ultima's avatar
Ultima committed
478
479
480
481
482
483
484
        if np.prod(self.data.shape) == 0:
            local = 0
            include = False
        else:
            local = function(self.data, **kwargs)
            include = True

485
        local_list = self.distributor._allgather(local)
486
        local_list = np.array(local_list, dtype=np.dtype(local_list[0]))
Ultima's avatar
Ultima committed
487
488
489
        include_list = np.array(self.distributor._allgather(include))
        work_list = local_list[include_list]
        if work_list.shape[0] == 0:
490
            raise ValueError("ERROR: Zero-size array to reduction operation " +
Ultima's avatar
Ultima committed
491
                             "which has no identity")
492
        else:
Ultima's avatar
Ultima committed
493
494
            result = function(work_list, axis=0)
            return result
495

496
497
498
    def min(self, **kwargs):
        return self.amin(**kwargs)

499
    def amin(self, **kwargs):
500
        return self._contraction_helper(np.amin, **kwargs)
501
502

    def nanmin(self, **kwargs):
503
        return self._contraction_helper(np.nanmin, **kwargs)
504

505
506
507
    def max(self, **kwargs):
        return self.amax(**kwargs)

508
    def amax(self, **kwargs):
509
        return self._contraction_helper(np.amax, **kwargs)
510

511
    def nanmax(self, **kwargs):
512
        return self._contraction_helper(np.nanmax, **kwargs)
513

514
    def sum(self, **kwargs):
515
516
        if self.shape == (0,):
            return self.dtype.type(0)
517
518
519
        return self._contraction_helper(np.sum, **kwargs)

    def prod(self, **kwargs):
520
521
        if self.shape == (0,):
            return self.dtype.type(1)
522
523
        return self._contraction_helper(np.prod, **kwargs)

524
    def mean(self, power=1):
525
526
        if self.shape == (0,):
            return np.mean(np.array([], dtype=self.dtype))
527
        # compute the local means and the weights for the mean-mean.
Ultima's avatar
Ultima committed
528
529
530
531
532
533
        if np.prod(self.data.shape) == 0:
            local_mean = 0
            include = False
        else:
            local_mean = np.mean(self.data**power)
            include = True
534

535
        local_weight = np.prod(self.data.shape)
536
        # collect the local means and cast the result to a ndarray
Ultima's avatar
Ultima committed
537
538
        local_mean_list = self.distributor._allgather(local_mean)
        local_weight_list = self.distributor._allgather(local_weight)
539

540
541
        local_mean_list = np.array(local_mean_list,
                                   dtype=np.dtype(local_mean_list[0]))
542
543
        local_weight_list = np.array(local_weight_list)
        # extract the parts from the non-empty nodes
Ultima's avatar
Ultima committed
544
545
546
547
548
        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.")
549
550
        else:
            # compute the denominator for the weighted mean-mean
Ultima's avatar
Ultima committed
551
            global_weight = np.sum(work_weight_list)
552
            # compute the numerator
Ultima's avatar
Ultima committed
553
            numerator = np.sum(work_mean_list * work_weight_list)
554
            global_mean = numerator / global_weight
Ultima's avatar
Ultima committed
555
            return global_mean
556
557

    def var(self):
558
559
        if self.shape == (0,):
            return np.var(np.array([], dtype=self.dtype))
560
561
562
563
564
565
566

        if issubclass(self.dtype.type, np.complexfloating):
            mean_of_the_square = abs(self**2).mean()
            square_of_the_mean = abs(self.mean())**2
        else:
            mean_of_the_square = self.mean(power=2)
            square_of_the_mean = self.mean()**2
567
        return mean_of_the_square - square_of_the_mean
568

569
    def std(self):
570
571
572
573
        if self.shape == (0,):
            return np.std(np.array([], dtype=self.dtype))
        if self.shape == (0,):
            return np.nan
574
        return np.sqrt(self.var())
575

Ultima's avatar
Ultima committed
576
    def argmin(self):
577
578
579
        if self.shape == (0,):
            raise ValueError(
                "ERROR: attempt to get argmin of an empty object")
Ultima's avatar
Ultima committed
580
581
582
583
584
585
        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)
586
            local_argmin_value = self.data[np.unravel_index(local_argmin,
Ultima's avatar
Ultima committed
587
                                                            self.data.shape)]
588

Ultima's avatar
Ultima committed
589
            globalized_local_argmin = self.distributor.globalize_flat_index(
590
591
592
593
                local_argmin)
        local_argmin_list = self.distributor._allgather(
                                                    (local_argmin_value,
                                                     globalized_local_argmin))
594
        local_argmin_list = np.array(local_argmin_list, dtype=[
595
596
            ('value', np.dtype('complex128')),
            ('index', np.dtype('float'))])
597
598
        local_argmin_list = np.sort(local_argmin_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
599
        return np.int(local_argmin_list[0][1])
600

Ultima's avatar
Ultima committed
601
    def argmax(self):
602
603
604
        if self.shape == (0,):
            raise ValueError(
                "ERROR: attempt to get argmax of an empty object")
Ultima's avatar
Ultima committed
605
606
607
608
609
610
        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)
611
            local_argmax_value = -self.data[np.unravel_index(local_argmax,
612
                                                             self.data.shape)]
Ultima's avatar
Ultima committed
613
            globalized_local_argmax = self.distributor.globalize_flat_index(
614
615
616
617
                local_argmax)
        local_argmax_list = self.distributor._allgather(
                                                  (local_argmax_value,
                                                   globalized_local_argmax))
618
        local_argmax_list = np.array(local_argmax_list, dtype=[
619
620
            ('value', np.dtype('complex128')),
            ('index', np.dtype('float'))])
621
622
        local_argmax_list = np.sort(local_argmax_list,
                                    order=['value', 'index'])
Ultima's avatar
Ultima committed
623
        return np.int(local_argmax_list[0][1])
624

625
    def argmin_nonflat(self):
Ultima's avatar
Ultima committed
626
        return np.unravel_index(self.argmin(), self.shape)
627

Ultima's avatar
Ultima committed
628
629
    def argmax_nonflat(self):
        return np.unravel_index(self.argmax(), self.shape)
630

631
632
633
634
635
636
637
    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

    def conj(self):
638
639
        return self.conjugate()

640
    def median(self):
641
        about.warnings.cprint(
642
643
644
            "WARNING: The current implementation of median is very expensive!")
        median = np.median(self.get_full_data())
        return median
645

646
    def _is_helper(self, function):
Ultima's avatar
Ultima committed
647
        temp_d2o = self.copy_empty(dtype=np.dtype('bool'))
648
        temp_d2o.set_local_data(function(self.data))
649
        return temp_d2o
650

651
652
653
    def iscomplex(self):
        return self._is_helper(np.iscomplex)

654
    def isreal(self):
655
656
657
658
659
660
661
662
663
664
665
666
667
668
        return self._is_helper(np.isreal)

    def isnan(self):
        return self._is_helper(np.isnan)

    def isinf(self):
        return self._is_helper(np.isinf)

    def isfinite(self):
        return self._is_helper(np.isfinite)

    def nan_to_num(self):
        temp_d2o = self.copy_empty()
        temp_d2o.set_local_data(np.nan_to_num(self.data))
669
        return temp_d2o
670

671
672
673
674
675
676
677
678
    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)
679
        return np.any(global_any)
680

681
682
683
684
685
    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)
686

687
    def bincount(self, weights=None, minlength=None):
688
        if self.dtype not in [np.dtype('int16'), np.dtype('int32'),
689
690
                              np.dtype('int64'),  np.dtype('uint16'),
                              np.dtype('uint32'), np.dtype('uint64')]:
691
            raise TypeError(about._errors.cstring(
692
693
                "ERROR: Distributed-data-object must be of integer datatype!"))

694
        minlength = max(self.amax() + 1, minlength)
695

696
697
        if weights is not None:
            local_weights = self.distributor.extract_local_data(weights).\
698
                                flatten()
699
700
        else:
            local_weights = None
701

702
        local_counts = np.bincount(self.get_local_data().flatten(),
703
704
                                   weights=local_weights,
                                   minlength=minlength)
Ultima's avatar
Ultima committed
705
706
707
708
        if self.distribution_strategy == 'not':
            return local_counts
        else:
            list_of_counts = self.distributor._allgather(local_counts)
709
            counts = np.sum(list_of_counts, axis=0)
Ultima's avatar
Ultima committed
710
            return counts
711

Ultima's avatar
Ultima committed
712
713
    def where(self):
        return self.distributor.where(self.data)
714

715
    def set_local_data(self, data, hermitian=False, copy=True):
ultimanet's avatar
ultimanet committed
716
        """
717
            Stores data directly in the local data attribute. No distribution
ultimanet's avatar
ultimanet committed
718
719
720
721
722
            is done. The shape of the data must fit the local data attributes
            shape.

            Parameters
            ----------
723
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
724
                The data which should be stored in the local data attribute.
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
        if copy is True:
Ultima's avatar
Ultima committed
733
734
            self.data[:] = data
        else:
735
736
737
738
            self.data = np.array(data,
                                 dtype=self.dtype,
                                 copy=False,
                                 order='C').reshape(self.local_shape)
739

Ultima's avatar
Ultima committed
740
    def set_data(self, data, to_key, from_key=None, local_keys=False,
741
                 hermitian=False, copy=True, **kwargs):
ultimanet's avatar
ultimanet committed
742
        """
743
            Stores the supplied data in the region which is specified by key.
ultimanet's avatar
ultimanet committed
744
            The data is distributed according to the distribution strategy. If
745
            the individual nodes get different key-arguments. Their data is
ultimanet's avatar
ultimanet committed
746
            processed one-by-one.
747

ultimanet's avatar
ultimanet committed
748
749
            Parameters
            ----------
750
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
751
752
                The data which should be distributed.
            key : int, slice, tuple of int or slice
753
754
755
                The key is the object which specifies the region, where data
                will be stored in.

ultimanet's avatar
ultimanet committed
756
757
758
            Returns
            -------
            None
759

ultimanet's avatar
ultimanet committed
760
        """
Ultimanet's avatar
Ultimanet committed
761
        self.hermitian = hermitian
762
763
764
765
766
767
        self.distributor.disperse_data(data=self.data,
                                       to_key=to_key,
                                       data_update=data,
                                       from_key=from_key,
                                       local_keys=local_keys,
                                       copy=copy,
Ultima's avatar
Ultima committed
768
                                       **kwargs)
769

770
    def set_full_data(self, data, hermitian=False, copy=True, **kwargs):
ultimanet's avatar
ultimanet committed
771
        """
772
            Distributes the supplied data to the nodes. The shape of data must
ultimanet's avatar
ultimanet committed
773
            match the shape of the distributed_data_object.
774

ultimanet's avatar
ultimanet committed
775
776
            Parameters
            ----------
777
            data : tuple, list, numpy.ndarray
ultimanet's avatar
ultimanet committed
778
                The data which should be distributed.
779

ultimanet's avatar
ultimanet committed
780
781
            Notes
            -----
782
            set_full_data(foo) is equivalent to set_data(foo,slice(None)) but
ultimanet's avatar
ultimanet committed
783
            faster.
784

ultimanet's avatar
ultimanet committed
785
786
787
            Returns
            -------
            None
788

ultimanet's avatar
ultimanet committed
789
        """
Ultimanet's avatar
Ultimanet committed
790
        self.hermitian = hermitian
791
        self.data = self.distributor.distribute_data(data=data, copy=copy,
792
                                                     **kwargs)
ultimanet's avatar
ultimanet committed
793

Ultimanet's avatar
Ultimanet committed
794
    def get_local_data(self, key=(slice(None),), copy=True):
ultimanet's avatar
ultimanet committed
795
        """
796
797
            Loads data directly from the local data attribute. No consolidation
            is done.
ultimanet's avatar
ultimanet committed
798
799
800
801

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

ultimanet's avatar
ultimanet committed
804
805
806
            Returns
            -------
            self.data[key] : numpy.ndarray
807

Ultimanet's avatar
Ultimanet committed
808
        """
809
        if copy is True:
810
            return self.data[key]
811
        if copy is False:
Ultimanet's avatar
Ultimanet committed
812
            return self.data
813

814
    def get_data(self, key, local_keys=False, **kwargs):
ultimanet's avatar
ultimanet committed
815
        """
816
817
            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
818
            individual nodes get different key-arguments, they get individual
819
820
            data.

ultimanet's avatar
ultimanet committed
821
822
            Parameters
            ----------
823

ultimanet's avatar
ultimanet committed
824
            key : int, slice, tuple of int or slice
825
826
827
                The key is the object which specifies the region, where data
                will be loaded from.

ultimanet's avatar
ultimanet committed
828
829
830
            Returns
            -------
            global_data[key] : numpy.ndarray
831

ultimanet's avatar
ultimanet committed
832
        """
Ultima's avatar
Ultima committed
833
834
835
836
837
838
839
840
841
842
843
844
        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

845
846
        return self.distributor.collect_data(self.data,
                                             key,
847
                                             local_keys=local_keys,
848
                                             **kwargs)